diff --git a/Lib/test/test_free_threading/test_collections.py b/Lib/test/test_free_threading/test_collections.py index a0a73c965d072a2..d6efde55e698d0a 100644 --- a/Lib/test/test_free_threading/test_collections.py +++ b/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 diff --git a/Misc/NEWS.d/next/Library/2026-06-18-11-26-00.gh-issue-151633.uJzHdc.rst b/Misc/NEWS.d/next/Library/2026-06-18-11-26-00.gh-issue-151633.uJzHdc.rst new file mode 100644 index 000000000000000..73d533c0d109bbe --- /dev/null +++ b/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. diff --git a/Modules/_collectionsmodule.c b/Modules/_collectionsmodule.c index 0626cc6a70bc8f7..c63dc0661bd69e7 100644 --- a/Modules/_collectionsmodule.c +++ b/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;