Return ThumbnailInfo in more places (#16438)

Improves type hints by using concrete types instead of
dictionaries.

6 files changed, 90 insertions, 80 deletions

diff --git a/changelog.d/16438.misc b/changelog.d/16438.misc
new file mode 100644
index 0000000000..bd7cdd42af
--- /dev/null
+++ b/changelog.d/16438.misc
@@ -0,0 +1 @@
+Reduce memory allocations.
diff --git a/synapse/media/_base.py b/synapse/media/_base.py
index d103b43449..13345acf75 100644
--- a/synapse/media/_base.py
+++ b/synapse/media/_base.py
@@ -332,7 +332,7 @@ class ThumbnailInfo:
     # Content type of thumbnail, e.g. image/png
     type: str
     # The size of the media file, in bytes.
-    length: Optional[int] = None
+    length: int
 
 
 @attr.s(slots=True, frozen=True, auto_attribs=True)
diff --git a/synapse/media/media_repository.py b/synapse/media/media_repository.py
index d11c2ff4ee..7fd46901f7 100644
--- a/synapse/media/media_repository.py
+++ b/synapse/media/media_repository.py
@@ -624,6 +624,7 @@ class MediaRepository:
                         height=t_height,
                         method=t_method,
                         type=t_type,
+                        length=t_byte_source.tell(),
                     ),
                 )
 
@@ -694,6 +695,7 @@ class MediaRepository:
                         height=t_height,
                         method=t_method,
                         type=t_type,
+                        length=t_byte_source.tell(),
                     ),
                 )
 
@@ -839,6 +841,7 @@ class MediaRepository:
                         height=t_height,
                         method=t_method,
                         type=t_type,
+                        length=t_byte_source.tell(),
                     ),
                 )
 
diff --git a/synapse/rest/media/thumbnail_resource.py b/synapse/rest/media/thumbnail_resource.py
index f9cd773f77..85b6bdbe72 100644
--- a/synapse/rest/media/thumbnail_resource.py
+++ b/synapse/rest/media/thumbnail_resource.py
@@ -15,7 +15,7 @@
 
 import logging
 import re
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 from synapse.api.errors import Codes, SynapseError, cs_error
 from synapse.config.repository import THUMBNAIL_SUPPORTED_MEDIA_FORMAT_MAP
@@ -159,30 +159,24 @@ class ThumbnailResource(RestServlet):
 
         thumbnail_infos = await self.store.get_local_media_thumbnails(media_id)
         for info in thumbnail_infos:
-            t_w = info["thumbnail_width"] == desired_width
-            t_h = info["thumbnail_height"] == desired_height
-            t_method = info["thumbnail_method"] == desired_method
-            t_type = info["thumbnail_type"] == desired_type
+            t_w = info.width == desired_width
+            t_h = info.height == desired_height
+            t_method = info.method == desired_method
+            t_type = info.type == desired_type
 
             if t_w and t_h and t_method and t_type:
                 file_info = FileInfo(
                     server_name=None,
                     file_id=media_id,
                     url_cache=media_info["url_cache"],
-                    thumbnail=ThumbnailInfo(
-                        width=info["thumbnail_width"],
-                        height=info["thumbnail_height"],
-                        type=info["thumbnail_type"],
-                        method=info["thumbnail_method"],
-                    ),
+                    thumbnail=info,
                 )
 
-                t_type = file_info.thumbnail_type
-                t_length = info["thumbnail_length"]
-
                 responder = await self.media_storage.fetch_media(file_info)
                 if responder:
-                    await respond_with_responder(request, responder, t_type, t_length)
+                    await respond_with_responder(
+                        request, responder, info.type, info.length
+                    )
                     return
 
         logger.debug("We don't have a thumbnail of that size. Generating")
@@ -222,29 +216,23 @@ class ThumbnailResource(RestServlet):
         file_id = media_info["filesystem_id"]
 
         for info in thumbnail_infos:
