Răsfoiți Sursa

Images : une préparation annulée garde sa place jusqu'à la fin du décodage

Le sémaphore qui limite à une les images préparées à la fois était rendu à la
fin de l'attente (async with). Si la tâche qui attendait était annulée, la
place se libérait alors que le thread de prepare_image continuait à décoder :
l'image suivante démarrait en même temps, et la mémoire n'était plus bornée.

La place est maintenant rendue quand le thread finit (rappel sur la tâche,
attendue derrière asyncio.shield). L'erreur éventuelle est lue dans ce rappel,
pour ne pas être signalée « never retrieved » après une annulation.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Johnounet 1 zi în urmă
părinte
comite
cae0a50f98
2 a modificat fișierele cu 46 adăugiri și 6 ștergeri
  1. 11 2
      bot/attachments.py
  2. 35 4
      tests/test_attachments_prep.py

+ 11 - 2
bot/attachments.py

@@ -139,8 +139,17 @@ class AttachmentReader:
 
         # Réduit / convertit l'image si besoin (dans un thread : ne bloque pas le bot),
         # une à la fois pour borner la mémoire. Les téléchargements, eux, restent en parallèle.
-        async with self._image_slots:
-            return await asyncio.to_thread(prepare_image, data)
+        await self._image_slots.acquire()
+        preparing = asyncio.ensure_future(asyncio.to_thread(prepare_image, data))
+        # Place rendue à la fin du thread, pas de l'attente : si l'attente est annulée,
+        # le thread, lui, continue de décoder.
+        preparing.add_done_callback(self._release_image_slot)
+        return await asyncio.shield(preparing)
+
+    def _release_image_slot(self, preparing: asyncio.Future) -> None:
+        self._image_slots.release()
+        if not preparing.cancelled():
+            preparing.exception()  # lue : après une annulation, personne d'autre ne la lira
 
     async def download_images(self, attachments: list[discord.Attachment]) -> list[ImageData]:
         """Télécharge plusieurs images en parallèle et les prépare (lève ImageError si l'une pose problème).

+ 35 - 4
tests/test_attachments_prep.py

@@ -176,16 +176,47 @@ class ImagesAtOnceTest(unittest.IsolatedAsyncioTestCase):
                 running -= 1
             return ImageData("image/png", data)
 
-        reader = AttachmentReader(SimpleNamespace(max_image_size=10_000_000))
-        files = [SimpleNamespace(filename=f"{i}.png", content_type="image/png", size=10,
-                                 read=mock.AsyncMock(return_value=bytes([i])))
-                 for i in range(4)]
+        reader, files = self.reader_and_files(4)
         with mock.patch.object(attachments, "prepare_image", slow_prepare):
             results = await asyncio.gather(reader.download_images(files[:2]),
                                            reader.quoted_images(files[2:], "cité"))
         self.assertEqual(peak, attachments.MAX_IMAGES_AT_ONCE)  # tous messages confondus
         self.assertEqual([r.data for batch in results for r in batch], [b"\0", b"\1", b"\2", b"\3"])
 
+    async def test_slot_kept_until_cancelled_image_is_decoded(self):
+        # Annuler l'attente n'arrête pas le thread : sa place reste prise jusqu'à ce qu'il finisse.
+        started = []
+        release = threading.Event()
+
+        def blocking_prepare(data: bytes) -> ImageData:
+            started.append(data)
+            release.wait(5)
+            return ImageData("image/png", data)
+
+        reader, files = self.reader_and_files(2)
+        with mock.patch.object(attachments, "prepare_image", blocking_prepare):
+            try:
+                first = asyncio.create_task(reader.download_image(files[0]))
+                while not started:
+                    await asyncio.sleep(0.01)
+                first.cancel()
+                with self.assertRaises(asyncio.CancelledError):
+                    await first
+                second = asyncio.create_task(reader.download_image(files[1]))
+                await asyncio.sleep(0.1)
+                self.assertEqual(started, [b"\0"])  # la première décode encore : la deuxième attend
+            finally:
+                release.set()
+            self.assertEqual((await second).data, b"\1")
+
+    @staticmethod
+    def reader_and_files(count: int) -> tuple[AttachmentReader, list[SimpleNamespace]]:
+        reader = AttachmentReader(SimpleNamespace(max_image_size=10_000_000))
+        files = [SimpleNamespace(filename=f"{i}.png", content_type="image/png", size=10,
+                                 read=mock.AsyncMock(return_value=bytes([i])))
+                 for i in range(count)]
+        return reader, files
+
 
 def pdf(pages: int, password: str | None = None, owner_password: str | None = None) -> bytes:
     writer = PdfWriter()