PYTHON-XXXX Reuse a zstd compressor per thread

pymongo/compression_support.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import sys
+import threading
 import warnings
 from collections.abc import Iterable
 from typing import Any, Optional, Union
@@ -154,6 +155,8 @@ def compress(self, data: bytes) -> bytes:
 
 class ZstdContext:
     compressor_id = 3
+    # One compressor per thread: creating a zstd context costs more than compressing a small message.
+    _local = threading.local()
 
     @staticmethod
     def compress(data: bytes) -> bytes:
@@ -162,7 +165,14 @@ def compress(data: bytes) -> bytes:
         else:
             from backports import zstd
 
-        return zstd.compress(data)
+        compressor = getattr(ZstdContext._local, "compressor", None)
+        if compressor is None:
+            compressor = ZstdContext._local.compressor = zstd.ZstdCompressor()
+        try:
+            return compressor.compress(data, zstd.ZstdCompressor.FLUSH_FRAME)
+        except BaseException:
+            ZstdContext._local.compressor = None
+            raise
 
 
 def _snappy_uncompressed_length(data: bytes | memoryview) -> int:

test/test_compression_support.py

@@ -219,6 +219,49 @@ def test_zstd_roundtrip(self):
         result = decompress(compressed, ZstdContext.compressor_id, max_message_size=2**20)
         self.assertEqual(result, data)
 
+    def test_zstd_compressor_reused_across_messages(self):
+        if not _have_zstd():
+            self.skipTest("zstd not available")
+        if sys.version_info >= (3, 14):
+            from compression import zstd
+        else:
+            from backports import zstd
+
+        messages = [bytes([i % 251]) * (1 + i * 997) for i in range(64)]
+        for data in messages:
+            compressed = ZstdContext.compress(data)
+            self.assertEqual(zstd.decompress(compressed), data)
+            self.assertEqual(
+                decompress(compressed, ZstdContext.compressor_id, max_message_size=2**20), data
+            )
+
+    def test_zstd_compress_from_many_threads(self):
+        if not _have_zstd():
+            self.skipTest("zstd not available")
+        import threading
+
+        errors = []
+
+        def worker(seed):
+            try:
+                for i in range(200):
+                    data = (b"%d-%d " % (seed, i)) * (1 + (seed * 31 + i) % 500)
+                    compressed = ZstdContext.compress(data)
+                    result = decompress(
+                        compressed, ZstdContext.compressor_id, max_message_size=2**20
+                    )
+                    if result != data:
+                        errors.append((seed, i))
+            except Exception as exc:
+                errors.append(exc)
+
+        threads = [threading.Thread(target=worker, args=(seed,)) for seed in range(8)]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+        self.assertEqual(errors, [])
+
 
 class TestDecompressSizeLimit(unittest.TestCase):
     def test_decompression_peak_memory_bounded(self):