-            t_w = info["thumbnail_width"] == desired_width
-            t_h = info["thumbnail_height"] == desired_height
-            t_method = info["thumbnail_method"] == desired_method
-            t_type = info["thumbnail_type"] == desired_type
+            t_w = info.width == desired_width
+            t_h = info.height == desired_height
+            t_method = info.method == desired_method
+            t_type = info.type == desired_type
 
             if t_w and t_h and t_method and t_type:
                 file_info = FileInfo(
                     server_name=server_name,
                     file_id=media_info["filesystem_id"],
-                    thumbnail=ThumbnailInfo(
-                        width=info["thumbnail_width"],
-                        height=info["thumbnail_height"],
-                        type=info["thumbnail_type"],
-                        method=info["thumbnail_method"],
-                    ),
+                    thumbnail=info,
                 )
 
-                t_type = file_info.thumbnail_type
-                t_length = info["thumbnail_length"]
-
                 responder = await self.media_storage.fetch_media(file_info)
                 if responder:
-                    await respond_with_responder(request, responder, t_type, t_length)
+                    await respond_with_responder(
+                        request, responder, info.type, info.length
+                    )
                     return
 
         logger.debug("We don't have a thumbnail of that size. Generating")
@@ -304,7 +292,7 @@ class ThumbnailResource(RestServlet):
         desired_height: int,
         desired_method: str,
         desired_type: str,
-        thumbnail_infos: List[Dict[str, Any]],
+        thumbnail_infos: List[ThumbnailInfo],
         media_id: str,
         file_id: str,
         url_cache: bool,
@@ -319,7 +307,7 @@ class ThumbnailResource(RestServlet):
             desired_height: The desired height, the returned thumbnail may be larger than this.
             desired_method: The desired method used to generate the thumbnail.
             desired_type: The desired content-type of the thumbnail.
-            thumbnail_infos: A list of dictionaries of candidate thumbnails.
+            thumbnail_infos: A list of thumbnail info of candidate thumbnails.
             file_id: The ID of the media that a thumbnail is being requested for.
             url_cache: True if this is from a URL cache.
             server_name: The server name, if this is a remote thumbnail.
@@ -443,7 +431,7 @@ class ThumbnailResource(RestServlet):
         desired_height: int,
         desired_method: str,
         desired_type: str,
-        thumbnail_infos: List[Dict[str, Any]],
+        thumbnail_infos: List[ThumbnailInfo],
         file_id: str,
         url_cache: bool,
         server_name: Optional[str],
@@ -456,7 +444,7 @@ class ThumbnailResource(RestServlet):
             desired_height: The desired height, the returned thumbnail may be larger than this.
             desired_method: The desired method used to generate the thumbnail.
             desired_type: The desired content-type of the thumbnail.
-            thumbnail_infos: A list of dictionaries of candidate thumbnails.
+            thumbnail_infos: A list of thumbnail infos of candidate thumbnails.
             file_id: The ID of the media that a thumbnail is being requested for.
             url_cache: True if this is from a URL cache.
             server_name: The server name, if this is a remote thumbnail.
@@ -474,21 +462,25 @@ class ThumbnailResource(RestServlet):
 
         if desired_method == "crop":
             # Thumbnails that match equal or larger sizes of desired width/height.
-            crop_info_list: List[Tuple[int, int, int, bool, int, Dict[str, Any]]] = []
+            crop_info_list: List[
+                Tuple[int, int, int, bool, Optional[int], ThumbnailInfo]
+            ] = []
             # Other thumbnails.
-            crop_info_list2: List[Tuple[int, int, int, bool, int, Dict[str, Any]]] = []
+            crop_info_list2: List[
+                Tuple[int, int, int, bool, Optional[int], ThumbnailInfo]
+            ] = []
             for info in thumbnail_infos:
                 # Skip thumbnails generated with different methods.
-                if info["thumbnail_method"] != "crop":
+                if info.method != "crop":
                     continue
 
-                t_w = info["thumbnail_width"]
-                t_h = info["thumbnail_height"]
+                t_w = info.width
+                t_h = info.height
                 aspect_quality = abs(d_w * t_h - d_h * t_w)
                 min_quality = 0 if d_w <= t_w and d_h <= t_h else 1
                 size_quality = abs((d_w - t_w) * (d_h - t_h))
