Put locks around relevant functions and use PyDict_GetItemRef

Author: emmatyping

Modules/_zstd/compressor.c

@@ -155,8 +155,9 @@ _get_CDict(ZstdDict *self, int compressionLevel)
 {
     assert(PyMutex_IsLocked(&self->lock));
     PyObject *level = NULL;
-    PyObject *capsule;
+    PyObject *capsule = NULL;
     ZSTD_CDict *cdict;
+    int ret;
 
 
     /* int level object */
@@ -166,11 +167,12 @@ _get_CDict(ZstdDict *self, int compressionLevel)
     }
 
     /* Get PyCapsule object from self->c_dicts */
-    capsule = PyDict_GetItemWithError(self->c_dicts, level);
+    ret = PyDict_GetItemRef(self->c_dicts, level, &capsule);
+    if (ret < 0) {
+        Py_XDECREF(capsule);
+        goto error;
+    }
     if (capsule == NULL) {
-        if (PyErr_Occurred()) {
-            goto error;
-        }
         /* Create ZSTD_CDict instance */
         char *dict_buffer = PyBytes_AS_STRING(self->dict_content);
         Py_ssize_t dict_len = Py_SIZE(self->dict_content);
@@ -197,9 +199,8 @@ _get_CDict(ZstdDict *self, int compressionLevel)
             goto error;
         }
 
-        /* Add PyCapsule object to self->c_dicts if not already inserted */
-        int ret = PyDict_SetDefaultRef(self->c_dicts, level, capsule, NULL);
-        Py_DECREF(capsule);
+        /* Add PyCapsule object to self->c_dicts */
+        ret = PyDict_SetItem(self->c_dicts, level, capsule);
         if (ret < 0) {
             goto error;
         }
@@ -214,6 +215,7 @@ _get_CDict(ZstdDict *self, int compressionLevel)
     cdict = NULL;
 success:
     Py_XDECREF(level);
+    Py_XDECREF(capsule);
     return cdict;
 }
 
@@ -422,6 +424,7 @@ static PyObject *
 compress_lock_held(ZstdCompressor *self, Py_buffer *data,
                    ZSTD_EndDirective end_directive)
 {
+    assert(PyMutex_IsLocked(&self->lock));
     ZSTD_inBuffer in;
     ZSTD_outBuffer out;
     _BlocksOutputBuffer buffer = {.list = NULL};
@@ -504,6 +507,7 @@ mt_continue_should_break(ZSTD_inBuffer *in, ZSTD_outBuffer *out)
 static PyObject *
 compress_mt_continue_lock_held(ZstdCompressor *self, Py_buffer *data)
 {
+    assert(PyMutex_IsLocked(&self->lock));
     ZSTD_inBuffer in;
     ZSTD_outBuffer out;
     _BlocksOutputBuffer buffer = {.list = NULL};

Modules/_zstd/decompressor.c

@@ -343,6 +343,7 @@ decompress_lock_held(ZstdDecompressor *self, ZSTD_inBuffer *in,
 static void
 decompressor_reset_session_lock_held(ZstdDecompressor *self)
 {
+    assert(PyMutex_IsLocked(&self->lock));
 
     /* Reset variables */
     self->in_begin = 0;
@@ -362,6 +363,7 @@ static PyObject *
 stream_decompress_lock_held(ZstdDecompressor *self, Py_buffer *data,
                             Py_ssize_t max_length)
 {
+    assert(PyMutex_IsLocked(&self->lock));
     ZSTD_inBuffer in;
     PyObject *ret = NULL;
     int use_input_buffer;