Skip to content

Commit 61839ca

Browse files
author
Ahmed Magdy Ali
committed
gh-150505: fix data race in Unpickler_set_memo under free-threading
Unpickler_set_memo modified self->memo and self->memo_size without any critical section, allowing concurrent threads to race on the same Unpickler instance. Wrap the memo population loop and the pointer swap in Py_BEGIN_CRITICAL_SECTION/Py_END_CRITICAL_SECTION to serialize access.
1 parent 9242700 commit 61839ca

3 files changed

Lines changed: 53 additions & 6 deletions

File tree

Lib/test/test_pickle.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from test import support
1818
from test.support import cpython_only, import_helper, os_helper
1919
from test.support.import_helper import ensure_lazy_imports
20+
from test.support import threading_helper
2021

2122
from test.pickletester import AbstractHookTests
2223
from test.pickletester import AbstractUnpickleTests
@@ -808,6 +809,20 @@ def test_unknown_flag(self):
808809
self.assertStartsWith(stderr.getvalue(), 'usage: ')
809810

810811

812+
class UnpicklerMemoThreadSafetyTest(unittest.TestCase):
813+
@unittest.skipUnless(support.Py_GIL_DISABLED, 'this test can only possibly fail with GIL disabled')
814+
@threading_helper.reap_threads
815+
@threading_helper.requires_working_threading()
816+
def test_memo_assignment_race(self):
817+
from threading import Thread
818+
shared = pickle.Unpickler(io.BytesIO(b''))
819+
def assign():
820+
for _ in range(10000):
821+
shared.memo = {j: j for j in range(100)}
822+
threads = [Thread(target=assign) for _ in range(8)]
823+
for t in threads: t.start()
824+
for t in threads: t.join()
825+
811826
def load_tests(loader, tests, pattern):
812827
tests.addTest(doctest.DocTestSuite(pickle))
813828
return tests
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Fix data race in :class:`pickle.Unpickler` when multiple threads assign to
2+
the ``memo`` attribute concurrently under free-threading. The memo swap and
3+
population loop are now protected with a critical section.

Modules/_pickle.c

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7786,24 +7786,38 @@ Unpickler_set_memo(PyObject *op, PyObject *obj, void *Py_UNUSED(closure))
77867786
if (new_memo == NULL)
77877787
return -1;
77887788

7789+
bool goto_error = false;
7790+
Py_BEGIN_CRITICAL_SECTION(self);
77897791
while (PyDict_Next(obj, &i, &key, &value)) {
77907792
Py_ssize_t idx;
77917793
if (!PyLong_Check(key)) {
77927794
PyErr_SetString(PyExc_TypeError,
77937795
"memo key must be integers");
7794-
goto error;
7796+
goto_error = true;
7797+
break;
77957798
}
77967799
idx = PyLong_AsSsize_t(key);
7797-
if (idx == -1 && PyErr_Occurred())
7798-
goto error;
7800+
if (idx == -1 && PyErr_Occurred()) {
7801+
goto_error = true;
7802+
break;
7803+
}
77997804
if (idx < 0) {
78007805
PyErr_SetString(PyExc_ValueError,
78017806
"memo key must be positive integers.");
7802-
goto error;
7807+
goto_error = true;
7808+
break;
78037809
}
7810+
78047811
if (_Unpickler_MemoPut(self, idx, value) < 0)
7805-
goto error;
7812+
{
7813+
goto_error = true;
7814+
break;
7815+
}
78067816
}
7817+
Py_END_CRITICAL_SECTION();
7818+
if (goto_error)
7819+
goto error;
7820+
78077821
}
78087822
else {
78097823
PyErr_Format(PyExc_TypeError,
@@ -7812,9 +7826,24 @@ Unpickler_set_memo(PyObject *op, PyObject *obj, void *Py_UNUSED(closure))
78127826
return -1;
78137827
}
78147828

7815-
_Unpickler_MemoCleanup(self);
7829+
Py_ssize_t old_memo_size, i;
7830+
PyObject **old_memo;
7831+
7832+
Py_BEGIN_CRITICAL_SECTION(self);
7833+
old_memo_size = self->memo_size;
7834+
old_memo = self->memo;
78167835
self->memo_size = new_memo_size;
78177836
self->memo = new_memo;
7837+
Py_END_CRITICAL_SECTION();
7838+
7839+
if (old_memo == NULL) {
7840+
return 0;
7841+
}
7842+
i = old_memo_size;
7843+
while (--i >= 0) {
7844+
Py_XDECREF(old_memo[i]);
7845+
}
7846+
PyMem_Free(old_memo);
78187847

78197848
return 0;
78207849

0 commit comments

Comments
 (0)