From ac03980d2ab855bd1bb7201c996764d84a5102bb Mon Sep 17 00:00:00 2001
From: Saul Cooperman <saulcoops@gmail.com>
Date: Mon, 18 May 2026 12:40:58 -0700
Subject: [PATCH] Improve locking granularity in pickle's batch_dict_exact
 (free-threading)

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.
---
 ...-05-18-15-30-34.gh-issue-146452.RM0EVJ.rst |  2 +
 Modules/_pickle.c                             | 49 +++++++++++--------
 2 files changed, 31 insertions(+), 20 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2026-05-18-15-30-34.gh-issue-146452.RM0EVJ.rst

diff --git a/Misc/NEWS.d/next/Library/2026-05-18-15-30-34.gh-issue-146452.RM0EVJ.rst b/Misc/NEWS.d/next/Library/2026-05-18-15-30-34.gh-issue-146452.RM0EVJ.rst
new file mode 100644
index 00000000000000..66f9acf6c710a7
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2026-05-18-15-30-34.gh-issue-146452.RM0EVJ.rst
@@ -0,0 +1,2 @@
+Fix race condition when pickling dictionaries in free threaded builds. Also
+reduce critical section cover.
diff --git a/Modules/_pickle.c b/Modules/_pickle.c
index 7b87be23269d40..15d95c658d6f90 100644
--- a/Modules/_pickle.c
+++ b/Modules/_pickle.c
@@ -3450,9 +3450,12 @@ batch_dict(PickleState *state, PicklerObject *self, PyObject *iter, PyObject *or
  * Returns 0 on success, -1 on error.
  *
  * Note that this currently doesn't work for protocol 0.
+
+ * gh-146452: Wrap the dict iteration in a critical sections to prevent
+ * concurrent mutation from invalidating PyDict_Next() iteration state.
  */
 static int
-batch_dict_exact_impl(PickleState *state, PicklerObject *self, PyObject *obj)
+batch_dict_exact(PickleState *state, PicklerObject *self, PyObject *obj)
 {
     PyObject *key = NULL, *value = NULL;
     int i;
@@ -3466,15 +3469,24 @@ batch_dict_exact_impl(PickleState *state, PicklerObject *self, PyObject *obj)
     assert(self->proto > 0);
 
     dict_size = PyDict_GET_SIZE(obj);
-    assert(dict_size);
 
     /* Write in batches of BATCHSIZE. */
     Py_ssize_t total = 0;
     do {
         if (dict_size - total == 1) {
-            PyDict_Next(obj, &ppos, &key, &value);
-            Py_INCREF(key);
-            Py_INCREF(value);
+            int next;
+            Py_BEGIN_CRITICAL_SECTION(obj);
+            next = PyDict_Next(obj, &ppos, &key, &value);
+            if (next) {
+                Py_INCREF(key);
+                Py_INCREF(value);
+            }
+            Py_END_CRITICAL_SECTION();
+            if (!next) {
+                PyErr_SetString(PyExc_RuntimeError,
+                                "dictionary changed size during iteration");
+                goto error;
+            }
             if (save(state, self, key, 0) < 0) {
                 goto error;
             }
@@ -3492,9 +3504,18 @@ batch_dict_exact_impl(PickleState *state, PicklerObject *self, PyObject *obj)
         i = 0;
         if (_Pickler_Write(self, &mark_op, 1) < 0)
             return -1;
-        while (PyDict_Next(obj, &ppos, &key, &value)) {
-            Py_INCREF(key);
-            Py_INCREF(value);
+        int next;
+        while (1) {
+            Py_BEGIN_CRITICAL_SECTION(obj);
+            next = PyDict_Next(obj, &ppos, &key, &value);
+            if (next) {
+                Py_INCREF(key);
+                Py_INCREF(value);
+            }
+            Py_END_CRITICAL_SECTION();
+            if (!next) {
+                break;
+            }
             if (save(state, self, key, 0) < 0) {
                 goto error;
             }
@@ -3525,18 +3546,6 @@ batch_dict_exact_impl(PickleState *state, PicklerObject *self, PyObject *obj)
     return -1;
 }
 
-/* gh-146452: Wrap the dict iteration in a critical section to prevent
-   concurrent mutation from invalidating PyDict_Next() iteration state. */
-static int
-batch_dict_exact(PickleState *state, PicklerObject *self, PyObject *obj)
-{
-    int ret;
-    Py_BEGIN_CRITICAL_SECTION(obj);
-    ret = batch_dict_exact_impl(state, self, obj);
-    Py_END_CRITICAL_SECTION();
-    return ret;
-}
-
 static int
 save_dict(PickleState *state, PicklerObject *self, PyObject *obj)
 {