gh-158600: Preserve concurrent removals for several operands in set.intersection_update()

Author: CaQtiml

Lib/test/test_free_threading/test_set.py

@@ -210,6 +210,49 @@ def intersect(first, second):
 
             self.assertEqual(target, expected)
 
+    def test_intersection_update_suspended_lock(self):
+        """Test an update while a later operand's lock is held."""
+        NUM_ITERS = 200
+        BLOCK_SIZE = self.SET_SIZE * 1_000
+        HOLD_WORK = 10_000
+
+        initial = set(range(BLOCK_SIZE))
+        blocked_operand = set(initial)
+        evens = set(range(0, BLOCK_SIZE, 2))
+        threes = set(range(0, BLOCK_SIZE, 3))
+        expected = set(range(0, BLOCK_SIZE, 6))
+
+        class SlowEmpty:
+            def __iter__(self):
+                for _ in range(HOLD_WORK):
+                    pass
+                return iter(())
+
+        # Repeat because the lock suspension depends on thread scheduling.
+        for _ in range(NUM_ITERS):
+            target = set(initial)
+
+            def hold_operand():
+                blocked_operand.intersection(SlowEmpty())
+
+            def first_update():
+                target.intersection_update(evens, blocked_operand)
+
+            def second_update():
+                target.intersection_update(threes)
+
+            threads = [
+                Thread(target=hold_operand),
+                Thread(target=first_update),
+                Thread(target=second_update),
+            ]
+            for thread in threads:
+                thread.start()
+            for thread in threads:
+                thread.join()
+
+            self.assertEqual(target, expected)
+
     def test_iand_concurrent(self):
         """Test concurrent &= operations on one shared set."""
         NUM_ITERS = 10

Objects/setobject.c

@@ -1887,7 +1887,19 @@ set_intersection_update_multi_impl(PySetObject *so, PyObject * const *others,
             Py_SETREF(result, newresult);
         }
         if (result != NULL) {
-            set_swap_bodies(so, (PySetObject *)result);
+            if (others_length == 1) {
+                set_swap_bodies(so, (PySetObject *)result);
+            }
+            else {
+                /* A later operand's critical section may have suspended
+                   so's lock. Preserve concurrent removals by intersecting
+                   the result with so's current contents. */
+                PyObject *updated = set_intersection_update(so, result);
+                if (updated == NULL) {
+                    Py_CLEAR(result);
+                }
+                Py_XDECREF(updated);
+            }
         }
     }
     Py_END_CRITICAL_SECTION2();