From 0088faddddf59cad6f53f46a05d7705a5565502f Mon Sep 17 00:00:00 2001
From: miruka <miruka@disroot.org>
Date: Tue, 21 Jul 2020 23:35:16 -0400
Subject: [PATCH] Compress images in a separate process

Compression with Pillow can take long, especially with large
clipboard PNG images.
Doing this in a separate process prevents the async event loop from
getting blocked, and allows multiple compression operations to run in
parallel.
---
 TODO.md                      |  1 -
 src/backend/matrix_client.py | 25 ++++++++++++-------------
 src/backend/utils.py         | 28 ++++++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/TODO.md b/TODO.md
index 3cd42c4b..4d3ad263 100644
--- a/TODO.md
+++ b/TODO.md
@@ -4,7 +4,6 @@
     - hflickable: support kinetic scrolling disabler 
     - settings
 
-- compress png in a thread
 - verify upload cancellation
 - clipboard preview doesn't update when copied image changes until second time
 - Avatar tooltip can get displayed in front of presence menu
diff --git a/src/backend/matrix_client.py b/src/backend/matrix_client.py
index bed92840..a84ffd04 100644
--- a/src/backend/matrix_client.py
+++ b/src/backend/matrix_client.py
@@ -597,11 +597,11 @@ class MatrixClient(nio.AsyncClient):
         with NamedTemporaryFile(prefix=prefix, suffix=".png") as temp:
 
             async def get_path() -> Path:
-                with io.BytesIO(image) as inp, io.BytesIO() as buffer:
-                    PILImage.open(inp).save(buffer, "PNG")
+                # optimize is too slow for large images
+                compressed = await utils.compress_image(image, optimize=False)
 
-                    async with aiofiles.open(temp.name, "wb") as file:
-                        await file.write(buffer.getvalue())
+                async with aiofiles.open(temp.name, "wb") as file:
+                    await file.write(compressed)
 
                 return Path(temp.name)
 
@@ -1308,16 +1308,15 @@ class MatrixClient(nio.AsyncClient):
         if not small:
             thumb.thumbnail((800, 600))
 
-        with io.BytesIO() as out:
-            if thumb.mode in png_modes:
-                thumb.save(out, "PNG", optimize=True)
-                mime = "image/png"
-            else:
-                thumb.convert("RGB").save(out, "JPEG", optimize=True)
-                mime = "image/jpeg"
+        if thumb.mode in png_modes:
+            thumb_data = await utils.compress_image(thumb)
+            mime       = "image/png"
+        else:
+            thumb      = thumb.convert("RGB")
+            thumb_data = await utils.compress_image(thumb, "JPEG")
+            mime       = "image/jpeg"
 
-            thumb_data = out.getvalue()
-            thumb_size = len(thumb_data)
+        thumb_size = len(thumb_data)
 
         if thumb_size >= len(data) and is_jpg_png and not is_svg:
             raise UneededThumbnail()
diff --git a/src/backend/utils.py b/src/backend/utils.py
index 91e3b479..f2ce3efb 100644
--- a/src/backend/utils.py
+++ b/src/backend/utils.py
@@ -2,6 +2,7 @@
 
 """Various utilities that are used throughout the package."""
 
+import asyncio
 import collections
 import html
 import inspect
@@ -9,6 +10,7 @@ import io
 import json
 import sys
 import xml.etree.cElementTree as xml_etree  # FIXME: bandit warning
+from concurrent.futures import ProcessPoolExecutor
 from datetime import datetime, timedelta
 from enum import Enum
 from enum import auto as autostr
@@ -24,6 +26,8 @@ import aiofiles
 import filetype
 from aiofiles.threadpool.binary import AsyncBufferedReader
 from aiofiles.threadpool.text import AsyncTextIOWrapper
+from PIL import Image as PILImage
+
 from nio.crypto import AsyncDataT as File
 from nio.crypto import async_generator_from_data
 
@@ -34,8 +38,11 @@ else:
 
 AsyncOpenFile = Union[AsyncTextIOWrapper, AsyncBufferedReader]
 Size          = Tuple[int, int]
+BytesOrPIL    = Union[bytes, PILImage.Image]
 auto          = autostr
 
+COMPRESSION_POOL = ProcessPoolExecutor()
+
 
 class AutoStrEnum(Enum):
     """An Enum where auto() assigns the member's name instead of an integer.
@@ -236,3 +243,24 @@ async def atomic_write(
             temp_path.replace(path)
         else:
             temp_path.unlink()
+
+
+def _compress(image: BytesOrPIL, fmt: str, optimize: bool) -> bytes:
+    if isinstance(image, bytes):
+        pil_image = PILImage.open(io.BytesIO(image))
+    else:
+        pil_image = image
+
+    with io.BytesIO() as buffer:
+        pil_image.save(buffer, fmt, optimize=optimize)
+        return buffer.getvalue()
+
+
+async def compress_image(
+    image: BytesOrPIL, fmt: str = "PNG", optimize: bool = True,
+) -> bytes:
+    """Compress image in a separate process, without blocking event loop."""
+
+    return await asyncio.get_event_loop().run_in_executor(
+        COMPRESSION_POOL, _compress, image, fmt, optimize,
+    )