Issue #28176: Fix callbacks race in asyncio.SelectorLoop.sock_connect.

1 files changed, 10 insertions, 17 deletions

diff --git a/Lib/asyncio/selector_events.py b/Lib/asyncio/selector_events.py
index c18885e..2f02d76 100644
--- a/Lib/asyncio/selector_events.py
+++ b/Lib/asyncio/selector_events.py
@@ -400,6 +400,7 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
                 data = data[n:]
             self.add_writer(fd, self._sock_sendall, fut, True, sock, data)
 
+    @coroutine
     def sock_connect(self, sock, address):
         """Connect to a remote socket at address.
 
@@ -408,24 +409,16 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
         if self._debug and sock.gettimeout() != 0:
             raise ValueError("the socket must be non-blocking")
 
-        fut = self.create_future()
-        if hasattr(socket, 'AF_UNIX') and sock.family == socket.AF_UNIX:
-            self._sock_connect(fut, sock, address)
-        else:
+        if not hasattr(socket, 'AF_UNIX') or sock.family != socket.AF_UNIX:
             resolved = base_events._ensure_resolved(
                 address, family=sock.family, proto=sock.proto, loop=self)
-            resolved.add_done_callback(
-                lambda resolved: self._on_resolved(fut, sock, resolved))
-
-        return fut
-
-    def _on_resolved(self, fut, sock, resolved):
-        try:
+            if not resolved.done():
+                yield from resolved
             _, _, _, _, address = resolved.result()[0]
-        except Exception as exc:
-            fut.set_exception(exc)
-        else:
-            self._sock_connect(fut, sock, address)
+
+        fut = self.create_future()
+        self._sock_connect(fut, sock, address)
+        return (yield from fut)
 
     def _sock_connect(self, fut, sock, address):
         fd = sock.fileno()
@@ -436,8 +429,8 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
             # connection runs in background. We have to wait until the socket
             # becomes writable to be notified when the connection succeed or
             # fails.
-            fut.add_done_callback(functools.partial(self._sock_connect_done,
-                                                    fd))
+            fut.add_done_callback(
+                functools.partial(self._sock_connect_done, fd))
             self.add_writer(fd, self._sock_connect_cb, fut, sock, address)
         except Exception as exc:
             fut.set_exception(exc)