Add cancellation callback

Author: asvetlov

Lib/asyncio/unix_events.py

@@ -328,13 +328,13 @@ async def _sock_sendfile_native(self, sock, file, offset, count):
             return 0  # empty file
 
         fut = self.create_future()
-        fd = sock.fileno()
-        self._sock_sendfile_native_impl(fut, None, fd, fileno,
+        self._sock_sendfile_native_impl(fut, None, sock, fileno,
                                         offset, count, blocksize, 0)
         return await fut
 
-    def _sock_sendfile_native_impl(self, fut, registered_fd, fd, fileno,
+    def _sock_sendfile_native_impl(self, fut, registered_fd, sock, fileno,
                                    offset, count, blocksize, total_sent):
+        fd = sock.fileno()
         if registered_fd is not None:
             # Remove the callback early.  It should be rare that the
             # selector says the fd is ready but the call still returns
@@ -354,8 +354,11 @@ def _sock_sendfile_native_impl(self, fut, registered_fd, fd, fileno,
         try:
             sent = os.sendfile(fd, fileno, offset, blocksize)
         except (BlockingIOError, InterruptedError):
-            self.add_writer(fd, self._sock_sendfile_native_impl, fut, fd, fd,
-                            fileno, offset, count, blocksize. total_sent)
+            if registered_fd is None:
+                self._sock_add_cancellation_callback(fut, sock)
+            self.add_writer(fd, self._sock_sendfile_native_impl, fut,
+                            fd, sock, fileno,
+                            offset, count, blocksize. total_sent)
         except OSError as exc:
             if total_sent == 0:
                 # We can get here for different reasons, the main
@@ -380,14 +383,24 @@ def _sock_sendfile_native_impl(self, fut, registered_fd, fd, fileno,
             else:
                 offset += sent
                 total_sent += sent
+                if registered_fd is None:
+                    self._sock_add_cancellation_callback(fut, sock)
                 self.add_writer(fd, self._sock_sendfile_native_impl, fut,
-                                fd, fd, fileno,
+                                fd, sock, fileno,
                                 offset, count, blocksize, total_sent)
 
     def _sock_sendfile_update_filepos(self, fileno, offset, total_sent):
         if total_sent > 0:
             os.lseek(fileno, offset, os.SEEK_SET)
 
+    def _sock_add_cancellation_callback(self, fut, sock):
+        def cb(fut):
+            if fut.cancelled():
+                fd = sock.fileno()
+                if fd!= -1:
+                    self.remove_writer(fd)
+        fut.add_done_callback(cb)
+
 
 class _UnixReadPipeTransport(transports.ReadTransport):
 

Lib/test/test_asyncio/test_unix_events.py

@@ -430,9 +430,11 @@ def __init__(self, loop):
             self.closed = False
             self.data = bytearray()
             self.fut = loop.create_future()
+            self.transport = None
 
         def connection_made(self, transport):
             self.started = True
+            self.transport = transport
 
         def data_received(self, data):
             self.data.extend(data)
@@ -480,6 +482,9 @@ def prepare(self):
         self.run_loop(self.loop.sock_connect(sock, (support.HOST, port)))
 
         def cleanup():
+            proto.transport.close()
+            self.run_loop(proto.wait_closed())
+
             server.close()
             self.run_loop(server.wait_closed())