Repository: dask/distributed

Issue:
Zstandard compression is not thread safe
**Describe the issue**:

In `compression.py` we create a singleton `ZstdCompressor` and `ZstdDecompressor` object. The docs for the Zstandard library [state](https://github.com/indygreg/python-zstandard/blob/main/zstandard/backend_cffi.py#L1764):

> Unless specified otherwise, assume that no two methods of ``ZstdCompressor`` instances can be called from multiple Python threads simultaneously. In other words, assume instances are not thread safe unless stated otherwise.

**Minimal Complete Verifiable Example**:
This can trivially be reproduced by introducing threading to the `test_maybe_compress` test in `test_protocol.py`

Some errors that I've reproduced range from compression/decompression failure, to seg faults. Examples for searchability:

```
Fatal Python error: Segmentation fault

Thread 0x000000030df25000 (most recent call first):
  File "/Users/dwragge/code/distributed/distributed/protocol/compression.py", line 92 in zstd_compress
  File "/Users/dwragge/code/distributed/distributed/protocol/compression.py", line 186 in maybe_compress
  File "/Users/dwragge/code/distributed/distributed/protocol/tests/test_protocol.py", line 148 in test_compress_decompress
  File "/Users/dwragge/miniconda/envs/dask-distributed/lib/python3.8/threading.py", line 870 in run
  File "/Users/dwragge/miniconda/envs/dask-distributed/lib/python3.8/threading.py", line 932 in _bootstrap_inner
  File "/Users/dwragge/miniconda/envs/dask-distributed/lib/python3.8/threading.py", line 890 in _bootstrap
```

-  `zstd.ZstdError: cannot compress: Src size is incorrect`
- `zstd.ZstdError: cannot compress: Operation not authorized at current processing stage`
- `zstd.ZstdError: decompression error: Corrupted block detected`
- `zstd.ZstdError: error determining content size from frame header`
- `zstd.ZstdError: decompression error: Unknown frame descriptor`
- `zstd.ZstdError: decompression error: Destination buffer is too small`

This manifests in production as errors such as:

- `distributed.core - ERROR - Exception while handling op get_data`
- `distributed.worker - ERROR - failed during get data with tcp://a.b.c.d:39195 -> tcp://d.e.f.g:43787`
- `distributed.protocol.core - CRITICAL - Failed to deserialize`

Amongst others, including clear lost workers, and `CommClosedError`s.

**Anything else we need to know?**:

The fix is trivial (instantiate a new Zstd(De)Compressor instance for each compression. I will raise a PR to fix plus tests.


Discussion:

Resolving patch:
diff --git a/distributed/protocol/compression.py b/distributed/protocol/compression.py
--- a/distributed/protocol/compression.py
+++ b/distributed/protocol/compression.py
@@ -81,17 +81,15 @@
     if parse_version(zstandard.__version__) < parse_version("0.9.0"):
         raise ImportError("Need zstandard >= 0.9.0")
 
-    zstd_compressor = zstandard.ZstdCompressor(
-        level=dask.config.get("distributed.comm.zstd.level"),
-        threads=dask.config.get("distributed.comm.zstd.threads"),
-    )
-
-    zstd_decompressor = zstandard.ZstdDecompressor()
-
     def zstd_compress(data):
+        zstd_compressor = zstandard.ZstdCompressor(
+            level=dask.config.get("distributed.comm.zstd.level"),
+            threads=dask.config.get("distributed.comm.zstd.threads"),
+        )
         return zstd_compressor.compress(data)
 
     def zstd_decompress(data):
+        zstd_decompressor = zstandard.ZstdDecompressor()
         return zstd_decompressor.decompress(data)
 
     compressions["zstd"] = {"compress": zstd_compress, "decompress": zstd_decompress}


Test patch:
diff --git a/distributed/protocol/tests/test_protocol.py b/distributed/protocol/tests/test_protocol.py
--- a/distributed/protocol/tests/test_protocol.py
+++ b/distributed/protocol/tests/test_protocol.py
@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from threading import Thread
+from time import sleep
+
 import pytest
 
 import dask
@@ -121,6 +124,44 @@ def test_maybe_compress(lib, compression):
         assert compressions[rc]["decompress"](rd) == payload
 
 
+@pytest.mark.parametrize(
+    "lib,compression",
+    [(None, None), ("zlib", "zlib"), ("lz4", "lz4"), ("zstandard", "zstd")],
+)
+def test_compression_thread_safety(lib, compression):
+    if lib:
+        pytest.importorskip(lib)
+
+    try_converters = [bytes, memoryview]
+    start = (
+        False  # signal variable to increase likelihood of hitting thread-safety issues
+    )
+
+    def test_compress_decompress(fn):
+        while not start:
+            sleep(0.001)
+
+        payload = b"123"
+        assert maybe_compress(fn(payload), compression=compression) == (None, payload)
+
+        payload = b"0" * 10000
+        rc, rd = maybe_compress(fn(payload), compression=compression)
+        assert rc == compression
+        assert compressions[rc]["decompress"](rd) == payload
+
+    for fn in try_converters:
+        threads = []
+        for _ in range(0, 100):
+            start = False
+            for _ in range(0, 10):
+                thread = Thread(target=test_compress_decompress, args=(fn,))
+                thread.start()
+                threads.append(thread)
+            start = True
+            for thr in threads:
+                thr.join()
+
+
 @pytest.mark.parametrize(
     "lib,compression",
     [(None, None), ("zlib", "zlib"), ("lz4", "lz4"), ("zstandard", "zstd")],