Remove reader/writer when sock_recv(), sock_sendall(), etc are cancelled

Author: 1st1

tests/test_sockets.py

@@ -505,6 +505,65 @@ def test_socket_close_remove_writer(self):
             s.close()
             self.assertEqual(s.fileno(), -1)
 
+    def test_socket_cancel_sock_recv(self):
+        def srv_gen(sock):
+            time.sleep(1.2)
+            sock.send(b'helo')
+
+        async def kill(fut):
+            await asyncio.sleep(0.2, loop=self.loop)
+            fut.cancel()
+
+        async def client(sock, addr):
+            await self.loop.sock_connect(sock, addr)
+
+            f = asyncio.ensure_future(self.loop.sock_recv(sock, 10),
+                                      loop=self.loop)
+            self.loop.create_task(kill(f))
+            with self.assertRaises(asyncio.CancelledError):
+                await f
+            sock.close()
+            self.assertEqual(sock.fileno(), -1)
+
+        with self.tcp_server(srv_gen) as srv:
+
+            sock = socket.socket()
+            with sock:
+                sock.setblocking(False)
+                c = client(sock, srv.addr)
+                w = asyncio.wait_for(c, timeout=5.0, loop=self.loop)
+                self.loop.run_until_complete(w)
+
+    def test_socket_cancel_sock_sendall(self):
+        def srv_gen(sock):
+            time.sleep(1.2)
+            sock.recv_all(4)
+
+        async def kill(fut):
+            await asyncio.sleep(0.2, loop=self.loop)
+            fut.cancel()
+
+        async def client(sock, addr):
+            await self.loop.sock_connect(sock, addr)
+
+            f = asyncio.ensure_future(
+                self.loop.sock_sendall(sock, b'helo' * (1024 * 1024 * 50)),
+                loop=self.loop)
+            self.loop.create_task(kill(f))
+            with self.assertRaises(asyncio.CancelledError):
+                await f
+            sock.close()
+            self.assertEqual(sock.fileno(), -1)
+
+        with self.tcp_server(srv_gen) as srv:
+
+            sock = socket.socket()
+            with sock:
+                sock.setblocking(False)
+                c = client(sock, srv.addr)
+                w = asyncio.wait_for(c, timeout=5.0, loop=self.loop)
+                self.loop.run_until_complete(w)
+
 
 class TestAIOSockets(_TestSockets, tb.AIOTestCase):
     pass

uvloop/loop.pxd

@@ -167,6 +167,8 @@ cdef class Loop:
     cdef _fileobj_to_fd(self, fileobj)
     cdef _ensure_fd_no_transport(self, fd)
 
+    cdef _new_reader_future(self, sock)
+    cdef _new_writer_future(self, sock)
     cdef _add_reader(self, fd, Handle handle)
     cdef _remove_reader(self, fd)
 
@@ -178,7 +180,7 @@ cdef class Loop:
     cdef _sock_sendall(self, fut, sock, data)
     cdef _sock_accept(self, fut, sock)
 
-    cdef _sock_connect(self, fut, sock, address)
+    cdef _sock_connect(self, sock, address)
     cdef _sock_connect_cb(self, fut, sock, address)
 
     cdef _sock_set_reuseport(self, int fd)

uvloop/loop.pyx

@@ -794,14 +794,28 @@ cdef class Loop:
         nr.query(addr, flags)
         return fut
 
+    cdef _new_reader_future(self, sock):
+        def _on_cancel(fut):
+            if fut.cancelled():
+                self._remove_reader(sock)
+
+        fut = self._new_future()
+        fut.add_done_callback(_on_cancel)
+        return fut
+
+    cdef _new_writer_future(self, sock):
+        def _on_cancel(fut):
+            if fut.cancelled():
+                self._remove_writer(sock)
+
+        fut = self._new_future()
+        fut.add_done_callback(_on_cancel)
+        return fut
+
     cdef _sock_recv(self, fut, sock, n):
         cdef:
             Handle handle
 
-        if fut.cancelled():
-            self._remove_reader(sock)
-            return
-
         try:
             data = sock.recv(n)
         except (BlockingIOError, InterruptedError):
