snapshot of https://github.com/zulip/zulip.git @ 48c69911c7dd4e03b20552eb3bd77ad8151c1357 (2026-09-15)

zerver/worker/thumbnail.py

@@ -0,0 +1,255 @@
+import logging
+import time
+from dataclasses import asdict, dataclass
+from io import BytesIO
+from typing import Any
+
+import pyvips
+from django.db import transaction
+from typing_extensions import override
+
+from zerver.actions.message_edit import do_update_embedded_data
+from zerver.lib.mime_types import guess_type
+from zerver.lib.thumbnail import (
+    IMAGE_MAX_ANIMATED_PIXELS,
+    MarkdownImageMetadata,
+    StoredThumbnailFormat,
+    get_default_thumbnail_url,
+    get_image_thumbnail_path,
+    get_transcoded_format,
+    missing_thumbnails,
+    rewrite_thumbnailed_images,
+)
+from zerver.lib.upload import save_attachment_contents, store_message_attachment
+from zerver.models import ArchivedMessage, ImageAttachment, Message
+from zerver.worker.base import QueueProcessingWorker, assign_queue
+
+logger = logging.getLogger(__name__)
+
+
+@assign_queue("thumbnail")
+class ThumbnailWorker(QueueProcessingWorker):
+    @override
+    def consume(self, event: dict[str, Any]) -> None:
+        start_time = time.perf_counter()
+        with transaction.atomic(savepoint=False):
+            logger.info("Starting thumbnailing for %s", event.get("path_id", f"id: {event['id']}"))
+            try:
+                # This lock prevents us from racing with the on-demand
+                # rendering that can be triggered if a request is made
+                # directly to a thumbnail URL we have not made yet.
+                # This may mean that we may generate 0 thumbnail
+                # images once we get the lock.
+                # If the image data is bad and thumbnailing fails, the
+                # ImageAttachment row might be deleted; so we need a FOR UPDATE
+                # lock.
+                row = ImageAttachment.objects.select_for_update(of=("self",), no_key=False).get(
+                    id=event["id"]
+                )
+            except ImageAttachment.DoesNotExist:  # nocoverage
+                logger.info(
+                    "ImageAttachment row %d missing [%dms]",
+                    event["id"],
+                    (time.perf_counter() - start_time) * 1000,
+                )
+                return
+            lock_time = time.perf_counter() - start_time
+            result = ensure_thumbnails(row)
+            commit_time = time.perf_counter()
+        end_time = time.perf_counter()
+        logger.info(
+            "Processed %d thumbnails [%dms; lock=%dms, fetch=%dms, thumb=%dms, upload=%dms, rewrite=%dms, commit=%dms]",
+            result.generated_thumbnail_count,
+            (end_time - start_time) * 1000,
+            lock_time * 1000,
+            result.fetching_seconds * 1000,
+            result.processing_seconds * 1000,
+            result.uploading_seconds * 1000,
+            result.message_update_seconds * 1000,
+            (end_time - commit_time) * 1000,
+        )
+
+
+@dataclass
+class ThumbnailingResult:
+    generated_thumbnail_count: int = 0
+    fetching_seconds: float = 0
+    processing_seconds: float = 0
+    uploading_seconds: float = 0
+    message_update_seconds: float = 0
+
+
+def ensure_thumbnails(image_attachment: ImageAttachment) -> ThumbnailingResult:
+    needed_thumbnails = missing_thumbnails(image_attachment)
+
+    if not needed_thumbnails:
+        return ThumbnailingResult()
+
+    written_images = 0
+    with BytesIO() as f:
+        start_time = time.perf_counter()
+        save_attachment_contents(image_attachment.path_id, f)
+        image_bytes = f.getvalue()
+        fetching_seconds = time.perf_counter() - start_time
+    try:
+        # TODO: We could save some computational time by using the same
+        # bytes if multiple resolutions are larger than the source
+        # image.  That is, if the input is 10x10, a 100x100.jpg is
+        # going to be the same as a 200x200.jpg, since those set the
+        # max dimensions, and we do not scale up.
+        processing_seconds = 0.0
+        uploading_seconds = 0.0
+        for thumbnail_format in needed_thumbnails:
+            # This will scale to fit within the given dimensions; it
+            # may be smaller in one or more of them.
+            logger.debug(
+                "Resizing to %d x %d, from %d x %d",
+                thumbnail_format.max_width,
+                thumbnail_format.max_height,
+                image_attachment.original_width_px,
+                image_attachment.original_height_px,
+            )
+            load_opts = ""
+            if image_attachment.frames > 1:
+                # If the original has multiple frames, we want to load
+                # one of them if we're outputting to a static format,
+                # otherwise we load them all.
+                if thumbnail_format.animated:
+                    # We compute how many frames to thumbnail based on
+                    # how many frames it will take us to get to
+                    # IMAGE_MAX_ANIMATED_PIXELS
+                    pixels_per_frame = (
+                        image_attachment.original_width_px * image_attachment.original_height_px
+                    )
+                    if pixels_per_frame * image_attachment.frames < IMAGE_MAX_ANIMATED_PIXELS:
+                        load_opts = "n=-1"
+                    else:
+                        load_opts = f"n={IMAGE_MAX_ANIMATED_PIXELS // pixels_per_frame}"
+                else:
+                    load_opts = "n=1"
+            start_time = time.perf_counter()
+            resized = pyvips.Image.thumbnail_buffer(
+                image_bytes,
+                thumbnail_format.max_width,
+                height=thumbnail_format.max_height,
+                option_string=load_opts,
+                size=pyvips.Size.DOWN,
+            )
+            thumbnailed_bytes = resized.write_to_buffer(
+                f".{thumbnail_format.extension}[{thumbnail_format.opts}]"
+            )
+            processing_seconds += time.perf_counter() - start_time
+            content_type = guess_type(f"image.{thumbnail_format.extension}")[0]
+            assert content_type is not None
+            thumbnail_path = get_image_thumbnail_path(image_attachment, thumbnail_format)
+            logger.debug("Uploading %d bytes to %s", len(thumbnailed_bytes), thumbnail_path)
+            start_time = time.perf_counter()
+            store_message_attachment(
+                thumbnail_path,
+                str(thumbnail_format),
+                content_type,
+                thumbnailed_bytes,
+                None,
+                None,
+            )
+            uploading_seconds += time.perf_counter() - start_time
+            height = resized.get("page-height") if thumbnail_format.animated else resized.height
+            image_attachment.thumbnail_metadata.append(
+                asdict(
+                    StoredThumbnailFormat(
+                        extension=thumbnail_format.extension,
+                        content_type=content_type,
+                        max_width=thumbnail_format.max_width,
+                        max_height=thumbnail_format.max_height,
+                        animated=thumbnail_format.animated,
+                        width=resized.width,
+                        height=height,
+                        byte_size=len(thumbnailed_bytes),
+                    )
+                )
+            )
+            written_images += 1
+
+    except pyvips.Error as e:
+        logger.exception(
+            "Failed to thumbnail %s, %d byte %s @ %dx%d: %s",
+            image_attachment.path_id,
+            len(image_bytes),
+            image_attachment.content_type,
+            image_attachment.original_width_px,
+            image_attachment.original_height_px,
+            e.message,
+        )
+
+        if written_images == 0 and len(image_attachment.thumbnail_metadata) == 0:
+            # We have never thumbnailed this -- it most likely had
+            # bad data.  Remove the ImageAttachment row, since it is
+            # not valid for thumbnailing.
+            start_time = time.perf_counter()
+            update_message_rendered_content(
+                image_attachment.realm_id, image_attachment.path_id, None
+            )
+            message_update_seconds = time.perf_counter() - start_time
+            image_attachment.delete()
+            return ThumbnailingResult(
+                0, fetching_seconds, processing_seconds, uploading_seconds, message_update_seconds
+            )
+        else:  # nocoverage
+            # TODO: Clean up any dangling thumbnails we may have
+            # produced?  Seems unlikely that we'd fail on one size,
+            # but not another, but anything's possible.
+            pass
+
+    image_attachment.save(update_fields=["thumbnail_metadata"])
+    start_time = time.perf_counter()
+    url, is_animated = get_default_thumbnail_url(image_attachment)
+    update_message_rendered_content(
+        image_attachment.realm_id,
+        image_attachment.path_id,
+        MarkdownImageMetadata(
+            url=url,
+            is_animated=is_animated,
+            original_width_px=image_attachment.original_width_px,
+            original_height_px=image_attachment.original_height_px,
+            original_content_type=image_attachment.content_type,
+            transcoded_image=get_transcoded_format(image_attachment),
+        ),
+    )
+    message_update_seconds = time.perf_counter() - start_time
+    return ThumbnailingResult(
+        written_images,
+        fetching_seconds,
+        processing_seconds,
+        uploading_seconds,
+        message_update_seconds,
+    )
+
+
+def update_message_rendered_content(
+    realm_id: int, path_id: str, image_data: MarkdownImageMetadata | None
+) -> None:
+    for message_class in (Message, ArchivedMessage):
+        messages_with_image = (
+            message_class.objects.filter(realm_id=realm_id, attachment__path_id=path_id)
+            .select_for_update(of=("self",), no_key=True)
+            .order_by("id")
+        )
+        for message in messages_with_image:
+            assert message.rendered_content is not None
+            rendered_content = rewrite_thumbnailed_images(
+                message.rendered_content,
+                {} if image_data is None else {path_id: image_data},
+                {path_id} if image_data is None else set(),
+            )[0]
+            if rendered_content is None:
+                # There were no updates -- for instance, if we re-run
+                # ensure_thumbnails on an ImageAttachment we already
+                # ran it on once.  Do not bother to no-op update
+                # clients.
+                continue
+            if isinstance(message, Message):
+                # Perform a silent update push to the clients
+                do_update_embedded_data(message.sender, message, rendered_content)
+            else:
+                message.rendered_content = rendered_content
+                message.save(update_fields=["rendered_content"])