bpo-41317: Remove reader on cancellation in asyncio.loop.sock_accept() (GH-21595)

(cherry picked from commit 0dd98c2d00a75efbec19c2ed942923981bc06683)

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>

3 files changed, 27 insertions, 7 deletions

diff --git a/Lib/asyncio/selector_events.py b/Lib/asyncio/selector_events.py
index 884a58f..8495a79 100644
--- a/Lib/asyncio/selector_events.py
+++ b/Lib/asyncio/selector_events.py
@@ -555,20 +555,19 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
         if self._debug and sock.gettimeout() != 0:
             raise ValueError("the socket must be non-blocking")
         fut = self.create_future()
-        self._sock_accept(fut, False, sock)
+        self._sock_accept(fut, sock)
         return await fut
 
-    def _sock_accept(self, fut, registered, sock):
+    def _sock_accept(self, fut, sock):
         fd = sock.fileno()
-        if registered:
-            self.remove_reader(fd)
-        if fut.done():
-            return
         try:
             conn, address = sock.accept()
             conn.setblocking(False)
         except (BlockingIOError, InterruptedError):
-            self.add_reader(fd, self._sock_accept, fut, True, sock)
+            self._ensure_fd_no_transport(fd)
+            handle = self._add_reader(fd, self._sock_accept, fut, sock)
+            fut.add_done_callback(
+                functools.partial(self._sock_read_done, fd, handle=handle))
         except (SystemExit, KeyboardInterrupt):
             raise
         except BaseException as exc:
diff --git a/Lib/test/test_asyncio/test_sock_lowlevel.py b/Lib/test/test_asyncio/test_sock_lowlevel.py
index e339ee9..d8a5df8 100644
--- a/Lib/test/test_asyncio/test_sock_lowlevel.py
+++ b/Lib/test/test_asyncio/test_sock_lowlevel.py
@@ -415,6 +415,25 @@ class BaseSockTestsMixin:
         conn.close()
         listener.close()
 
+    def test_cancel_sock_accept(self):
+        listener = socket.socket()
+        listener.setblocking(False)
+        listener.bind(('127.0.0.1', 0))
+        listener.listen(1)
+        sockaddr = listener.getsockname()
+        f = asyncio.wait_for(self.loop.sock_accept(listener), 0.1)
+        with self.assertRaises(asyncio.TimeoutError):
+            self.loop.run_until_complete(f)
+
+        listener.close()
+        client = socket.socket()
+        client.setblocking(False)
+        f = self.loop.sock_connect(client, sockaddr)
+        with self.assertRaises(ConnectionRefusedError):
+            self.loop.run_until_complete(f)
+
+        client.close()
+
     def test_create_connection_sock(self):
         with test_utils.run_test_server() as httpd:
             sock = None
diff --git a/Misc/NEWS.d/next/Library/2020-07-23-01-18-34.bpo-41317.O17Z6x.rst b/Misc/NEWS.d/next/Library/2020-07-23-01-18-34.bpo-41317.O17Z6x.rst
new file mode 100644
index 0000000..1af985e
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2020-07-23-01-18-34.bpo-41317.O17Z6x.rst
@@ -0,0 +1,2 @@
+Use add_done_callback() in asyncio.loop.sock_accept() to unsubscribe reader
+early on cancellation.