Updates based on Sam's review

emmatyping · emmatyping · commit d142aa849f56 · 2025-05-21T16:52:49.000-04:00
- Assert locks are held in `_get_(C/D)Dict`
- Update self-&gt;d_dict outside `Py_BEGIN_ALLOW_THREADS`
- removed critical section AC I missed
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;
 

Original file line number	Diff line number	Diff line change
`@@ -153,6 +153,7 @@ capsule_free_cdict(PyObject *capsule)`
`153`	`153`	`ZSTD_CDict *`
`154`	`154`	`_get_CDict(ZstdDict *self, int compressionLevel)`
`155`	`155`	`{`
	`156`	`+ assert(PyMutex_IsLocked(&self->lock));`
`156`	`157`	`PyObject *level = NULL;`
`157`	`158`	`PyObject *capsule;`
`158`	`159`	`ZSTD_CDict *cdict;`
`@@ -197,11 +198,14 @@ _get_CDict(ZstdDict *self, int compressionLevel)`
`197`	`198`	`}`
`198`	`199`
`199`	`200`	`/* Add PyCapsule object to self->c_dicts if not already inserted */`
`200`		`- if (PyDict_SetItem(self->c_dicts, level, capsule) < 0) {`
`201`		`- Py_DECREF(capsule);`
	`201`	`+ PyObject *capsule_dict;`
	`202`	`+ int ret = PyDict_SetDefaultRef(self->c_dicts, level, capsule,`
	`203`	`+ &capsule_dict);`
	`204`	`+ Py_XDECREF(capsule_dict);`
	`205`	`+ Py_DECREF(capsule);`
	`206`	`+ if (ret < 0) {`
`202`	`207`	`goto error;`
`203`	`208`	`}`
`204`		`- Py_DECREF(capsule);`
`205`	`209`	`}`
`206`	`210`	`else {`
`207`	`211`	`/* ZSTD_CDict instance already exists */`