-                type_quality = desired_type != info["thumbnail_type"]
-                length_quality = info["thumbnail_length"]
+                type_quality = desired_type != info.type
+                length_quality = info.length
                 if t_w >= d_w or t_h >= d_h:
                     crop_info_list.append(
                         (
@@ -513,7 +505,7 @@ class ThumbnailResource(RestServlet):
                     )
             # Pick the most appropriate thumbnail. Some values of `desired_width` and
             # `desired_height` may result in a tie, in which case we avoid comparing on
-            # the thumbnail info dictionary and pick the thumbnail that appears earlier
+            # the thumbnail info and pick the thumbnail that appears earlier
             # in the list of candidates.
             if crop_info_list:
                 thumbnail_info = min(crop_info_list, key=lambda t: t[:-1])[-1]
@@ -521,20 +513,20 @@ class ThumbnailResource(RestServlet):
                 thumbnail_info = min(crop_info_list2, key=lambda t: t[:-1])[-1]
         elif desired_method == "scale":
             # Thumbnails that match equal or larger sizes of desired width/height.
-            info_list: List[Tuple[int, bool, int, Dict[str, Any]]] = []
+            info_list: List[Tuple[int, bool, int, ThumbnailInfo]] = []
             # Other thumbnails.
-            info_list2: List[Tuple[int, bool, int, Dict[str, Any]]] = []
+            info_list2: List[Tuple[int, bool, int, ThumbnailInfo]] = []
 
             for info in thumbnail_infos:
                 # Skip thumbnails generated with different methods.
-                if info["thumbnail_method"] != "scale":
+                if info.method != "scale":
                     continue
 
-                t_w = info["thumbnail_width"]
-                t_h = info["thumbnail_height"]
+                t_w = info.width
+                t_h = info.height
                 size_quality = abs((d_w - t_w) * (d_h - t_h))
-                type_quality = desired_type != info["thumbnail_type"]
-                length_quality = info["thumbnail_length"]
+                type_quality = desired_type != info.type
+                length_quality = info.length
                 if t_w >= d_w or t_h >= d_h:
                     info_list.append((size_quality, type_quality, length_quality, info))
                 else:
@@ -543,7 +535,7 @@ class ThumbnailResource(RestServlet):
                     )
             # Pick the most appropriate thumbnail. Some values of `desired_width` and
             # `desired_height` may result in a tie, in which case we avoid comparing on
-            # the thumbnail info dictionary and pick the thumbnail that appears earlier
+            # the thumbnail info and pick the thumbnail that appears earlier
             # in the list of candidates.
             if info_list:
                 thumbnail_info = min(info_list, key=lambda t: t[:-1])[-1]
@@ -555,13 +547,7 @@ class ThumbnailResource(RestServlet):
                 file_id=file_id,
                 url_cache=url_cache,
                 server_name=server_name,
-                thumbnail=ThumbnailInfo(
-                    width=thumbnail_info["thumbnail_width"],
-                    height=thumbnail_info["thumbnail_height"],
-                    type=thumbnail_info["thumbnail_type"],
-                    method=thumbnail_info["thumbnail_method"],
-                    length=thumbnail_info["thumbnail_length"],
-                ),
+                thumbnail=thumbnail_info,
             )
 
         # No matching thumbnail was found.
