gh-89240: Enable multiprocessing on Windows to use large process pools (GH-107873)

We add _winapi.BatchedWaitForMultipleObjects to wait for larger numbers of handles.
This is an internal module, hence undocumented, and should be used with caution.
Check the docstring for info before using BatchedWaitForMultipleObjects.

3 files changed, 125 insertions, 1 deletions

diff --git a/Lib/multiprocessing/connection.py b/Lib/multiprocessing/connection.py
index c6a66a1..58d697f 100644
--- a/Lib/multiprocessing/connection.py
+++ b/Lib/multiprocessing/connection.py
@@ -1011,8 +1011,20 @@ if sys.platform == 'win32':
         # returning the first signalled might create starvation issues.)
         L = list(handles)
         ready = []
+        # Windows limits WaitForMultipleObjects at 64 handles, and we use a
+        # few for synchronisation, so we switch to batched waits at 60.
+        if len(L) > 60:
+            try:
+                res = _winapi.BatchedWaitForMultipleObjects(L, False, timeout)
+            except TimeoutError:
+                return []
+            ready.extend(L[i] for i in res)
+            if res:
+                L = [h for i, h in enumerate(L) if i > res[0] & i not in res]
+            timeout = 0
         while L:
-            res = _winapi.WaitForMultipleObjects(L, False, timeout)
+            short_L = L[:60] if len(L) > 60 else L
+            res = _winapi.WaitForMultipleObjects(short_L, False, timeout)
             if res == WAIT_TIMEOUT:
                 break
             elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
diff --git a/Lib/test/_test_multiprocessing.py b/Lib/test/_test_multiprocessing.py
index c0d3ca5..94ce85c 100644
--- a/Lib/test/_test_multiprocessing.py
+++ b/Lib/test/_test_multiprocessing.py
@@ -6113,6 +6113,24 @@ class MiscTestCase(unittest.TestCase):
         self.assertEqual(rc, 0)
         self.assertFalse(err, msg=err.decode('utf-8'))
 
+    def test_large_pool(self):
+        #
+        # gh-89240: Check that large pools are always okay
+        #
+        testfn = os_helper.TESTFN
+        self.addCleanup(os_helper.unlink, testfn)
+        with open(testfn, 'w', encoding='utf-8') as f:
+            f.write(textwrap.dedent('''\
+                import multiprocessing
+                def f(x): return x*x
+                if __name__ == '__main__':
+                    with multiprocessing.Pool(200) as p:
+                        print(sum(p.map(f, range(1000))))
+            '''))
+        rc, out, err = script_helper.assert_python_ok(testfn)
+        self.assertEqual("332833500", out.decode('utf-8').strip())
+        self.assertFalse(err, msg=err.decode('utf-8'))
+
 
 #
 # Mixins
diff --git a/Lib/test/test_winapi.py b/Lib/test/test_winapi.py
new file mode 100644
index 0000000..014aeea
--- /dev/null
+++ b/Lib/test/test_winapi.py
@@ -0,0 +1,94 @@
+# Test the Windows-only _winapi module
+
+import random
+import threading
+import time
+import unittest
+from test.support import import_helper
+
+_winapi = import_helper.import_module('_winapi', required_on=['win'])
+
+MAXIMUM_WAIT_OBJECTS = 64
+MAXIMUM_BATCHED_WAIT_OBJECTS = (MAXIMUM_WAIT_OBJECTS - 1) ** 2
+
+class WinAPIBatchedWaitForMultipleObjectsTests(unittest.TestCase):
+    def _events_waitall_test(self, n):
+        evts = [_winapi.CreateEventW(0, False, False, None) for _ in range(n)]
+
+        with self.assertRaises(TimeoutError):
+            _winapi.BatchedWaitForMultipleObjects(evts, True, 100)
+
+        # Ensure no errors raised when all are triggered
+        for e in evts:
+            _winapi.SetEvent(e)
+        try:
+            _winapi.BatchedWaitForMultipleObjects(evts, True, 100)
+        except TimeoutError:
+            self.fail("expected wait to complete immediately")
+
+        # Choose 8 events to set, distributed throughout the list, to make sure
+        # we don't always have them in the first chunk
+        chosen = [i * (len(evts) // 8) for i in range(8)]
+
+        # Replace events with invalid handles to make sure we fail
+        for i in chosen:
+            old_evt = evts[i]
+            evts[i] = -1
+            with self.assertRaises(OSError):
+                _winapi.BatchedWaitForMultipleObjects(evts, True, 100)
+            evts[i] = old_evt
+
+
+    def _events_waitany_test(self, n):
+        evts = [_winapi.CreateEventW(0, False, False, None) for _ in range(n)]
+
+        with self.assertRaises(TimeoutError):
+            _winapi.BatchedWaitForMultipleObjects(evts, False, 100)
+
+        # Choose 8 events to set, distributed throughout the list, to make sure
+        # we don't always have them in the first chunk
+        chosen = [i * (len(evts) // 8) for i in range(8)]
+
+        # Trigger one by one. They are auto-reset events, so will only trigger once
+        for i in chosen:
+            with self.subTest(f"trigger event {i} of {len(evts)}"):
+                _winapi.SetEvent(evts[i])
+                triggered = _winapi.BatchedWaitForMultipleObjects(evts, False, 10000)
+                self.assertSetEqual(set(triggered), {i})
+
+        # Trigger all at once. This may require multiple calls
+        for i in chosen:
+            _winapi.SetEvent(evts[i])
+        triggered = set()
+        while len(triggered) < len(chosen):
+            triggered.update(_winapi.BatchedWaitForMultipleObjects(evts, False, 10000))
+        self.assertSetEqual(triggered, set(chosen))
+
+        # Replace events with invalid handles to make sure we fail
+        for i in chosen:
+            with self.subTest(f"corrupt event {i} of {len(evts)}"):
+                old_evt = evts[i]
+                evts[i] = -1
+                with self.assertRaises(OSError):
+                    _winapi.BatchedWaitForMultipleObjects(evts, False, 100)
+                evts[i] = old_evt
+
+
+    def test_few_events_waitall(self):
+        self._events_waitall_test(16)
+
+    def test_many_events_waitall(self):
+        self._events_waitall_test(256)
+
+    def test_max_events_waitall(self):
+        self._events_waitall_test(MAXIMUM_BATCHED_WAIT_OBJECTS)
+
+
+    def test_few_events_waitany(self):
+        self._events_waitany_test(16)
+
+    def test_many_events_waitany(self):
+        self._events_waitany_test(256)
+
+    def test_max_events_waitany(self):
+        self._events_waitany_test(MAXIMUM_BATCHED_WAIT_OBJECTS)