Updates based on Sam's review

- Assert locks are held in `_get_(C/D)Dict` - Update self->d_dict outside `Py_BEGIN_ALLOW_THREADS` - removed critical section AC I missed
python · colesbury · May 23, 2025 · May 19, 2025 · May 19, 2025 · May 20, 2025
commit d142aa849f56972b20a8df1e588b5dd540d7276c
diff --git a/Modules/_zstd/clinic/decompressor.c.h b/Modules/_zstd/clinic/decompressor.c.h
diff --git a/Modules/_zstd/compressor.c b/Modules/_zstd/compressor.c
@@ -153,6 +153,7 @@ capsule_free_cdict(PyObject *capsule)
 ZSTD_CDict *
 _get_CDict(ZstdDict *self, int compressionLevel)
 {
+    assert(PyMutex_IsLocked(&self->lock));
     PyObject *level = NULL;
     PyObject *capsule;
     ZSTD_CDict *cdict;
@@ -197,11 +198,14 @@ _get_CDict(ZstdDict *self, int compressionLevel)
         }
 
         /* Add PyCapsule object to self->c_dicts if not already inserted */
-        if (PyDict_SetItem(self->c_dicts, level, capsule) < 0) {
-            Py_DECREF(capsule);
+        PyObject *capsule_dict;
+        int ret = PyDict_SetDefaultRef(self->c_dicts, level, capsule,
+                                       &capsule_dict);
+        Py_XDECREF(capsule_dict);
+        Py_DECREF(capsule);
+        if (ret < 0) {
             goto error;
         }
-        Py_DECREF(capsule);
     }
     else {
         /* ZSTD_CDict instance already exists */

diff --git a/Modules/_zstd/decompressor.c b/Modules/_zstd/decompressor.c
@@ -58,6 +58,7 @@ typedef struct {
 static inline ZSTD_DDict *
 _get_DDict(ZstdDict *self)
 {
+    assert(PyMutex_IsLocked(&self->lock));
     ZSTD_DDict *ret;
 
     /* Already created */
@@ -70,9 +71,9 @@ _get_DDict(ZstdDict *self)
         char *dict_buffer = PyBytes_AS_STRING(self->dict_content);
         Py_ssize_t dict_len = Py_SIZE(self->dict_content);
         Py_BEGIN_ALLOW_THREADS
-        self->d_dict = ZSTD_createDDict(dict_buffer,
-                                        dict_len);
+        ret = ZSTD_createDDict(dict_buffer, dict_len);
         Py_END_ALLOW_THREADS
+        self->d_dict = ret;
 
         if (self->d_dict == NULL) {
             _zstd_state* const mod_state = PyType_GetModuleState(Py_TYPE(self));
@@ -84,10 +85,7 @@ _get_DDict(ZstdDict *self)
         }
     }
 
-    /* Don't lose any exception */
-    ret = self->d_dict;
-
-    return ret;
+    return self->d_dict;
 }
 
 /* Set decompression parameters to decompression context */
@@ -629,7 +627,6 @@ ZstdDecompressor_dealloc(PyObject *ob)
 }
 
 /*[clinic input]
-@critical_section
 @getter
 _zstd.ZstdDecompressor.unused_data
 
@@ -641,7 +638,7 @@ decompressed, unused input data after the frame. Otherwise this will be b''.
 
 static PyObject *
 _zstd_ZstdDecompressor_unused_data_get_impl(ZstdDecompressor *self)
-/*[clinic end generated code: output=f3a20940f11b6b09 input=5233800bef00df04]*/
+/*[clinic end generated code: output=f3a20940f11b6b09 input=54d41ecd681a3444]*/
 {
     PyObject *ret;