diff --git a/synapse/storage/databases/main/media_repository.py b/synapse/storage/databases/main/media_repository.py
index 8cebeb5189..2e6b176bd2 100644
--- a/synapse/storage/databases/main/media_repository.py
+++ b/synapse/storage/databases/main/media_repository.py
@@ -28,6 +28,7 @@ from typing import (
 
 from synapse.api.constants import Direction
 from synapse.logging.opentracing import trace
+from synapse.media._base import ThumbnailInfo
 from synapse.storage._base import SQLBaseStore
 from synapse.storage.database import (
     DatabasePool,
@@ -435,8 +436,8 @@ class MediaRepositoryStore(MediaRepositoryBackgroundUpdateStore):
             desc="store_url_cache",
         )
 
-    async def get_local_media_thumbnails(self, media_id: str) -> List[Dict[str, Any]]:
-        return await self.db_pool.simple_select_list(
+    async def get_local_media_thumbnails(self, media_id: str) -> List[ThumbnailInfo]:
+        rows = await self.db_pool.simple_select_list(
             "local_media_repository_thumbnails",
             {"media_id": media_id},
             (
@@ -448,6 +449,16 @@ class MediaRepositoryStore(MediaRepositoryBackgroundUpdateStore):
             ),
             desc="get_local_media_thumbnails",
         )
+        return [
+            ThumbnailInfo(
+                width=row["thumbnail_width"],
+                height=row["thumbnail_height"],
+                method=row["thumbnail_method"],
+                type=row["thumbnail_type"],
+                length=row["thumbnail_length"],
+            )
+            for row in rows
+        ]
 
     @trace
     async def store_local_thumbnail(
@@ -556,8 +567,8 @@ class MediaRepositoryStore(MediaRepositoryBackgroundUpdateStore):
 
     async def get_remote_media_thumbnails(
         self, origin: str, media_id: str
-    ) -> List[Dict[str, Any]]:
-        return await self.db_pool.simple_select_list(
+    ) -> List[ThumbnailInfo]:
+        rows = await self.db_pool.simple_select_list(
             "remote_media_cache_thumbnails",
             {"media_origin": origin, "media_id": media_id},
             (
@@ -566,10 +577,19 @@ class MediaRepositoryStore(MediaRepositoryBackgroundUpdateStore):
                 "thumbnail_method",
                 "thumbnail_type",
                 "thumbnail_length",
-                "filesystem_id",
             ),
             desc="get_remote_media_thumbnails",
         )
+        return [
+            ThumbnailInfo(
+                width=row["thumbnail_width"],
+                height=row["thumbnail_height"],
+                method=row["thumbnail_method"],
+                type=row["thumbnail_type"],
+                length=row["thumbnail_length"],
+            )
+            for row in rows
+        ]
 
     @trace
     async def get_remote_media_thumbnail(
diff --git a/tests/media/test_media_storage.py b/tests/media/test_media_storage.py
index ba00e35a9e..15f5d644e4 100644
--- a/tests/media/test_media_storage.py
+++ b/tests/media/test_media_storage.py
@@ -34,7 +34,7 @@ from synapse.api.errors import Codes
 from synapse.events import EventBase
 from synapse.http.types import QueryParams
 from synapse.logging.context import make_deferred_yieldable
-from synapse.media._base import FileInfo
+from synapse.media._base import FileInfo, ThumbnailInfo
 from synapse.media.filepath import MediaFilePaths
 from synapse.media.media_storage import MediaStorage, ReadableFileWrapper
 from synapse.media.storage_provider import FileStorageProviderBackend
@@ -605,6 +605,8 @@ class MediaRepoTests(unittest.HomeserverTestCase):
         """Test that choosing between thumbnails with the same quality rating succeeds.
 
         We are not particular about which thumbnail is chosen."""
+
+        content_type = self.test_image.content_type.decode()
         media_repo = self.hs.get_media_repository()
         thumbnail_resouce = ThumbnailResource(
             self.hs, media_repo, media_repo.media_storage
@@ -615,26 +617,24 @@ class MediaRepoTests(unittest.HomeserverTestCase):
                 desired_width=desired_size,
                 desired_height=desired_size,
                 desired_method=method,
-                desired_type=self.test_image.content_type,  # type: ignore[arg-type]
+                desired_type=content_type,
                 # Provide two identical thumbnails which are guaranteed to have the same
                 # quality rating.
                 thumbnail_infos=[
-                    {
-                        "thumbnail_width": 32,
-                        "thumbnail_height": 32,
-                        "thumbnail_method": method,
-                        "thumbnail_type": self.test_image.content_type,
-                        "thumbnail_length": 256,
-                        "filesystem_id": f"thumbnail1{self.test_image.extension.decode()}",
-                    },
-                    {
-                        "thumbnail_width": 32,
-                        "thumbnail_height": 32,
-                        "thumbnail_method": method,
-                        "thumbnail_type": self.test_image.content_type,
-                        "thumbnail_length": 256,
-                        "filesystem_id": f"thumbnail2{self.test_image.extension.decode()}",
-                    },
+                    ThumbnailInfo(
+                        width=32,
+                        height=32,
+                        method=method,
+                        type=content_type,
+                        length=256,
+                    ),
+                    ThumbnailInfo(
+                        width=32,
+                        height=32,
+                        method=method,
+                        type=content_type,
+                        length=256,
+                    ),
                 ],
                 file_id=f"image{self.test_image.extension.decode()}",
                 url_cache=False,