python/aqmp: refactor _do_accept() into two distinct steps

Refactor _do_accept() into _do_start_server() and _do_accept(). As of
this commit, the former calls the latter, but in subsequent commits
they'll be split apart.

(So please forgive the misnomer for _do_start_server(); it will live up
to its name shortly, and the docstring will be updated then too. I'm
just cutting down on some churn.)

Signed-off-by: John Snow <jsnow@redhat.com>
Acked-by: Kevin Wolf <kwolf@redhat.com>
Reviewed-by: Daniel P. Berrangé <berrange@redhat.com>
Message-id: 20220225205948.3693480-7-jsnow@redhat.com
Signed-off-by: John Snow <jsnow@redhat.com>

2 files changed, 26 insertions, 7 deletions

diff --git a/python/qemu/aqmp/protocol.py b/python/qemu/aqmp/protocol.py
index 631bcda..e2bdad5 100644
--- a/python/qemu/aqmp/protocol.py
+++ b/python/qemu/aqmp/protocol.py
@@ -295,7 +295,7 @@ class AsyncProtocol(Generic[T]):
             session, the wrapped error may also be an `QMPError`.
         """
         await self._session_guard(
-            self._do_accept(address, ssl),
+            self._do_start_server(address, ssl),
             'Failed to establish connection')
         await self._session_guard(
             self._establish_session(),
@@ -509,8 +509,8 @@ class AsyncProtocol(Generic[T]):
         self._sock = sock
 
     @upper_half
-    async def _do_accept(self, address: SocketAddrT,
-                         ssl: Optional[SSLContext] = None) -> None:
+    async def _do_start_server(self, address: SocketAddrT,
+                               ssl: Optional[SSLContext] = None) -> None:
         """
         Acting as the transport server, accept a single connection.
 
@@ -551,9 +551,28 @@ class AsyncProtocol(Generic[T]):
         # otherwise yield.
         await asyncio.sleep(0)
 
-        self._server = await coro    # Starts listening
-        await self._accepted.wait()  # Waits for the callback to finish
+        # This will start the server (bind(2), listen(2)). It will also
+        # call accept(2) if we yield, but we don't block on that here.
+        self._server = await coro
+
+        # Just for this one commit, wait for a peer.
+        # This gets split out in the next patch.
+        await self._do_accept()
+
+    @upper_half
+    async def _do_accept(self) -> None:
+        """
+        Wait for and accept an incoming connection.
+
+        Requires that we have not yet accepted an incoming connection
+        from the upper_half, but it's OK if the server is no longer
+        running because the bottom_half has already accepted the
+        connection.
+        """
+        assert self._accepted is not None
+        await self._accepted.wait()
         assert self._server is None
+        self._accepted = None
         self._sock = None
 
         self.logger.debug("Connection accepted.")
diff --git a/python/tests/protocol.py b/python/tests/protocol.py
index 8dd26c4..5e442e1 100644
--- a/python/tests/protocol.py
+++ b/python/tests/protocol.py
@@ -41,12 +41,12 @@ class NullProtocol(AsyncProtocol[None]):
         self.trigger_input = asyncio.Event()
         await super()._establish_session()
 
-    async def _do_accept(self, address, ssl=None):
+    async def _do_start_server(self, address, ssl=None):
         if self.fake_session:
             self._set_state(Runstate.CONNECTING)
             await asyncio.sleep(0)
         else:
-            await super()._do_accept(address, ssl)
+            await super()._do_start_server(address, ssl)
 
     async def _do_connect(self, address, ssl=None):
         if self.fake_session: