Try to fix asyncio.Server.wait_closed() again

I identified the condition that `wait_closed()` is intended to wait for: the server is closed *and* there are no more active connections. When this condition first becomes true, `_wakeup()` is called (either from `close()` or from `_detach()`) and it sets `_waiters` to `None`. So we just check for `self._waiters is None`; if it's not `None`, we know we have to wait, and do so. A problem was that the new test introduced in 3.12 explicitly tested that `wait_closed()` returns immediately when the server is *not* closed but there are currently no active connections. This was a mistake (probably a misunderstanding of the intended semantics). I've fixed the test, and added a separate test that checks exactly for this scenario. I also fixed an oddity where in `_wakeup()` the result of the waiter was set to the waiter itself. This result is not used anywhere and I changed this to `None`, to avoid a GC cycle.
python · willingc · Oct 28, 2023 · Oct 25, 2023 · Oct 28, 2023 · Oct 25, 2023
commit 503f0bdb53fb305a49864c5125575657e43e7b28
@@ -1619,8 +1619,9 @@ Do not instantiate the :class:`Server` class directly.
      The sockets that represent existing incoming client connections
      are left open.

-      The server is closed asynchronously, use the :meth:`wait_closed`
-      coroutine to wait until the server is closed.
+      The server is closed asynchronously; use the :meth:`wait_closed`
+      coroutine to wait until the server is closed (and no more
+      connections are active).

   .. method:: get_loop()

@@ -1678,7 +1679,8 @@ Do not instantiate the :class:`Server` class directly.

   .. coroutinemethod:: wait_closed()

-      Wait until the :meth:`close` method completes.
+      Wait until the :meth:`close` method completes and all active
+      connections have finished.

   .. attribute:: sockets


@@ -305,7 +305,7 @@ def _wakeup(self):
        self._waiters = None
        for waiter in waiters:
            if not waiter.done():
-                waiter.set_result(waiter)
+                waiter.set_result(None)

    def _start_serving(self):
        if self._serving:
@@ -377,7 +377,26 @@ async def serve_forever(self):
            self._serving_forever_fut = None

    async def wait_closed(self):
-        if self._waiters is None or self._active_count == 0:
+        """Wait until server is closed and all connections are dropped.
+
+        - If the server is not closed, wait.
+        - If it is closed, but there are still active connections, wait.
+
+        Anyone waiting here will be unblocked once both conditions
+        have become true, in either order.
+
+        Historical note: In 3.11 and before, this was broken, returning
+        immediately if the server was already closed, even if there
+        were still active connections. An attempted fix in 3.12.0 was
+        still broken, returning immediately if the server was still
+        open and there were no active connections. Hopefully in 3.12.1
+        we have it right.
+        """
+        # Waiters are unblocked by self._wakeup(), which is called
+        # from two places: self.close() and self._detach(), but only
+        # when both conditions have become true. To signal that this
+        # has happened, self._wakeup() sets self._waiters to None.
+        if self._waiters is None:
            return
        waiter = self._loop.create_future()
        self._waiters.append(waiter)

@@ -122,29 +122,59 @@ async def main(srv):

 class TestServer2(unittest.IsolatedAsyncioTestCase):

-    async def test_wait_closed(self):
+    async def test_wait_closed_basic(self):
        async def serve(*args):
            pass

        srv = await asyncio.start_server(serve, socket_helper.HOSTv4, 0)
+        self.addCleanup(srv.close)

-        # active count = 0
+        # active count = 0, not closed: should block
        task1 = asyncio.create_task(srv.wait_closed())
        await asyncio.sleep(0)
-        self.assertTrue(task1.done())
+        self.assertFalse(task1.done())

-        # active count != 0
+        # active count != 0, not closed: should block
        srv._attach()
        task2 = asyncio.create_task(srv.wait_closed())
        await asyncio.sleep(0)
+        self.assertFalse(task1.done())
        self.assertFalse(task2.done())

        srv.close()
        await asyncio.sleep(0)
+        # active count != 0, closed: should block
+        task3 = asyncio.create_task(srv.wait_closed())
+        await asyncio.sleep(0)
+        self.assertFalse(task1.done())
        self.assertFalse(task2.done())
+        self.assertFalse(task3.done())

        srv._detach()
+        # active count == 0, closed: should unblock
+        await task1
        await task2
+        await task3
+        await srv.wait_closed()  # Return immediately
+
+    async def test_wait_closed_race(self):
+        # Test a regression in 3.12.0, should be fixed in 3.12.1
+        async def serve(*args):
+            pass
+
+        srv = await asyncio.start_server(serve, socket_helper.HOSTv4, 0)
+        self.addCleanup(srv.close)
+
+        task = asyncio.create_task(srv.wait_closed())
+        await asyncio.sleep(0)
+        self.assertFalse(task.done())
+        srv._attach()
+        loop = asyncio.get_running_loop()
+        loop.call_soon(srv.close)
+        loop.call_soon(srv._detach)
+        await srv.wait_closed()
+
+


 @unittest.skipUnless(hasattr(asyncio, 'ProactorEventLoop'), 'Windows only')

diff --git a/Misc/NEWS.d/next/Library/2023-10-25-11-54-00.gh-issue-79033.5ePgFl.rst b/Misc/NEWS.d/next/Library/2023-10-25-11-54-00.gh-issue-79033.5ePgFl.rst
@@ -0,0 +1,6 @@
+Another attempt at fixing :func:`asyncio.Server.wait_closed()`. It now
+blocks until both conditions are true: the server is closed, *and* there
+are no more active connections. (This means that in some cases where in
+3.12.0 this function would *incorrectly* have returned immediately,
+it will now block; in particular, when there are no active connections
+but the server hasn't been closed yet.)