gh-120973: Fix thread-safety issues with `threading.local` (#121655)

This is a small refactoring to the current design that allows us to
avoid manually iterating over threads.

This should also fix gh-118490.

Include/cpython/pystate.h

@@ -192,6 +192,14 @@ struct _ts {
     PyObject *previous_executor;
 
     uint64_t dict_global_version;
+
+    /* Used to store/retrieve `threading.local` keys/values for this thread */
+    PyObject *threading_local_key;
+
+    /* Used by `threading.local`s to be remove keys/values for dying threads.
+       The PyThreadObject must hold the only reference to this value.
+    */
+    PyObject *threading_local_sentinel;
 };
 
 #ifdef Py_DEBUG

Modules/_threadmodule.c

@@ -1350,33 +1350,44 @@ newlockobject(PyObject *module)
    Our implementation uses small "localdummy" objects in order to break
    the reference chain. These trivial objects are hashable (using the
    default scheme of identity hashing) and weakrefable.
-   Each thread-state holds a separate localdummy for each local object
+
-   (as a /strong reference/),
+   Each thread-state holds two separate localdummy objects:
-   and each thread-local object holds a dict mapping /weak references/
+
-   of localdummies to local dicts.
+   - `threading_local_key` is used as a key to retrieve the locals dictionary
+     for the thread in any `threading.local` object.
+   - `threading_local_sentinel` is used to signal when a thread is being
+     destroyed. Consequently, the associated thread-state must hold the only
+     reference.
+
+   Each `threading.local` object contains a dict mapping localdummy keys to
+   locals dicts and a set containing weak references to localdummy
+   sentinels. Each sentinel weak reference has a callback that removes itself
+   and the locals dict for the key from the `threading.local` object when
+   called.
 
    Therefore:
-   - only the thread-state dict holds a strong reference to the dummies
+   - The thread-state only holds strong references to localdummy objects, which
-   - only the thread-local object holds a strong reference to the local dicts
+     cannot participate in cycles.
-   - only outside objects (application- or library-level) hold strong
+   - Only outside objects (application- or library-level) hold strong
-     references to the thread-local objects
+     references to the thread-local objects.
-   - as soon as a thread-state dict is destroyed, the weakref callbacks of all
+   - As soon as thread-state's sentinel dummy is destroyed the callbacks for
-     dummies attached to that thread are called, and destroy the corresponding
+     all weakrefs attached to the sentinel are called, and destroy the
-     local dicts from thread-local objects
+     corresponding local dicts from thread-local objects.
-   - as soon as a thread-local object is destroyed, its local dicts are
+   - As soon as a thread-local object is destroyed, its local dicts are
-     destroyed and its dummies are manually removed from all thread states
+     destroyed.
-   - the GC can do its work correctly when a thread-local object is dangling,
+   - The GC can do its work correctly when a thread-local object is dangling,
-     without any interference from the thread-state dicts
+     without any interference from the thread-state dicts.
 
-   As an additional optimization, each localdummy holds a borrowed reference
+   This dual key arrangement is necessary to ensure that `threading.local`
-   to the corresponding localdict.  This borrowed reference is only used
+   values can be retrieved from finalizers. If we were to only keep a mapping
-   by the thread-local object which has created the localdummy, which should
+   of localdummy weakrefs to locals dicts it's possible that the weakrefs would
-   guarantee that the localdict still exists when accessed.
+   be cleared before finalizers were called (GC currently clears weakrefs that
+   are garbage before invoking finalizers), causing lookups in finalizers to
+   fail.
 */
 
 typedef struct {
     PyObject_HEAD
-    PyObject *localdict;        /* Borrowed reference! */
     PyObject *weakreflist;      /* List of weak references to self */
 } localdummyobject;
 
@@ -1413,80 +1424,60 @@ static PyType_Spec local_dummy_type_spec = {
 
 typedef struct {
     PyObject_HEAD
-    PyObject *key;
     PyObject *args;
     PyObject *kw;
     PyObject *weakreflist;      /* List of weak references to self */
-    /* A {localdummy weakref -> localdict} dict */
+    /* A {localdummy -> localdict} dict */
-    PyObject *dummies;
+    PyObject *localdicts;
-    /* The callback for weakrefs to localdummies */
+    /* A set of weakrefs to thread sentinels localdummies*/
-    PyObject *wr_callback;
+    PyObject *thread_watchdogs;
 } localobject;
 
 /* Forward declaration */
-static PyObject *_ldict(localobject *self, thread_module_state *state);
+static int create_localsdict(localobject *self, thread_module_state *state,
-static PyObject *_localdummy_destroyed(PyObject *meth_self, PyObject *dummyweakref);
+                             PyObject **localsdict, PyObject **sentinel_wr);
+static PyObject *clear_locals(PyObject *meth_self, PyObject *dummyweakref);
 
-/* Create and register the dummy for the current thread.
+/* Create a weakref to the sentinel localdummy for the current thread */
-   Returns a borrowed reference of the corresponding local dict */
 static PyObject *
-_local_create_dummy(localobject *self, thread_module_state *state)
+create_sentinel_wr(localobject *self)
 {
-    PyObject *ldict = NULL, *wr = NULL;
+    static PyMethodDef wr_callback_def = {
-    localdummyobject *dummy = NULL;
+        "clear_locals", (PyCFunction) clear_locals, METH_O
-    PyTypeObject *type = state->local_dummy_type;
+    };
 
-    PyObject *tdict = PyThreadState_GetDict();
+    PyThreadState *tstate = PyThreadState_Get();
-    if (tdict == NULL) {
+
-        PyErr_SetString(PyExc_SystemError,
+    /* We use a weak reference to self in the callback closure
-                        "Couldn't get thread-state dictionary");
+       in order to avoid spurious reference cycles */
-        goto err;
+    PyObject *self_wr = PyWeakref_NewRef((PyObject *) self, NULL);
+    if (self_wr == NULL) {
+        return NULL;
     }
 
-    ldict = PyDict_New();
+    PyObject *args = PyTuple_New(2);
-    if (ldict == NULL) {
+    if (args == NULL) {
-        goto err;
+        Py_DECREF(self_wr);
+        return NULL;
     }
-    dummy = (localdummyobject *) type->tp_alloc(type, 0);
+    PyTuple_SET_ITEM(args, 0, self_wr);
-    if (dummy == NULL) {
+    PyTuple_SET_ITEM(args, 1, Py_NewRef(tstate->threading_local_key));
-        goto err;
+
-    }
+    PyObject *cb = PyCFunction_New(&wr_callback_def, args);
-    dummy->localdict = ldict;
+    Py_DECREF(args);
-    wr = PyWeakref_NewRef((PyObject *) dummy, self->wr_callback);
+    if (cb == NULL) {
-    if (wr == NULL) {
+        return NULL;
-        goto err;
     }
 
-    /* As a side-effect, this will cache the weakref's hash before the
+    PyObject *wr = PyWeakref_NewRef(tstate->threading_local_sentinel, cb);
-       dummy gets deleted */
+    Py_DECREF(cb);
-    int r = PyDict_SetItem(self->dummies, wr, ldict);
-    if (r < 0) {
-        goto err;
-    }
-    Py_CLEAR(wr);
-    r = PyDict_SetItem(tdict, self->key, (PyObject *) dummy);
-    if (r < 0) {
-        goto err;
-    }
-    Py_CLEAR(dummy);
 
-    Py_DECREF(ldict);
+    return wr;
-    return ldict;
-
-err:
-    Py_XDECREF(ldict);
-    Py_XDECREF(wr);
-    Py_XDECREF(dummy);
-    return NULL;
 }
 
 static PyObject *
 local_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 {
-    static PyMethodDef wr_callback_def = {
-        "_localdummy_destroyed", (PyCFunction) _localdummy_destroyed, METH_O
-    };
-
     if (type->tp_init == PyBaseObject_Type.tp_init) {
         int rc = 0;
         if (args != NULL)
@@ -1513,30 +1504,25 @@ local_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 
     self->args = Py_XNewRef(args);
     self->kw = Py_XNewRef(kw);
-    self->key = PyUnicode_FromFormat("thread.local.%p", self);
+
-    if (self->key == NULL) {
+    self->localdicts = PyDict_New();
+    if (self->localdicts == NULL) {
         goto err;
     }
 
-    self->dummies = PyDict_New();
+    self->thread_watchdogs = PySet_New(NULL);
-    if (self->dummies == NULL) {
+    if (self->thread_watchdogs == NULL) {
         goto err;
     }
 
-    /* We use a weak reference to self in the callback closure
+    PyObject *localsdict = NULL;
-       in order to avoid spurious reference cycles */
+    PyObject *sentinel_wr = NULL;
-    PyObject *wr = PyWeakref_NewRef((PyObject *) self, NULL);
+    if (create_localsdict(self, state, &localsdict, &sentinel_wr) < 0) {
-    if (wr == NULL) {
-        goto err;
-    }
-    self->wr_callback = PyCFunction_NewEx(&wr_callback_def, wr, NULL);
-    Py_DECREF(wr);
-    if (self->wr_callback == NULL) {
-        goto err;
-    }
-    if (_local_create_dummy(self, state) == NULL) {
         goto err;
     }
+    Py_DECREF(localsdict);
+    Py_DECREF(sentinel_wr);
+
     return (PyObject *)self;
 
   err:
@@ -1550,7 +1536,8 @@ local_traverse(localobject *self, visitproc visit, void *arg)
     Py_VISIT(Py_TYPE(self));
     Py_VISIT(self->args);
     Py_VISIT(self->kw);
-    Py_VISIT(self->dummies);
+    Py_VISIT(self->localdicts);
+    Py_VISIT(self->thread_watchdogs);
     return 0;
 }
 
@@ -1559,27 +1546,8 @@ local_clear(localobject *self)
 {
     Py_CLEAR(self->args);
     Py_CLEAR(self->kw);
-    Py_CLEAR(self->dummies);
+    Py_CLEAR(self->localdicts);
-    Py_CLEAR(self->wr_callback);
+    Py_CLEAR(self->thread_watchdogs);
-    /* Remove all strong references to dummies from the thread states */
-    if (self->key) {
-        PyInterpreterState *interp = _PyInterpreterState_GET();
-        _PyRuntimeState *runtime = &_PyRuntime;
-        HEAD_LOCK(runtime);
-        PyThreadState *tstate = PyInterpreterState_ThreadHead(interp);
-        HEAD_UNLOCK(runtime);
-        while (tstate) {
-            if (tstate->dict) {
-                if (PyDict_Pop(tstate->dict, self->key, NULL) < 0) {
-                    // Silently ignore error
-                    PyErr_Clear();
-                }
-            }
-            HEAD_LOCK(runtime);
-            tstate = PyThreadState_Next(tstate);
-            HEAD_UNLOCK(runtime);
-        }
-    }
     return 0;
 }
 
@@ -1595,48 +1563,142 @@ local_dealloc(localobject *self)
     PyObject_GC_UnTrack(self);
 
     local_clear(self);
-    Py_XDECREF(self->key);
 
     PyTypeObject *tp = Py_TYPE(self);
     tp->tp_free((PyObject*)self);
     Py_DECREF(tp);
 }
 
-/* Returns a borrowed reference to the local dict, creating it if necessary */
+/* Create the TLS key and sentinel if they don't exist */
+static int
+create_localdummies(thread_module_state *state)
+{
+    PyThreadState *tstate = _PyThreadState_GET();
+
+    if (tstate->threading_local_key != NULL) {
+        return 0;
+    }
+
+    PyTypeObject *ld_type = state->local_dummy_type;
+    tstate->threading_local_key = ld_type->tp_alloc(ld_type, 0);
+    if (tstate->threading_local_key == NULL) {
+        return -1;
+    }
+
+    tstate->threading_local_sentinel = ld_type->tp_alloc(ld_type, 0);
+    if (tstate->threading_local_sentinel == NULL) {
+        Py_CLEAR(tstate->threading_local_key);
+        return -1;
+    }
+
+    return 0;
+}
+
+/* Insert a localsdict and sentinel weakref for the current thread, placing
+   strong references in localsdict and sentinel_wr, respectively.
+*/
+static int
+create_localsdict(localobject *self, thread_module_state *state,
+                  PyObject **localsdict, PyObject **sentinel_wr)
+{
+    PyThreadState *tstate = _PyThreadState_GET();
+    PyObject *ldict = NULL;
+    PyObject *wr = NULL;
+
+    if (create_localdummies(state) < 0) {
+        goto err;
+    }
+
+    /* Create and insert the locals dict and sentinel weakref */
+    ldict = PyDict_New();
+    if (ldict == NULL) {
+        goto err;
+    }
+
+    if (PyDict_SetItem(self->localdicts, tstate->threading_local_key, ldict) <
+        0) {
+        goto err;
+    }
+
+    wr = create_sentinel_wr(self);
+    if (wr == NULL) {
+        PyObject *exc = PyErr_GetRaisedException();
+        if (PyDict_DelItem(self->localdicts, tstate->threading_local_key) <
+            0) {
+            PyErr_WriteUnraisable((PyObject *)self);
+        }
+        PyErr_SetRaisedException(exc);
+        goto err;
+    }
+
+    if (PySet_Add(self->thread_watchdogs, wr) < 0) {
+        PyObject *exc = PyErr_GetRaisedException();
+        if (PyDict_DelItem(self->localdicts, tstate->threading_local_key) <
+            0) {
+            PyErr_WriteUnraisable((PyObject *)self);
+        }
+        PyErr_SetRaisedException(exc);
+        goto err;
+    }
+
+    *localsdict = ldict;
+    *sentinel_wr = wr;
+    return 0;
+
+err:
+    Py_XDECREF(ldict);
+    Py_XDECREF(wr);
+    return -1;
+}
+
+/* Return a strong reference to the locals dict for the current thread,
+   creating it if necessary.
+*/
 static PyObject *
 _ldict(localobject *self, thread_module_state *state)
 {
-    PyObject *tdict = PyThreadState_GetDict();
+    if (create_localdummies(state) < 0) {
-    if (tdict == NULL) {
-        PyErr_SetString(PyExc_SystemError,
-                        "Couldn't get thread-state dictionary");
         return NULL;
     }
 
+    /* Check if a localsdict already exists */
     PyObject *ldict;
-    PyObject *dummy = PyDict_GetItemWithError(tdict, self->key);
+    PyThreadState *tstate = _PyThreadState_GET();
-    if (dummy == NULL) {
+    if (PyDict_GetItemRef(self->localdicts, tstate->threading_local_key,
-        if (PyErr_Occurred()) {
+                          &ldict) < 0) {
-            return NULL;
+        return NULL;
-        }
+    }
-        ldict = _local_create_dummy(self, state);
+    if (ldict != NULL) {
-        if (ldict == NULL)
+        return ldict;
-            return NULL;
+    }
 
-        if (Py_TYPE(self)->tp_init != PyBaseObject_Type.tp_init &&
+    /* threading.local hasn't been instantiated for this thread */
-            Py_TYPE(self)->tp_init((PyObject*)self,
+    PyObject *wr;
-                                   self->args, self->kw) < 0) {
+    if (create_localsdict(self, state, &ldict, &wr) < 0) {
-            /* we need to get rid of ldict from thread so
+        return NULL;
-               we create a new one the next time we do an attr
+    }
-               access */
+
-            PyDict_DelItem(tdict, self->key);
+    /* run __init__ if we're a subtype of `threading.local` */
-            return NULL;
+    if (Py_TYPE(self)->tp_init != PyBaseObject_Type.tp_init &&
+        Py_TYPE(self)->tp_init((PyObject *)self, self->args, self->kw) < 0) {
+        /* we need to get rid of ldict from thread so
+           we create a new one the next time we do an attr
+           access */
+        PyObject *exc = PyErr_GetRaisedException();
+        if (PyDict_DelItem(self->localdicts, tstate->threading_local_key) <
+            0) {
+            PyErr_WriteUnraisable((PyObject *)self);
+            PyErr_Clear();
         }
+        if (PySet_Discard(self->thread_watchdogs, wr) < 0) {
+            PyErr_WriteUnraisable((PyObject *)self);
+        }
+        PyErr_SetRaisedException(exc);
+        Py_DECREF(ldict);
+        Py_DECREF(wr);
+        return NULL;
     }
-    else {
+    Py_DECREF(wr);
-        assert(Py_IS_TYPE(dummy, state->local_dummy_type));
-        ldict = ((localdummyobject *) dummy)->localdict;
-    }
 
     return ldict;
 }
@@ -1650,21 +1712,28 @@ local_setattro(localobject *self, PyObject *name, PyObject *v)
 
     PyObject *ldict = _ldict(self, state);
     if (ldict == NULL) {
-        return -1;
+        goto err;
     }
 
     int r = PyObject_RichCompareBool(name, &_Py_ID(__dict__), Py_EQ);
     if (r == -1) {
-        return -1;
+        goto err;
     }
     if (r == 1) {
         PyErr_Format(PyExc_AttributeError,
                      "'%.100s' object attribute '%U' is read-only",
                      Py_TYPE(self)->tp_name, name);
-        return -1;
+        goto err;
     }
 
-    return _PyObject_GenericSetAttrWithDict((PyObject *)self, name, v, ldict);
+    int st =
+        _PyObject_GenericSetAttrWithDict((PyObject *)self, name, v, ldict);
+    Py_DECREF(ldict);
+    return st;
+
+err:
+    Py_XDECREF(ldict);
+    return -1;
 }
 
 static PyObject *local_getattro(localobject *, PyObject *);
@@ -1707,34 +1776,42 @@ local_getattro(localobject *self, PyObject *name)
 
     int r = PyObject_RichCompareBool(name, &_Py_ID(__dict__), Py_EQ);
     if (r == 1) {
-        return Py_NewRef(ldict);
+        return ldict;
     }
     if (r == -1) {
+        Py_DECREF(ldict);
         return NULL;
     }
 
     if (!Py_IS_TYPE(self, state->local_type)) {
         /* use generic lookup for subtypes */
-        return _PyObject_GenericGetAttrWithDict((PyObject *)self, name,
+        PyObject *res =
-                                                ldict, 0);
+            _PyObject_GenericGetAttrWithDict((PyObject *)self, name, ldict, 0);
+        Py_DECREF(ldict);
+        return res;
     }
 
     /* Optimization: just look in dict ourselves */
     PyObject *value;
     if (PyDict_GetItemRef(ldict, name, &value) != 0) {
         // found or error
+        Py_DECREF(ldict);
         return value;
     }
 
     /* Fall back on generic to get __class__ and __dict__ */
-    return _PyObject_GenericGetAttrWithDict(
+    PyObject *res =
-        (PyObject *)self, name, ldict, 0);
+        _PyObject_GenericGetAttrWithDict((PyObject *)self, name, ldict, 0);
+    Py_DECREF(ldict);
+    return res;
 }
 
-/* Called when a dummy is destroyed. */
+/* Called when a dummy is destroyed, indicating that the owning thread is being
+ * cleared. */
 static PyObject *
-_localdummy_destroyed(PyObject *localweakref, PyObject *dummyweakref)
+clear_locals(PyObject *locals_and_key, PyObject *dummyweakref)
 {
+    PyObject *localweakref = PyTuple_GetItem(locals_and_key, 0);
     localobject *self = (localobject *)_PyWeakref_GET_REF(localweakref);
     if (self == NULL) {
         Py_RETURN_NONE;
@@ -1742,11 +1819,18 @@ _localdummy_destroyed(PyObject *localweakref, PyObject *dummyweakref)
 
     /* If the thread-local object is still alive and not being cleared,
        remove the corresponding local dict */
-    if (self->dummies != NULL) {
+    if (self->localdicts != NULL) {
-        if (PyDict_Pop(self->dummies, dummyweakref, NULL) < 0) {
+        PyObject *key = PyTuple_GetItem(locals_and_key, 1);
+        if (PyDict_Pop(self->localdicts, key, NULL) < 0) {
             PyErr_WriteUnraisable((PyObject*)self);
         }
     }
+    if (self->thread_watchdogs != NULL) {
+        if (PySet_Discard(self->thread_watchdogs, dummyweakref) < 0) {
+            PyErr_WriteUnraisable((PyObject *)self);
+        }
+    }
+
     Py_DECREF(self);
     Py_RETURN_NONE;
 }

Python/pystate.c

@@ -1702,6 +1702,9 @@ PyThreadState_Clear(PyThreadState *tstate)
 
     /* Don't clear tstate->pyframe: it is a borrowed reference */
 
+    Py_CLEAR(tstate->threading_local_key);
+    Py_CLEAR(tstate->threading_local_sentinel);
+
     Py_CLEAR(((_PyThreadStateImpl *)tstate)->asyncio_running_loop);
 
     Py_CLEAR(tstate->dict);