@@ -819,10 +833,6 @@ cdef class Loop:
         cdef:
             Handle handle
 
-        if fut.cancelled():
-            self._remove_reader(sock)
-            return
-
         try:
             data = sock.recv_into(buf)
         except (BlockingIOError, InterruptedError):
@@ -841,10 +851,6 @@ cdef class Loop:
             Handle handle
             int n
 
-        if fut.cancelled():
-            self._remove_writer(sock)
-            return
-
         try:
             n = sock.send(data)
         except (BlockingIOError, InterruptedError):
@@ -878,10 +884,6 @@ cdef class Loop:
         cdef:
             Handle handle
 
-        if fut.cancelled():
-            self._remove_reader(sock)
-            return
-
         try:
             conn, address = sock.accept()
             conn.setblocking(False)
@@ -896,31 +898,29 @@ cdef class Loop:
             fut.set_result((conn, address))
             self._remove_reader(sock)
 
-    cdef _sock_connect(self, fut, sock, address):
+    cdef _sock_connect(self, sock, address):
         cdef:
             Handle handle
 
         try:
             sock.connect(address)
         except (BlockingIOError, InterruptedError):
-            # Issue #23618: When the C function connect() fails with EINTR, the
-            # 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(lambda fut: self._remove_writer(sock))
+            pass
+        else:
+            return
 
-            handle = new_MethodHandle3(
-                self,
-                "Loop._sock_connect",
-                <method3_t>self._sock_connect_cb,
-                self,
-                fut, sock, address)
+        fut = self._new_future()
+        fut.add_done_callback(lambda fut: self._remove_writer(sock))
 
-            self._add_writer(sock, handle)
-        except Exception as exc:
-            fut.set_exception(exc)
-        else:
-            fut.set_result(None)
+        handle = new_MethodHandle3(
+            self,
+            "Loop._sock_connect",
+            <method3_t>self._sock_connect_cb,
+            self,
+            fut, sock, address)
+
+        self._add_writer(sock, handle)
+        return fut
 
     cdef _sock_connect_cb(self, fut, sock, address):
         if fut.cancelled():
@@ -2087,7 +2087,7 @@ cdef class Loop:
         if self._debug and sock.gettimeout() != 0:
             raise ValueError("the socket must be non-blocking")
 
-        fut = self._new_future()
+        fut = self._new_reader_future(sock)
         handle = new_MethodHandle3(
             self,
             "Loop._sock_recv",
@@ -2112,7 +2112,7 @@ cdef class Loop:
         if self._debug and sock.gettimeout() != 0:
             raise ValueError("the socket must be non-blocking")
 
-        fut = self._new_future()
+        fut = self._new_reader_future(sock)
         handle = new_MethodHandle3(
             self,
             "Loop._sock_recv_into",
@@ -2162,7 +2162,7 @@ cdef class Loop:
                     data = memoryview(data)
                 data = data[n:]
 
-            fut = self._new_future()
+            fut = self._new_writer_future(sock)
             handle = new_MethodHandle3(
                 self,
                 "Loop._sock_sendall",
@@ -2191,7 +2191,7 @@ cdef class Loop:
         if self._debug and sock.gettimeout() != 0:
             raise ValueError("the socket must be non-blocking")
 
-        fut = self._new_future()
+        fut = self._new_reader_future(sock)
         handle = new_MethodHandle2(
             self,
             "Loop._sock_accept",
@@ -2212,15 +2212,13 @@ cdef class Loop:
 
         socket_inc_io_ref(sock)
         try:
-            fut = self._new_future()
             if sock.family == uv.AF_UNIX:
-                self._sock_connect(fut, sock, address)
+                fut = self._sock_connect(sock, address)
+            else:
+                _, _, _, _, address = (await self.getaddrinfo(*address[:2]))[0]
+                fut = self._sock_connect(sock, address)
+            if fut is not None:
                 await fut
-                return
-
-            _, _, _, _, address = (await self.getaddrinfo(*address[:2]))[0]
-            self._sock_connect(fut, sock, address)
-            await fut
         finally:
             socket_dec_io_ref(sock)