Use critical sections to lock object state

This should avoid races and deadlocks.
python · emmatyping · Apr 26, 2025 · Apr 26, 2025 · Apr 26, 2025 · Apr 26, 2025
commit 02cd17a80f62c7df7c65fb7e613032c6311ea80a
diff --git a/Modules/_zstd/_zstdmodule.h b/Modules/_zstd/_zstdmodule.h
@@ -33,12 +33,6 @@ get_zstd_state(PyTypeObject *type) {
     return PyModule_GetState(module);
 }
 
-/* ------------------
-     Global macro
-   ------------------ */
-#define ACQUIRE_LOCK(obj) PyMutex_Lock(&obj->lock)
-#define RELEASE_LOCK(obj) PyMutex_Unlock(&obj->lock)
-
 extern 
10000
PyType_Spec zstddict_type_spec;
 extern PyType_Spec zstdcompressor_type_spec;
 extern PyType_Spec ZstdDecompressor_type_spec;
@@ -63,9 +57,6 @@ struct _zstd_state {
 typedef struct {
     PyObject_HEAD
 
-    /* Thread lock for generating ZSTD_CDict/ZSTD_DDict */
-    PyMutex lock;
-
     /* Reusable compress/decompress dictionary, they are created once and
        can be shared by multiple threads concurrently, since its usage is
        read-only.
@@ -85,9 +76,6 @@ typedef struct {
 typedef struct {
     PyObject_HEAD
 
-    /* Thread lock for compressing */
-    PyMutex lock;
-
     /* Compression context */
     ZSTD_CCtx *cctx;
 
@@ -110,9 +98,6 @@ typedef struct {
 typedef struct {
     PyObject_HEAD
 
-    /* Thread lock for compressing */
-    PyMutex lock;
-
     /* Decompression context */
     ZSTD_DCtx *dctx;
 

diff --git a/Modules/_zstd/compressor.c b/Modules/_zstd/compressor.c
@@ -135,7 +135,7 @@ _get_CDict(ZstdDict *self, int compressionLevel)
     PyObject *capsule;
     ZSTD_CDict *cdict;
 
-    ACQUIRE_LOCK(self);
+    Py_BEGIN_CRITICAL_SECTION(self);
 
     /* int level object */
     level = PyLong_FromLong(compressionLevel);
@@ -189,7 +189,7 @@ _get_CDict(ZstdDict *self, int compressionLevel)
     cdict = NULL;
 success:
     Py_XDECREF(level);
-    RELEASE_LOCK(self);
+    Py_END_CRITICAL_SECTION();
     return cdict;
 }
 
@@ -250,20 +250,26 @@ _PyZstd_load_c_dict(ZstdCompressor *self, PyObject *dict) {
         }
         /* Reference a prepared dictionary.
            It overrides some compression context's parameters. */
+        Py_BEGIN_CRITICAL_SECTION(self);
         zstd_ret = ZSTD_CCtx_refCDict(self->cctx, c_dict);
+        Py_END_CRITICAL_SECTION();
     } else if (type == DICT_TYPE_UNDIGESTED) {
         /* Load a dictionary.
            It doesn't override compression context's parameters. */
+        Py_BEGIN_CRITICAL_SECTION2(self, zd);
         zstd_ret = ZSTD_CCtx_loadDictionary(
                             self->cctx,
                             PyBytes_AS_STRING(zd->dict_content),
                             Py_SIZE(zd->dict_content));
+        Py_END_CRITICAL_SECTION2();
     } else if (type == DICT_TYPE_PREFIX) {
         /* Load a prefix */
+        Py_BEGIN_CRITICAL_SECTION2(self, zd);
         zstd_ret = ZSTD_CCtx_refPrefix(
                             self->cctx,
                             PyBytes_AS_STRING(zd->dict_content),
                             Py_SIZE(zd->dict_content));
+        Py_END_CRITICAL_SECTION2();
     } else {
         /* Impossible code path */
         PyErr_SetString(PyExc_SystemError,
@@ -327,11 +333,6 @@ ZstdCompressor_dealloc(ZstdCompressor *self)
     /* Py_XDECREF the dict after free the compression context */
     Py_XDECREF(self->dict);
 
-    /* Thread lock */
-    if (PyMutex_IsLocked(&self->lock)) {
-        PyMutex_Unlock(&self->lock);
-    }
-
     PyTypeObject *tp = Py_TYPE(self);
     tp->tp_free((PyObject*)self);
     Py_DECREF(tp);
@@ -564,7 +565,7 @@ _zstd_ZstdCompressor_compress_impl(ZstdCompressor *self, Py_buffer *data,
     }
 
     /* Thread-safe code */
-    ACQUIRE_LOCK(self);
+    Py_BEGIN_CRITICAL_SECTION(self);
 
     /* Compress */
     if (self->use_multithread && mode == ZSTD_e_continue) {
@@ -581,7 +582,7 @@ _zstd_ZstdCompressor_compress_impl(ZstdCompressor *self, Py_buffer *data,
         /* Resetting cctx's session never fail */
         ZSTD_CCtx_reset(self->cctx, ZSTD_reset_session_only);
     }
-    RELEASE_LOCK(self);
+    Py_END_CRITICAL_SECTION();
 
     return ret;
 }
@@ -616,7 +617,7 @@ _zstd_ZstdCompressor_flush_impl(ZstdCompressor *self, int mode)
     }
 
     /* Thread-safe code */
-    ACQUIRE_LOCK(self);
+    Py_BEGIN_CRITICAL_SECTION(self);
     ret = compress_impl(self, NULL, mode);
 
     if (ret) {
@@ -627,7 +628,7 @@ _zstd_ZstdCompressor_flush_impl(ZstdCompressor *self, int mode)
         /* Resetting cctx's session never fail */
         ZSTD_CCtx_reset(self->cctx, ZSTD_reset_session_only);
     }
-    RELEASE_LOCK(self);
+    Py_END_CRITICAL_SECTION();
 
     return ret;
 }

diff --git a/Modules/_zstd/decompressor.c b/Modules/_zstd/decompressor.c
@@ -38,7 +38,7 @@ _get_DDict(ZstdDict *self)
         return self->d_dict;
     }
 
-    ACQUIRE_LOCK(self);
+    Py_BEGIN_CRITICAL_SECTION(self);
     if (self->d_dict == NULL) {
         /* Create ZSTD_DDict instance from dictionary content */
         Py_BEGIN_ALLOW_THREADS
@@ -58,7 +58,7 @@ _get_DDict(ZstdDict *self)
 
     /* Don't lose any exception */
     ret = self->d_dict;
-    RELEASE_LOCK(self);
+    Py_END_CRITICAL_SECTION();
 
     return ret;
 }
@@ -175,19 +175,25 @@ _PyZstd_load_d_dict(ZstdDecompressor *self, PyObject *dict)
             return -1;
         }
         /* Reference a prepared dictionary */
