From 0dd98c2d00a75efbec19c2ed942923981bc06683 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alex=20Gr=C3=B6nholm?= <alex.gronholm@nextday.fi>
Date: Thu, 23 Jul 2020 22:45:08 +0300
Subject: [PATCH] bpo-41317: Remove reader on cancellation in
 asyncio.loop.sock_accept() (#21595)

---
 Lib/asyncio/selector_events.py                | 13 ++++++-------
 Lib/test/test_asyncio/test_sock_lowlevel.py   | 19 +++++++++++++++++++
 .../2020-07-23-01-18-34.bpo-41317.O17Z6x.rst  |  2 ++
 3 files changed, 27 insertions(+), 7 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2020-07-23-01-18-34.bpo-41317.O17Z6x.rst

diff --git a/Lib/asyncio/selector_events.py b/Lib/asyncio/selector_events.py
index 884a58f2ed6..8495a7901cd 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 e339ee9a4fc..d8a5df8ede1 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 00000000000..1af985e90e3
--- /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.