Skip to content
14 changes: 13 additions & 1 deletion Lib/test/test_free_threading/test_collections.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -49,6 +49,18 @@ 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,)
)


class TestOrderedDict(unittest.TestCase):
def test_iterator_update_clear_race(self):
# gh-151627: OrderedDict iterator construction must not race with
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fix a race condition under free-threading when multiple threads update the same :class:`~collections.Counter` concurrently.
54 changes: 33 additions & 21 deletions Modules/_collectionsmodule.c
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
#include "Python.h"
#include "pycore_call.h" // _PyObject_CallNoArgs()
#include "pycore_dict.h" // _PyDict_GetItem_KnownHash()
#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"
Expand Down Expand Up @@ -2547,7 +2547,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;
Expand Down Expand Up @@ -2594,24 +2593,36 @@ _collections__count_elements_impl(PyObject *module, PyObject *mapping,
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);
int found;
int cs_err = 0;
PyObject *newval = NULL;
Py_BEGIN_CRITICAL_SECTION(mapping);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The current PR looks good, but I think we can improve a bit more. The locking via the critical section now happens on each iteration, which seems a bit of a waste.

Can we place the critical section outside the while loop and then extract most (all?) of the code into a method _count_elements_dict_lock_held?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll look into this as soon as possible, thanks for reporting.

found = _PyDict_GetItemRef_KnownHash_LockHeld(
(PyDictObject *)mapping, key, hash, &oldval);
if (found < 0) {
cs_err = -1;
}
else if (found == 0) {
if (_PyDict_SetItem_KnownHash_LockHeld(
(PyDictObject *)mapping, key, one, hash) < 0) {
cs_err = -1;
}
}
else {
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);
if (newval == NULL) {
cs_err = -1;
}
else if (_PyDict_SetItem_KnownHash_LockHeld(
(PyDictObject *)mapping, key, newval, hash) < 0) {
cs_err = -1;
}
}
Py_END_CRITICAL_SECTION();
Py_CLEAR(newval);
if (cs_err < 0) {
goto done;
}
Py_DECREF(key);
}
Expand All @@ -2629,6 +2640,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 {
Expand All @@ -2637,9 +2649,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);
}
}
Expand All @@ -2649,7 +2662,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;
Expand Down
Loading