gh-151633: avoid use-after-free in Counter.update

Lib/test/test_free_threading/test_collections.py

@@ -1,5 +1,5 @@
 import unittest
-from collections import OrderedDict, deque
+from collections import Counter, OrderedDict, deque
 from copy import copy
 from test.support import threading_helper
 
@@ -49,6 +49,19 @@ def mutate():
         )
 
 
+class TestCounter(unittest.TestCase):
+    def test_update_concurrent(self):
+        # gh-151633: concurrent Counter.update calls must not cause use-after-free
+        NTHREADS = 4
+        PER_THREAD = 5000
+        c = Counter()
+        data = ('x',) * PER_THREAD
+        threading_helper.run_concurrently(
+            c.update, nthreads=NTHREADS, args=(data,)
+        )
+        self.assertEqual(c['x'], NTHREADS * PER_THREAD)
+
+
 class TestOrderedDict(unittest.TestCase):
     def test_iterator_update_clear_race(self):
         # gh-151627: OrderedDict iterator construction must not race with

Misc/NEWS.d/next/Library/2026-06-18-11-26-00.gh-issue-151633.uJzHdc.rst

@@ -0,0 +1,2 @@
+Fix a use-after-free under free-threading when multiple threads
+update the same :class:`~collections.Counter` concurrently.

Modules/_collectionsmodule.c

@@ -1,6 +1,7 @@
 #include "Python.h"
 #include "pycore_call.h"          // _PyObject_CallNoArgs()
-#include "pycore_dict.h"          // _PyDict_GetItem_KnownHash()
+#include "pycore_critical_section.h" // _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED()
+#include "pycore_dict.h"          // _PyDict_GetItemRef_KnownHash_LockHeld()
 #include "pycore_long.h"          // _PyLong_GetZero()
 #include "pycore_moduleobject.h"  // _PyModule_GetState()
 #include "pycore_pyatomic_ft_wrappers.h"
@@ -2525,6 +2526,66 @@ static PyType_Spec defdict_spec = {
 
 /* helper function for Counter  *********************************************/
 
+/* Fast path advantages:
+       1. Eliminate double hashing
+          (by re-using the same hash for both the get and set)
+       2. Avoid argument overhead of PyObject_CallFunctionObjArgs
+          (argument tuple creation and parsing)
+       3. Avoid indirection through a bound method object
+          (creates another argument tuple)
+       4. Avoid initial increment from zero
+          (reuse an existing one-object instead)
+*/
+static int
+_count_elements_dict_lock_held(PyDictObject *mapping, PyObject *it)
+{
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(mapping);
+    PyObject *one = _PyLong_GetOne();  // borrowed reference
+
+    while (1) {
+        Py_hash_t hash;
+        PyObject *oldval;
+        int status;
+
+        PyObject *key = PyIter_Next(it);
+        if (key == NULL) {
+            return PyErr_Occurred() ? -1 : 0;
+        }
+
+        hash = _PyObject_HashDictKey(key);
+        if (hash == -1) {
+            Py_DECREF(key);
+            return -1;
+        }
+
+        int found = _PyDict_GetItemRef_KnownHash_LockHeld(mapping, key, hash,
+                                                          &oldval);
+        if (found < 0) {
+            Py_DECREF(key);
+            return -1;
+        }
+        if (found == 0) {
+            status = _PyDict_SetItem_KnownHash_LockHeld(mapping, key, one,
+                                                        hash);
+        }
+        else {
+            PyObject *newval = PyNumber_Add(oldval, one);
+            Py_DECREF(oldval);
+            if (newval == NULL) {
+                Py_DECREF(key);
+                return -1;
+            }
+            status = _PyDict_SetItem_KnownHash_LockHeld(mapping, key, newval,
+                                                        hash);
+            Py_DECREF(newval);
+        }
+        Py_DECREF(key);
+        if (status < 0) {
+            return -1;
+        }
+    }
+}
+
 /*[clinic input]
 _collections._count_elements
 
@@ -2541,7 +2602,6 @@ _collections__count_elements_impl(PyObject *module, PyObject *mapping,
 /*[clinic end generated code: output=7e0c1789636b3d8f input=e79fad04534a0b45]*/
 {
     PyObject *it, *oldval;
-    PyObject *newval = NULL;
     PyObject *key = NULL;
     PyObject *bound_get = NULL;
     PyObject *mapping_get;
@@ -2566,48 +2626,12 @@ _collections__count_elements_impl(PyObject *module, PyObject *mapping,
         mapping_setitem != NULL && mapping_setitem == dict_setitem &&
         PyDict_Check(mapping))
     {
-        while (1) {
-            /* Fast path advantages:
-                   1. Eliminate double hashing
-                      (by re-using the same hash for both the get and set)
-                   2. Avoid argument overhead of PyObject_CallFunctionObjArgs
-                      (argument tuple creation and parsing)
-                   3. Avoid indirection through a bound method object
-                      (creates another argument tuple)
-                   4. Avoid initial increment from zero
-                      (reuse an existing one-object instead)
-            */
-            Py_hash_t hash;
-
-            key = PyIter_Next(it);
-            if (key == NULL)
-                break;
-
-            hash = _PyObject_HashDictKey(key);
-            if (hash == -1) {
-                goto done;
-            }
-
-            oldval = _PyDict_GetItem_KnownHash(mapping, key, hash);
-            if (oldval == NULL) {
-                if (PyErr_Occurred())
-                    goto done;
-                if (_PyDict_SetItem_KnownHash(mapping, key, one, hash) < 0)
-                    goto done;
-            } else {
-                /* oldval is a borrowed reference.  Keep it alive across
-                   PyNumber_Add(), which can execute arbitrary user code and
-                   mutate (or even clear) the underlying dict. */
-                Py_INCREF(oldval);
-                newval = PyNumber_Add(oldval, one);
-                Py_DECREF(oldval);
-                if (newval == NULL)
-                    goto done;
-                if (_PyDict_SetItem_KnownHash(mapping, key, newval, hash) < 0)
-                    goto done;
-                Py_CLEAR(newval);
-            }
-            Py_DECREF(key);
+        int status;
+        Py_BEGIN_CRITICAL_SECTION(mapping);
+        status = _count_elements_dict_lock_held((PyDictObject *)mapping, it);
+        Py_END_CRITICAL_SECTION();
+        if (status < 0) {
+            goto done;
         }
     }
     else {
@@ -2623,6 +2647,7 @@ _collections__count_elements_impl(PyObject *module, PyObject *mapping,
             oldval = PyObject_CallFunctionObjArgs(bound_get, key, zero, NULL);
             if (oldval == NULL)
                 break;
+            PyObject *newval;
             if (oldval == zero) {
                 newval = Py_NewRef(one);
             } else {
@@ -2631,9 +2656,10 @@ _collections__count_elements_impl(PyObject *module, PyObject *mapping,
             Py_DECREF(oldval);
             if (newval == NULL)
                 break;
-            if (PyObject_SetItem(mapping, key, newval) < 0)
+            int status = PyObject_SetItem(mapping, key, newval);
+            Py_DECREF(newval);
+            if (status < 0)
                 break;
-            Py_CLEAR(newval);
             Py_DECREF(key);
         }
     }
@@ -2643,7 +2669,6 @@ _collections__count_elements_impl(PyObject *module, PyObject *mapping,
     Py_XDECREF(mapping_setitem);
     Py_DECREF(it);
     Py_XDECREF(key);
-    Py_XDECREF(newval);
     Py_XDECREF(bound_get);
     if (PyErr_Occurred())
         return NULL;