Skip to content

Commit 1243cd6

Browse files
authored
[3.14] gh-146452: Improve locking granularity in pickle's batch_dict_… (#150062)
[3.14] gh-146452: Improve locking granularity in pickle's batch_dict_exact and fix race condition (GH-150025) Remove assertion that could fail in rare race condition. Replace the coarse critical section wrapping the entire function with fine-grained sections covering only PyDict_Next + Py_INCREF. Also handle PyDict_Next returning 0 in the single-item fast path. (cherry picked from commit 57a0e57)
1 parent 8e13025 commit 1243cd6

2 files changed

Lines changed: 31 additions & 19 deletions

File tree

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix race condition when pickling dictionaries in free threaded builds. Also
2+
reduce critical section cover.

Modules/_pickle.c

Lines changed: 29 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -3351,9 +3351,12 @@ batch_dict(PickleState *state, PicklerObject *self, PyObject *iter, PyObject *or
33513351
* Returns 0 on success, -1 on error.
33523352
*
33533353
* Note that this currently doesn't work for protocol 0.
3354+
3355+
* gh-146452: Wrap the dict iteration in a critical sections to prevent
3356+
* concurrent mutation from invalidating PyDict_Next() iteration state.
33543357
*/
33553358
static int
3356-
batch_dict_exact_impl(PickleState *state, PicklerObject *self, PyObject *obj)
3359+
batch_dict_exact(PickleState *state, PicklerObject *self, PyObject *obj)
33573360
{
33583361
PyObject *key = NULL, *value = NULL;
33593362
int i;
@@ -3370,9 +3373,19 @@ batch_dict_exact_impl(PickleState *state, PicklerObject *self, PyObject *obj)
33703373

33713374
/* Special-case len(d) == 1 to save space. */
33723375
if (dict_size == 1) {
3373-
PyDict_Next(obj, &ppos, &key, &value);
3374-
Py_INCREF(key);
3375-
Py_INCREF(value);
3376+
int next;
3377+
Py_BEGIN_CRITICAL_SECTION(obj);
3378+
next = PyDict_Next(obj, &ppos, &key, &value);
3379+
if (next) {
3380+
Py_INCREF(key);
3381+
Py_INCREF(value);
3382+
}
3383+
Py_END_CRITICAL_SECTION();
3384+
if (!next) {
3385+
PyErr_SetString(PyExc_RuntimeError,
3386+
"dictionary changed size during iteration");
3387+
goto error;
3388+
}
33763389
if (save(state, self, key, 0) < 0) {
33773390
goto error;
33783391
}
@@ -3392,9 +3405,18 @@ batch_dict_exact_impl(PickleState *state, PicklerObject *self, PyObject *obj)
33923405
i = 0;
33933406
if (_Pickler_Write(self, &mark_op, 1) < 0)
33943407
return -1;
3395-
while (PyDict_Next(obj, &ppos, &key, &value)) {
3396-
Py_INCREF(key);
3397-
Py_INCREF(value);
3408+
int next;
3409+
while (1) {
3410+
Py_BEGIN_CRITICAL_SECTION(obj);
3411+
next = PyDict_Next(obj, &ppos, &key, &value);
3412+
if (next) {
3413+
Py_INCREF(key);
3414+
Py_INCREF(value);
3415+
}
3416+
Py_END_CRITICAL_SECTION();
3417+
if (!next) {
3418+
break;
3419+
}
33983420
if (save(state, self, key, 0) < 0) {
33993421
goto error;
34003422
}
@@ -3424,18 +3446,6 @@ batch_dict_exact_impl(PickleState *state, PicklerObject *self, PyObject *obj)
34243446
return -1;
34253447
}
34263448

3427-
/* gh-146452: Wrap the dict iteration in a critical section to prevent
3428-
concurrent mutation from invalidating PyDict_Next() iteration state. */
3429-
static int
3430-
batch_dict_exact(PickleState *state, PicklerObject *self, PyObject *obj)
3431-
{
3432-
int ret;
3433-
Py_BEGIN_CRITICAL_SECTION(obj);
3434-
ret = batch_dict_exact_impl(state, self, obj);
3435-
Py_END_CRITICAL_SECTION();
3436-
return ret;
3437-
}
3438-
34393449
static int
34403450
save_dict(PickleState *state, PicklerObject *self, PyObject *obj)
34413451
{

0 commit comments

Comments
 (0)