Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Lock around usages of imaging memory arenas #8238

Merged
merged 8 commits into from
Aug 1, 2024
48 changes: 38 additions & 10 deletions src/_imaging.c
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand Down Expand Up @@ -3934,7 +3935,6 @@ static PyObject *
_get_stats(PyObject *self, PyObject *args) {
PyObject *d;
PyObject *v;
ImagingMemoryArena arena = &ImagingDefaultArena;

if (!PyArg_ParseTuple(args, ":get_stats")) {
return NULL;
Expand All @@ -3944,6 +3944,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);
Expand All @@ -3967,22 +3971,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;
Expand All @@ -3994,7 +4001,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 *
Expand All @@ -4003,7 +4013,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 *
Expand All @@ -4012,7 +4025,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 *
Expand All @@ -4032,7 +4048,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;
Expand All @@ -4055,7 +4073,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;
Expand All @@ -4071,14 +4091,20 @@ _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 >
SIZE_MAX / sizeof(ImagingDefaultArena.blocks_pool[0])) {
}

MUTEX_LOCK(&ImagingDefaultArena.mutex);
size_t blocksize = sizeof(ImagingDefaultArena.blocks_pool[0]);
lysnikolaou marked this conversation as resolved.
Show resolved Hide resolved
MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
if ((unsigned long)blocks_max > SIZE_MAX / blocksize) {
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();
}

Expand All @@ -4094,7 +4120,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;
Expand Down
12 changes: 12 additions & 0 deletions src/libImaging/Imaging.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand Down Expand Up @@ -690,6 +693,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
20 changes: 15 additions & 5 deletions src/libImaging/Storage.c
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,9 @@
break;
}

MUTEX_LOCK(&ImagingDefaultArena.mutex);
ImagingDefaultArena.stats_new_count += 1;
MUTEX_UNLOCK(&ImagingDefaultArena.mutex);

return im;
}
Expand Down Expand Up @@ -267,7 +269,8 @@
0,
0,
0,
0 // Stats
0, // Stats
{0},

Check failure on line 273 in src/libImaging/Storage.c

View workflow job for this annotation

GitHub Actions / ubuntu-latest Python 3.12

extra brace group at end of initializer

Check warning on line 273 in src/libImaging/Storage.c

View workflow job for this annotation

GitHub Actions / ubuntu-latest Python 3.12

excess elements in struct initializer
};

int
Expand Down Expand Up @@ -365,17 +368,18 @@

if (im->blocks) {
while (im->blocks[y].ptr) {
MUTEX_LOCK(&ImagingDefaultArena.mutex);
memory_return_block(&ImagingDefaultArena, im->blocks[y]);
MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
y += 1;
}
lysnikolaou marked this conversation as resolved.
Show resolved Hide resolved
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;
Expand Down Expand Up @@ -498,14 +502,20 @@
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;
}

Expand Down
Loading