Skip to content

Commit 66ade28

Browse files
[3.15] gh-146452: Improve locking granularity in pickle's batch_dict_exact and fix race condition (GH-150025) (#150039)
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) Co-authored-by: Saul Cooperman <58375603+scopreon@users.noreply.github.com>
1 parent c417fca commit 66ade28

2 files changed

Lines changed: 30 additions & 7 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: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3450,6 +3450,9 @@ batch_dict(PickleState *state, PicklerObject *self, PyObject *iter, PyObject *or
34503450
* Returns 0 on success, -1 on error.
34513451
*
34523452
* Note that this currently doesn't work for protocol 0.
3453+
3454+
* gh-146452: Wrap the dict iteration in a critical sections to prevent
3455+
* concurrent mutation from invalidating PyDict_Next() iteration state.
34533456
*/
34543457
static int
34553458
batch_dict_exact(PickleState *state, PicklerObject *self, PyObject *obj)
@@ -3466,15 +3469,24 @@ batch_dict_exact(PickleState *state, PicklerObject *self, PyObject *obj)
34663469
assert(self->proto > 0);
34673470

34683471
dict_size = PyDict_GET_SIZE(obj);
3469-
assert(dict_size);
34703472

34713473
/* Write in batches of BATCHSIZE. */
34723474
Py_ssize_t total = 0;
34733475
do {
34743476
if (dict_size - total == 1) {
3475-
PyDict_Next(obj, &ppos, &key, &value);
3476-
Py_INCREF(key);
3477-
Py_INCREF(value);
3477+
int next;
3478+
Py_BEGIN_CRITICAL_SECTION(obj);
3479+
next = PyDict_Next(obj, &ppos, &key, &value);
3480+
if (next) {
3481+
Py_INCREF(key);
3482+
Py_INCREF(value);
3483+
}
3484+
Py_END_CRITICAL_SECTION();
3485+
if (!next) {
3486+
PyErr_SetString(PyExc_RuntimeError,
3487+
"dictionary changed size during iteration");
3488+
goto error;
3489+
}
34783490
if (save(state, self, key, 0) < 0) {
34793491
goto error;
34803492
}
@@ -3492,9 +3504,18 @@ batch_dict_exact(PickleState *state, PicklerObject *self, PyObject *obj)
34923504
i = 0;
34933505
if (_Pickler_Write(self, &mark_op, 1) < 0)
34943506
return -1;
3495-
while (PyDict_Next(obj, &ppos, &key, &value)) {
3496-
Py_INCREF(key);
3497-
Py_INCREF(value);
3507+
int next;
3508+
while (1) {
3509+
Py_BEGIN_CRITICAL_SECTION(obj);
3510+
next = PyDict_Next(obj, &ppos, &key, &value);
3511+
if (next) {
3512+
Py_INCREF(key);
3513+
Py_INCREF(value);
3514+
}
3515+
Py_END_CRITICAL_SECTION();
3516+
if (!next) {
3517+
break;
3518+
}
34983519
if (save(state, self, key, 0) < 0) {
34993520
goto error;
35003521
}

0 commit comments

Comments
 (0)