Reduce memory arena contention

setup.py

@@ -14,12 +14,17 @@
 import struct
 import subprocess
 import sys
+import tempfile
 import warnings
 from collections.abc import Iterator
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
+from setuptools.errors import CompileError
+
+if TYPE_CHECKING:
+    import distutils.ccompiler
 
 
 def get_version() -> str:
@@ -292,6 +297,47 @@ def _pkg_config(name: str) -> tuple[list[str], list[str]] | None:
     return None
 
 
+def _try_compile(compiler: distutils.ccompiler.CCompiler, code: str) -> bool:
+    try:
+        with tempfile.TemporaryDirectory() as d:
+            fn = os.path.join(d, "test.c")
+            with open(fn, "w") as f:
+                f.write(code)
+            compiler.compile([fn], output_dir=d, extra_preargs=["-Werror"])
+        return True
+    except CompileError:
+        return False
+
+
+def _try_compile_attr(compiler: distutils.ccompiler.CCompiler, attr: str) -> bool:
+    code = f"""
+        #pragma GCC diagnostic error "-Wattributes"
+        #pragma clang diagnostic error "-Wattributes"
+
+        int {attr} foo;
+        int main() {{
+            return 0;
+        }}
+    """
+
+    return _try_compile(compiler, code)
+
+
+def _try_compile_tls_define_macros(
+    compiler: distutils.ccompiler.CCompiler,
+) -> list[tuple[str, str | None]]:
+    if _try_compile_attr(compiler, "thread_local"):  # C23
+        return [("HAVE_THREAD_LOCAL", None)]
+    elif _try_compile_attr(compiler, "_Thread_local"):  # C11/C17
+        return [("HAVE__THREAD_LOCAL", None)]
+    elif _try_compile_attr(compiler, "__thread"):  # GCC/clang
+        return [("HAVE___THREAD", None)]
+    elif _try_compile_attr(compiler, "__declspec(thread)"):  # MSVC
+        return [("HAVE___DECLSPEC_THREAD_", None)]
+    else:
+        return []
+
+
 class pil_build_ext(build_ext):
     class ext_feature:
         features = [
@@ -426,13 +472,14 @@ def finalize_options(self) -> None:
     def _update_extension(
         self,
         name: str,
-        libraries: list[str] | list[str | bool | None],
+        libraries: list[str] | list[str | bool | None] | None = None,
         define_macros: list[tuple[str, str | None]] | None = None,
         sources: list[str] | None = None,
     ) -> None:
         for extension in self.extensions:
             if extension.name == name:
-                extension.libraries += libraries
+                if libraries is not None:
+                    extension.libraries += libraries
                 if define_macros is not None:
                     extension.define_macros += define_macros
                 if sources is not None:
@@ -890,7 +937,10 @@ def build_extensions(self) -> None:
 
         defs.append(("PILLOW_VERSION", f'"{PILLOW_VERSION}"'))
 
-        self._update_extension("PIL._imaging", libs, defs)
+        tls_define_macros = _try_compile_tls_define_macros(self.compiler)
+        self._update_extension("PIL._imaging", libs, defs + tls_define_macros)
+        self._update_extension("PIL._imagingmath", define_macros=tls_define_macros)
+        self._update_extension("PIL._imagingmorph", define_macros=tls_define_macros)
 
         #
         # additional libraries
@@ -913,7 +963,9 @@ def build_extensions(self) -> None:
                         libs.append(feature.get("fribidi"))
                     else:  # building FriBiDi shim from src/thirdparty
                         srcs.append("src/thirdparty/fribidi-shim/fribidi.c")
-            self._update_extension("PIL._imagingft", libs, defs, srcs)
+            self._update_extension(
+                "PIL._imagingft", libs, defs + tls_define_macros, srcs
+            )
 
         else:
             self._remove_extension("PIL._imagingft")
@@ -922,19 +974,19 @@ def build_extensions(self) -> None:
             libs = [feature.get("lcms")]
             if sys.platform == "win32":
                 libs.extend(["user32", "gdi32"])
-            self._update_extension("PIL._imagingcms", libs)
+            self._update_extension("PIL._imagingcms", libs, tls_define_macros)
         else:
             self._remove_extension("PIL._imagingcms")
 
         webp = feature.get("webp")
         if isinstance(webp, str):
             libs = [webp, webp + "mux", webp + "demux"]
-            self._update_extension("PIL._webp", libs)
+            self._update_extension("PIL._webp", libs, tls_define_macros)
         else:
             self._remove_extension("PIL._webp")
 
         tk_libs = ["psapi"] if sys.platform in ("win32", "cygwin") else []
-        self._update_extension("PIL._imagingtk", tk_libs)
+        self._update_extension("PIL._imagingtk", tk_libs, tls_define_macros)
 
         build_ext.build_extensions(self)
 

src/_imaging.c

@@ -3822,34 +3822,49 @@ _get_stats(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    MUTEX_LOCK(&ImagingDefaultArena.mutex);
-    ImagingMemoryArena arena = &ImagingDefaultArena;
-
-    v = PyLong_FromLong(arena->stats_new_count);
+    long stats_new_count = 0;
+    long stats_allocated_blocks = 0;
+    long stats_reused_blocks = 0;
+    long stats_reallocated_blocks = 0;
+    long stats_freed_blocks = 0;
+    long blocks_cached = 0;
+
+    ImagingMemoryArena arena;
+    IMAGING_ARENAS_FOREACH(arena) {
+        MUTEX_LOCK(&arena->mutex);
+        stats_new_count += arena->stats_new_count;
+        stats_allocated_blocks += arena->stats_allocated_blocks;
+        stats_reused_blocks += arena->stats_reused_blocks;
+        stats_reallocated_blocks += arena->stats_reallocated_blocks;
+        stats_freed_blocks += arena->stats_freed_blocks;
+        blocks_cached += arena->blocks_cached;
+        MUTEX_UNLOCK(&arena->mutex);
+    }
+
+    v = PyLong_FromLong(stats_new_count);
     PyDict_SetItemString(d, "new_count", v ? v : Py_None);
     Py_XDECREF(v);
 
-    v = PyLong_FromLong(arena->stats_allocated_blocks);
+    v = PyLong_FromLong(stats_allocated_blocks);
     PyDict_SetItemString(d, "allocated_blocks", v ? v : Py_None);
     Py_XDECREF(v);
 
-    v = PyLong_FromLong(arena->stats_reused_blocks);
+    v = PyLong_FromLong(stats_reused_blocks);
     PyDict_SetItemString(d, "reused_blocks", v ? v : Py_None);
     Py_XDECREF(v);
 
-    v = PyLong_FromLong(arena->stats_reallocated_blocks);
+    v = PyLong_FromLong(stats_reallocated_blocks);
     PyDict_SetItemString(d, "reallocated_blocks", v ? v : Py_None);
     Py_XDECREF(v);
 
-    v = PyLong_FromLong(arena->stats_freed_blocks);
+    v = PyLong_FromLong(stats_freed_blocks);
     PyDict_SetItemString(d, "freed_blocks", v ? v : Py_None);
     Py_XDECREF(v);
 
-    v = PyLong_FromLong(arena->blocks_cached);
+    v = PyLong_FromLong(blocks_cached);
     PyDict_SetItemString(d, "blocks_cached", v ? v : Py_None);
     Py_XDECREF(v);
 
-    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
     return d;
 }
 
@@ -3859,14 +3874,16 @@ _reset_stats(PyObject *self, PyObject *args) {
         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);
+    ImagingMemoryArena arena;
+    IMAGING_ARENAS_FOREACH(arena) {
+        MUTEX_LOCK(&arena->mutex);
+        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(&arena->mutex);
+    }
 
     Py_RETURN_NONE;
 }
