gh-154971: Decompose IocpProactor._poll()

Lib/asyncio/windows_events.py

@@ -760,6 +760,46 @@ def _get_accept_socket(self, family):
         s.settimeout(0)
         return s
 
+    def _process_completion_status(self, status):
+        """Process a single status from the completion port.
+
+        A caller that waits on the completion port itself can pass each
+        status it receives here.
+        """
+        err, transferred, key, address = status
+        try:
+            f, ov, obj, callback = self._cache.pop(address)
+        except KeyError:
+            if self._loop.get_debug():
+                self._loop.call_exception_handler({
+                    'message': ('GetQueuedCompletionStatus() returned an '
+                                'unexpected event'),
+                    'status': ('err=%s transferred=%s key=%#x address=%#x'
+                               % (err, transferred, key, address)),
+                })
+
+            # key is either zero, or it is used to return a pipe
+            # handle which should be closed to avoid a leak.
+            if key not in (0, _overlapped.INVALID_HANDLE_VALUE):
+                _winapi.CloseHandle(key)
+            return
+
+        if obj in self._stopped_serving:
+            f.cancel()
+        # Don't call the callback if _register() already read the result or
+        # if the overlapped has been cancelled
+        elif not f.done():
+            try:
+                value = callback(transferred, key, ov)
+            except OSError as e:
+                f.set_exception(e)
+                self._results.append(f)
+            else:
+                f.set_result(value)
+                self._results.append(f)
+            finally:
+                f = None
+
     def _poll(self, timeout=None):
         if timeout is None:
             ms = INFINITE
@@ -778,39 +818,8 @@ def _poll(self, timeout=None):
                 break
             ms = 0
 
-            err, transferred, key, address = status
-            try:
-                f, ov, obj, callback = self._cache.pop(address)
-            except KeyError:
-                if self._loop.get_debug():
-                    self._loop.call_exception_handler({
-                        'message': ('GetQueuedCompletionStatus() returned an '
-                                    'unexpected event'),
-                        'status': ('err=%s transferred=%s key=%#x address=%#x'
-                                   % (err, transferred, key, address)),
-                    })
-
-                # key is either zero, or it is used to return a pipe
-                # handle which should be closed to avoid a leak.
-                if key not in (0, _overlapped.INVALID_HANDLE_VALUE):
-                    _winapi.CloseHandle(key)
-                continue
-
-            if obj in self._stopped_serving:
-                f.cancel()
-            # Don't call the callback if _register() already read the result or
-            # if the overlapped has been cancelled
-            elif not f.done():
-                try:
-                    value = callback(transferred, key, ov)
-                except OSError as e:
-                    f.set_exception(e)
-                    self._results.append(f)
-                else:
-                    f.set_result(value)
-                    self._results.append(f)
-                finally:
-                    f = None
+            # gh-154971: split out so custom event loops can call it directly
+            self._process_completion_status(status)
 
         # Remove unregistered futures
         for ov in self._unregistered:

Lib/test/test_asyncio/test_windows_events.py

@@ -327,6 +327,28 @@ def threadMain():
         stop.set()
         thr.join()
 
+    def test_custom_poll_integration(self):
+        # gh-154971: a caller can wait on the completion port and process statuses itself
+        proactor = self.loop._proactor
+
+        a, b = socket.socketpair()
+        self.addCleanup(a.close)
+        self.addCleanup(b.close)
+
+        fut = proactor.recv(a, 100)
+        self.assertFalse(fut.done())
+
+        b.send(b'data')
+
+        deadline = time.monotonic() + support.SHORT_TIMEOUT
+        while not fut.done() and time.monotonic() < deadline:
+            status = _overlapped.GetQueuedCompletionStatus(proactor._iocp, 100)
+            if status is not None:
+                proactor._process_completion_status(status)
+
+        self.assertTrue(fut.done())
+        self.assertEqual(fut.result(), b'data')
+
 
 class ProactorPipeObjectSupportTests(unittest.TestCase):