+        Py_BEGIN_CRITICAL_SECTION(self);
         zstd_ret = ZSTD_DCtx_refDDict(self->dctx, d_dict);
+        Py_END_CRITICAL_SECTION();
     } else if (type == DICT_TYPE_UNDIGESTED) {
         /* Load a dictionary */
+        Py_BEGIN_CRITICAL_SECTION2(self, zd);
         zstd_ret = ZSTD_DCtx_loadDictionary(
                             self->dctx,
                             PyBytes_AS_STRING(zd->dict_content),
                             Py_SIZE(zd->dict_content));
+        Py_END_CRITICAL_SECTION2();
     } else if (type == DICT_TYPE_PREFIX) {
         /* Load a prefix */
+        Py_BEGIN_CRITICAL_SECTION2(self, zd);
         zstd_ret = ZSTD_DCtx_refPrefix(
                             self->dctx,
                             PyBytes_AS_STRING(zd->dict_content),
                             Py_SIZE(zd->dict_content));
+        Py_END_CRITICAL_SECTION2();
     } else {
         /* Impossible code path */
         PyErr_SetString(PyExc_SystemError,
@@ -397,7 +403,7 @@ stream_decompress(ZstdDecompressor *self, Py_buffer *data, Py_ssize_t max_length
     int use_input_buffer;
 
     /* Thread-safe code */
-    ACQUIRE_LOCK(self);
+    Py_BEGIN_CRITICAL_SECTION(self);
 
     if (type == TYPE_DECOMPRESSOR) {
         /* Check .eof flag */
@@ -584,7 +590,7 @@ stream_decompress(ZstdDecompressor *self, Py_buffer *data, Py_ssize_t max_length
 
     Py_CLEAR(ret);
 success:
-    RELEASE_LOCK(self);
+    Py_END_CRITICAL_SECTION();
 
     return ret;
 }
@@ -651,11 +657,6 @@ ZstdDecompressor_dealloc(ZstdDecompressor *self)
     /* Free unused data */
     Py_XDECREF(self->unused_data);
 
-    /* Thread lock */
-    if (PyMutex_IsLocked(&self->lock)) {
-        PyMutex_Unlock(&self->lock);
-    }
-
     PyTypeObject *tp = Py_TYPE(self);
     tp->tp_free((PyObject*)self);
     Py_DECREF(tp);
@@ -726,7 +727,7 @@ _zstd_ZstdDecompressor_unused_data_get_impl(ZstdDecompressor *self)
     PyObject *ret;
 
     /* Thread-safe code */
-    ACQUIRE_LOCK(self);
+    Py_BEGIN_CRITICAL_SECTION(self);
 
     if (!self->eof) {
         _zstd_state* const _module_state = PyType_GetModuleState(Py_TYPE(self));
@@ -748,7 +749,7 @@ _zstd_ZstdDecompressor_unused_data_get_impl(ZstdDecompressor *self)
         }
     }
 
-    RELEASE_LOCK(self);
+    Py_END_CRITICAL_SECTION();
 
     return ret;
 }

diff --git a/Modules/_zstd/zdict.c b/Modules/_zstd/zdict.c
@@ -65,11 +65,6 @@ ZstdDict_dealloc(ZstdDict *self)
     /* Release dict_content after Free ZSTD_CDict/ZSTD_DDict instances */
     Py_XDECREF(self->dict_content);
 
-    /* Thread lock */
-    if (PyMutex_IsLocked(&self->lock)) {
-        PyMutex_Unlock(&self->lock);
-    }
-
     PyTypeObject *tp = Py_TYPE(self);
     tp->tp_free((PyObject*)self);
     Py_DECREF(tp);