@@ -3877,9 +3894,10 @@ _get_alignment(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    MUTEX_LOCK(&ImagingDefaultArena.mutex);
-    int alignment = ImagingDefaultArena.alignment;
-    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    ImagingMemoryArena arena = ImagingGetArena();
+    MUTEX_LOCK(&arena->mutex);
+    int alignment = arena->alignment;
+    MUTEX_UNLOCK(&arena->mutex);
     return PyLong_FromLong(alignment);
 }
 
@@ -3889,9 +3907,10 @@ _get_block_size(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    MUTEX_LOCK(&ImagingDefaultArena.mutex);
-    int block_size = ImagingDefaultArena.block_size;
-    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    ImagingMemoryArena arena = ImagingGetArena();
+    MUTEX_LOCK(&arena->mutex);
+    int block_size = arena->block_size;
+    MUTEX_UNLOCK(&arena->mutex);
     return PyLong_FromLong(block_size);
 }
 
@@ -3901,9 +3920,10 @@ _get_blocks_max(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    MUTEX_LOCK(&ImagingDefaultArena.mutex);
-    int blocks_max = ImagingDefaultArena.blocks_max;
-    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    ImagingMemoryArena arena = ImagingGetArena();
+    MUTEX_LOCK(&arena->mutex);
+    int blocks_max = arena->blocks_max;
+    MUTEX_UNLOCK(&arena->mutex);
     return PyLong_FromLong(blocks_max);
 }
 
