Merge pull request #8238 from lysnikolaou/arena-thread-safe

src/_imaging.c

@@ -92,6 +92,7 @@
 
 #define _USE_MATH_DEFINES
 #include <math.h>
+#include <stddef.h>
 
 /* Configuration stuff. Feel free to undef things you don't need. */
 #define WITH_IMAGECHOPS  /* ImageChops support */
@@ -3971,7 +3972,6 @@ static PyObject *
 _get_stats(PyObject *self, PyObject *args) {
     PyObject *d;
     PyObject *v;
-    ImagingMemoryArena arena = &ImagingDefaultArena;
 
     if (!PyArg_ParseTuple(args, ":get_stats")) {
         return NULL;
@@ -3981,6 +3981,10 @@ _get_stats(PyObject *self, PyObject *args) {
     if (!d) {
         return NULL;
     }
+
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
+    ImagingMemoryArena arena = &ImagingDefaultArena;
+
     v = PyLong_FromLong(arena->stats_new_count);
     PyDict_SetItemString(d, "new_count", v ? v : Py_None);
     Py_XDECREF(v);
@@ -4004,22 +4008,25 @@ _get_stats(PyObject *self, PyObject *args) {
     v = PyLong_FromLong(arena->blocks_cached);
     PyDict_SetItemString(d, "blocks_cached", v ? v : Py_None);
     Py_XDECREF(v);
+
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
     return d;
 }
 
 static PyObject *
 _reset_stats(PyObject *self, PyObject *args) {
-    ImagingMemoryArena arena = &ImagingDefaultArena;
-
     if (!PyArg_ParseTuple(args, ":reset_stats")) {
         return NULL;
     }
 
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
+    ImagingMemoryArena arena = &ImagingDefaultArena;
     arena->stats_new_count = 0;
     arena->stats_allocated_blocks = 0;
     arena->stats_reused_blocks = 0;
     arena->stats_reallocated_blocks = 0;
     arena->stats_freed_blocks = 0;
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
 
     Py_INCREF(Py_None);
     return Py_None;
@@ -4031,7 +4038,10 @@ _get_alignment(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    return PyLong_FromLong(ImagingDefaultArena.alignment);
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
+    int alignment = ImagingDefaultArena.alignment;
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    return PyLong_FromLong(alignment);
 }
 
 static PyObject *
@@ -4040,7 +4050,10 @@ _get_block_size(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    return PyLong_FromLong(ImagingDefaultArena.block_size);
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
+    int block_size = ImagingDefaultArena.block_size;
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    return PyLong_FromLong(block_size);
 }
 
 static PyObject *
@@ -4049,7 +4062,10 @@ _get_blocks_max(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    return PyLong_FromLong(ImagingDefaultArena.blocks_max);
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
+    int blocks_max = ImagingDefaultArena.blocks_max;
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    return PyLong_FromLong(blocks_max);
 }
 
 static PyObject *
@@ -4069,7 +4085,9 @@ _set_alignment(PyObject *self, PyObject *args) {
         return NULL;
     }
 
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
     ImagingDefaultArena.alignment = alignment;
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
 
     Py_INCREF(Py_None);
     return Py_None;
@@ -4092,7 +4110,9 @@ _set_block_size(PyObject *self, PyObject *args) {
         return NULL;
     }
 
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
     ImagingDefaultArena.block_size = block_size;
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
 
     Py_INCREF(Py_None);
     return Py_None;
@@ -4108,13 +4128,18 @@ _set_blocks_max(PyObject *self, PyObject *args) {
     if (blocks_max < 0) {
         PyErr_SetString(PyExc_ValueError, "blocks_max should be greater than 0");
         return NULL;
-    } else if ((unsigned long)blocks_max >
+    }
+
+    if ((unsigned long)blocks_max >
         SIZE_MAX / sizeof(ImagingDefaultArena.blocks_pool[0])) {
         PyErr_SetString(PyExc_ValueError, "blocks_max is too large");
         return NULL;
     }
 
-    if (!ImagingMemorySetBlocksMax(&ImagingDefaultArena, blocks_max)) {
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
+    int status = ImagingMemorySetBlocksMax(&ImagingDefaultArena, blocks_max);
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    if (!status) {
         return ImagingError_MemoryError();
     }
 
@@ -4130,7 +4155,9 @@ _clear_cache(PyObject *self, PyObject *args) {
         return NULL;
     }
 
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
     ImagingMemoryClearCache(&ImagingDefaultArena, i);
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
 
     Py_INCREF(Py_None);
     return Py_None;

src/libImaging/Imaging.h

@@ -161,6 +161,9 @@ typedef struct ImagingMemoryArena {
     int stats_reallocated_blocks; /* Number of blocks which were actually reallocated
                                      after retrieving */
     int stats_freed_blocks;       /* Number of freed blocks */
+#ifdef Py_GIL_DISABLED
+    PyMutex mutex;
+#endif
 } *ImagingMemoryArena;
 
 /* Objects */
@@ -710,6 +713,15 @@ _imaging_tell_pyFd(PyObject *fd);
 #include "ImagingUtils.h"
 extern UINT8 *clip8_lookups;
 
+/* Mutex lock/unlock helpers */
+#ifdef Py_GIL_DISABLED
+#define MUTEX_LOCK(m) PyMutex_Lock(m)
+#define MUTEX_UNLOCK(m) PyMutex_Unlock(m)
+#else
+#define MUTEX_LOCK(m)
+#define MUTEX_UNLOCK(m)
+#endif
+
 #if defined(__cplusplus)
 }
 #endif

src/libImaging/Storage.c

@@ -218,7 +218,9 @@ ImagingNewPrologueSubtype(const char *mode, int xsize, int ysize, int size) {
             break;
     }
 
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
     ImagingDefaultArena.stats_new_count += 1;
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
 
     return im;
 }
@@ -267,7 +269,10 @@ struct ImagingMemoryArena ImagingDefaultArena = {
     0,
     0,
     0,
-    0  // Stats
+    0,  // Stats
+#ifdef Py_GIL_DISABLED
+    {0},
+#endif
 };
 
 int
@@ -364,18 +369,19 @@ ImagingDestroyArray(Imaging im) {
     int y = 0;
 
     if (im->blocks) {
+        MUTEX_LOCK(&ImagingDefaultArena.mutex);
         while (im->blocks[y].ptr) {
             memory_return_block(&ImagingDefaultArena, im->blocks[y]);
             y += 1;
         }
+        MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
         free(im->blocks);
     }
 }
 
 Imaging
-ImagingAllocateArray(Imaging im, int dirty, int block_size) {
+ImagingAllocateArray(Imaging im, ImagingMemoryArena arena, int dirty, int block_size) {
     int y, line_in_block, current_block;
-    ImagingMemoryArena arena = &ImagingDefaultArena;
     ImagingMemoryBlock block = {NULL, 0};
     int aligned_linesize, lines_per_block, blocks_count;
     char *aligned_ptr = NULL;
@@ -498,14 +504,22 @@ ImagingNewInternal(const char *mode, int xsize, int ysize, int dirty) {
         return NULL;
     }
 
-    if (ImagingAllocateArray(im, dirty, ImagingDefaultArena.block_size)) {
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
+    Imaging tmp = ImagingAllocateArray(
+        im, &ImagingDefaultArena, dirty, ImagingDefaultArena.block_size
+    );
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    if (tmp) {
         return im;
     }
 
     ImagingError_Clear();
 
     // Try to allocate the image once more with smallest possible block size
-    if (ImagingAllocateArray(im, dirty, IMAGING_PAGE_SIZE)) {
+    MUTEX_LOCK(&ImagingDefaultArena.mutex);
+    tmp = ImagingAllocateArray(im, &ImagingDefaultArena, dirty, IMAGING_PAGE_SIZE);
+    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    if (tmp) {
         return im;
     }