gh-131757: allow lru_cache functions to execute concurrently

Misc/NEWS.d/next/Library/2025-03-26-10-56-22.gh-issue-131757.pFRdmN.rst

@@ -0,0 +1 @@
+Make ``functools.lru_cache`` call the cached function unlocked to allow concurrency.

Modules/_functoolsmodule.c

@@ -4,6 +4,7 @@
 #include "pycore_long.h"          // _PyLong_GetZero()
 #include "pycore_moduleobject.h"  // _PyModule_GetState()
 #include "pycore_object.h"        // _PyObject_GC_TRACK
+#include "pycore_pyatomic_ft_wrappers.h"
 #include "pycore_pystate.h"       // _PyThreadState_GET()
 #include "pycore_tuple.h"         // _PyTuple_ITEMS()
 
@@ -1175,7 +1176,11 @@ uncached_lru_cache_wrapper(lru_cache_object *self, PyObject *args, PyObject *kwd
 {
     PyObject *result;
 
+#ifdef Py_GIL_DISABLED
+    _Py_atomic_add_ssize(&self->misses, 1);
+#else
     self->misses++;
+#endif
     result = PyObject_Call(self->func, args, kwds);
     if (!result)
         return NULL;
@@ -1195,18 +1200,25 @@ infinite_lru_cache_wrapper(lru_cache_object *self, PyObject *args, PyObject *kwd
         Py_DECREF(key);
         return NULL;
     }
-    result = _PyDict_GetItem_KnownHash(self->cache, key, hash);
-    if (result) {
-        Py_INCREF(result);
+    int res = _PyDict_GetItemRef_KnownHash((PyDictObject *)self->cache, key, hash, &result);
+    if (res > 0) {
+#ifdef Py_GIL_DISABLED
+        _Py_atomic_add_ssize(&self->hits, 1);
+#else
         self->hits++;
+#endif
         Py_DECREF(key);
         return result;
     }
-    if (PyErr_Occurred()) {
+    if (res < 0) {
         Py_DECREF(key);
         return NULL;
     }
+#ifdef Py_GIL_DISABLED
+    _Py_atomic_add_ssize(&self->misses, 1);
+#else
     self->misses++;
+#endif
     result = PyObject_Call(self->func, args, kwds);
     if (!result) {
         Py_DECREF(key);
@@ -1281,37 +1293,53 @@ lru_cache_prepend_link(lru_cache_object *self, lru_list_elem *link)
    so that we know the cache is a consistent state.
  */
 
-static PyObject *
-bounded_lru_cache_wrapper(lru_cache_object *self, PyObject *args, PyObject *kwds)
+static int
+bounded_lru_cache_get_lock_held(lru_cache_object *self, PyObject *args, PyObject *kwds,
+                                PyObject **result, PyObject **key, Py_hash_t *hash)
 {
     lru_list_elem *link;
-    PyObject *key, *result, *testresult;
-    Py_hash_t hash;
 
-    key = lru_cache_make_key(self->kwd_mark, args, kwds, self->typed);
-    if (!key)
-        return NULL;
-    hash = PyObject_Hash(key);
-    if (hash == -1) {
-        Py_DECREF(key);
-        return NULL;
+    PyObject *key_ = *key = lru_cache_make_key(self->kwd_mark, args, kwds, self->typed);
+    if (!key_)
+        return -1;
+    Py_hash_t hash_ = *hash = PyObject_Hash(key_);
+    if (hash_ == -1) {
+        Py_DECREF(key_);
+        return -1;
     }
-    link  = (lru_list_elem *)_PyDict_GetItem_KnownHash(self->cache, key, hash);
+    link = (lru_list_elem *)_PyDict_GetItem_KnownHash(self->cache, key_, hash_);
     if (link != NULL) {
         lru_cache_extract_link(link);
         lru_cache_append_link(self, link);
-        result = link->result;
+        *result = link->result;
+#ifdef Py_GIL_DISABLED
+        _Py_atomic_add_ssize(&self->hits, 1);
+#else
         self->hits++;
-        Py_INCREF(result);
-        Py_DECREF(key);
-        return result;
+#endif
+        Py_INCREF(link->result);
+        Py_DECREF(key_);
+        return 1;
     }
     if (PyErr_Occurred()) {
-        Py_DECREF(key);
-        return NULL;
+        Py_DECREF(key_);
+        return -1;
     }
+#ifdef Py_GIL_DISABLED
+    _Py_atomic_add_ssize(&self->misses, 1);
+#else
     self->misses++;
-    result = PyObject_Call(self->func, args, kwds);
+#endif
+    return 0;
+}
+
+static PyObject *
+bounded_lru_cache_update_lock_held(lru_cache_object *self,
+                                              PyObject *result, PyObject *key, Py_hash_t hash)
+{
+    lru_list_elem *link;
+    PyObject *testresult;
+
     if (!result) {
         Py_DECREF(key);
         return NULL;
@@ -1447,6 +1475,33 @@ bounded_lru_cache_wrapper(lru_cache_object *self, PyObject *args, PyObject *kwds
     return result;
 }
 
+static PyObject *
+bounded_lru_cache_wrapper(lru_cache_object *self, PyObject *args, PyObject *kwds)
+{
+    PyObject *key, *result;
+    Py_hash_t hash;
+    int res;
+
+    Py_BEGIN_CRITICAL_SECTION(self);
+    res = bounded_lru_cache_get_lock_held(self, args, kwds, &result, &key, &hash);
+    Py_END_CRITICAL_SECTION();
+
+    if (res < 0) {
+        return NULL;
+    }
+    if (res > 0) {
+        return result;
+    }
+
+    result = PyObject_Call(self->func, args, kwds);
+
+    Py_BEGIN_CRITICAL_SECTION(self);
+    result = bounded_lru_cache_update_lock_held(self, result, key, hash);
+    Py_END_CRITICAL_SECTION();
+
+    return result;
+}
+
 static PyObject *
 lru_cache_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 {
@@ -1579,9 +1634,7 @@ lru_cache_call(PyObject *op, PyObject *args, PyObject *kwds)
 {
     lru_cache_object *self = lru_cache_object_CAST(op);
     PyObject *result;
-    Py_BEGIN_CRITICAL_SECTION(self);
     result = self->wrapper(self, args, kwds);
-    Py_END_CRITICAL_SECTION();
     return result;
 }
 
@@ -1608,11 +1661,15 @@ _functools__lru_cache_wrapper_cache_info_impl(PyObject *self)
     lru_cache_object *_self = (lru_cache_object *) self;
     if (_self->maxsize == -1) {
         return PyObject_CallFunction(_self->cache_info_type, "nnOn",
-                                     _self->hits, _self->misses, Py_None,
+                                     FT_ATOMIC_LOAD_SSIZE_RELAXED(_self->hits),
+                                     FT_ATOMIC_LOAD_SSIZE_RELAXED(_self->misses),
+                                     Py_None,
                                      PyDict_GET_SIZE(_self->cache));
     }
     return PyObject_CallFunction(_self->cache_info_type, "nnnn",
-                                 _self->hits, _self->misses, _self->maxsize,
+                                 FT_ATOMIC_LOAD_SSIZE_RELAXED(_self->hits),
+                                 FT_ATOMIC_LOAD_SSIZE_RELAXED(_self->misses),
+                                 _self->maxsize,
                                  PyDict_GET_SIZE(_self->cache));
 }
 
@@ -1629,7 +1686,8 @@ _functools__lru_cache_wrapper_cache_clear_impl(PyObject *self)
 {
     lru_cache_object *_self = (lru_cache_object *) self;
     lru_list_elem *list = lru_cache_unlink_list(_self);
-    _self->hits = _self->misses = 0;
+    FT_ATOMIC_STORE_SSIZE_RELAXED(_self->hits, 0);
+    FT_ATOMIC_STORE_SSIZE_RELAXED(_self->misses, 0);
     PyDict_Clear(_self->cache);
     lru_cache_clear_list(list);
     Py_RETURN_NONE;