@@ -3924,9 +3944,12 @@ _set_alignment(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    MUTEX_LOCK(&ImagingDefaultArena.mutex);
-    ImagingDefaultArena.alignment = alignment;
-    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    ImagingMemoryArena arena;
+    IMAGING_ARENAS_FOREACH(arena) {
+        MUTEX_LOCK(&arena->mutex);
+        arena->alignment = alignment;
+        MUTEX_UNLOCK(&arena->mutex);
+    }
 
     Py_RETURN_NONE;
 }
@@ -3948,9 +3971,12 @@ _set_block_size(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    MUTEX_LOCK(&ImagingDefaultArena.mutex);
-    ImagingDefaultArena.block_size = block_size;
-    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    ImagingMemoryArena arena;
+    IMAGING_ARENAS_FOREACH(arena) {
+        MUTEX_LOCK(&arena->mutex);
+        arena->block_size = block_size;
+        MUTEX_UNLOCK(&arena->mutex);
+    }
 
     Py_RETURN_NONE;
 }
@@ -3968,15 +3994,20 @@ _set_blocks_max(PyObject *self, PyObject *args) {
     }
 
     if ((unsigned long)blocks_max >
-        SIZE_MAX / sizeof(ImagingDefaultArena.blocks_pool[0])) {
+        SIZE_MAX / sizeof(ImagingGetArena()->blocks_pool[0])) {
         PyErr_SetString(PyExc_ValueError, "blocks_max is too large");
         return NULL;
     }
 
-    MUTEX_LOCK(&ImagingDefaultArena.mutex);
-    int status = ImagingMemorySetBlocksMax(&ImagingDefaultArena, blocks_max);
-    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
-    if (!status) {
+    int error = 0;
+    ImagingMemoryArena arena;
+    IMAGING_ARENAS_FOREACH(arena) {
+        MUTEX_LOCK(&arena->mutex);
+        error |= ImagingMemorySetBlocksMax(arena, blocks_max);
+        MUTEX_UNLOCK(&arena->mutex);
+    }
+
+    if (error) {
         return ImagingError_MemoryError();
     }
 
@@ -3991,9 +4022,12 @@ _clear_cache(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    MUTEX_LOCK(&ImagingDefaultArena.mutex);
-    ImagingMemoryClearCache(&ImagingDefaultArena, i);
-    MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
+    ImagingMemoryArena arena;
+    IMAGING_ARENAS_FOREACH(arena) {
+        MUTEX_LOCK(&arena->mutex);
+        ImagingMemoryClearCache(arena, i);
+        MUTEX_UNLOCK(&arena->mutex);
+    }
 
     Py_RETURN_NONE;
 }