gh-106527: asyncio: optimize to add/remove readers and writers (#106528)

Author: bdraco

Lib/asyncio/selector_events.py

@@ -274,9 +274,8 @@ def _ensure_fd_no_transport(self, fd):
     def _add_reader(self, fd, callback, *args):
         self._check_closed()
         handle = events.Handle(callback, args, self, None)
-        try:
-            key = self._selector.get_key(fd)
-        except KeyError:
+        key = self._selector.get_map().get(fd)
+        if key is None:
             self._selector.register(fd, selectors.EVENT_READ,
                                     (handle, None))
         else:
@@ -290,30 +289,27 @@ def _add_reader(self, fd, callback, *args):
     def _remove_reader(self, fd):
         if self.is_closed():
             return False
-        try:
-            key = self._selector.get_key(fd)
-        except KeyError:
+        key = self._selector.get_map().get(fd)
+        if key is None:
             return False
+        mask, (reader, writer) = key.events, key.data
+        mask &= ~selectors.EVENT_READ
+        if not mask:
+            self._selector.unregister(fd)
         else:
-            mask, (reader, writer) = key.events, key.data
-            mask &= ~selectors.EVENT_READ
-            if not mask:
-                self._selector.unregister(fd)
-            else:
-                self._selector.modify(fd, mask, (None, writer))
+            self._selector.modify(fd, mask, (None, writer))
 
-            if reader is not None:
-                reader.cancel()
-                return True
-            else:
-                return False
+        if reader is not None:
+            reader.cancel()
+            return True
+        else:
+            return False
 
     def _add_writer(self, fd, callback, *args):
         self._check_closed()
         handle = events.Handle(callback, args, self, None)
-        try:
-            key = self._selector.get_key(fd)
-        except KeyError:
+        key = self._selector.get_map().get(fd)
+        if key is None:
             self._selector.register(fd, selectors.EVENT_WRITE,
                                     (None, handle))
         else:
@@ -328,24 +324,22 @@ def _remove_writer(self, fd):
         """Remove a writer callback."""
         if self.is_closed():
             return False
-        try:
-            key = self._selector.get_key(fd)
-        except KeyError:
+        key = self._selector.get_map().get(fd)
+        if key is None:
             return False
+        mask, (reader, writer) = key.events, key.data
+        # Remove both writer and connector.
+        mask &= ~selectors.EVENT_WRITE
+        if not mask:
+            self._selector.unregister(fd)
         else:
-            mask, (reader, writer) = key.events, key.data
-            # Remove both writer and connector.
-            mask &= ~selectors.EVENT_WRITE
-            if not mask:
-                self._selector.unregister(fd)
-            else:
-                self._selector.modify(fd, mask, (reader, None))
+            self._selector.modify(fd, mask, (reader, None))
 
-            if writer is not None:
-                writer.cancel()
-                return True
-            else:
-                return False
+        if writer is not None:
+            writer.cancel()
+            return True
+        else:
+            return False
 
     def add_reader(self, fd, callback, *args):
         """Add a reader callback."""

Lib/test/test_asyncio/test_selector_events.py

@@ -178,7 +178,7 @@ def test_sock_connect_resolve_using_socket_params(self, m_gai):
                 sock.connect.assert_called_with(('127.0.0.1', 0))
 
     def test_add_reader(self):
-        self.loop._selector.get_key.side_effect = KeyError
+        self.loop._selector.get_map.return_value = {}
         cb = lambda: True
         self.loop.add_reader(1, cb)
 
@@ -192,8 +192,8 @@ def test_add_reader(self):
     def test_add_reader_existing(self):
         reader = mock.Mock()
         writer = mock.Mock()
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
-            1, 1, selectors.EVENT_WRITE, (reader, writer))
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
+            1, 1, selectors.EVENT_WRITE, (reader, writer))}
         cb = lambda: True
         self.loop.add_reader(1, cb)
 
@@ -208,8 +208,8 @@ def test_add_reader_existing(self):
 
     def test_add_reader_existing_writer(self):
         writer = mock.Mock()
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
-            1, 1, selectors.EVENT_WRITE, (None, writer))
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
+            1, 1, selectors.EVENT_WRITE, (None, writer))}
         cb = lambda: True
         self.loop.add_reader(1, cb)
 
@@ -222,18 +222,18 @@ def test_add_reader_existing_writer(self):
         self.assertEqual(writer, w)
 
     def test_remove_reader(self):
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
-            1, 1, selectors.EVENT_READ, (None, None))
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
+            1, 1, selectors.EVENT_READ, (None, None))}
         self.assertFalse(self.loop.remove_reader(1))
 
         self.assertTrue(self.loop._selector.unregister.called)
 
     def test_remove_reader_read_write(self):
         reader = mock.Mock()
         writer = mock.Mock()
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
             1, 1, selectors.EVENT_READ | selectors.EVENT_WRITE,
-            (reader, writer))
+            (reader, writer))}
         self.assertTrue(
             self.loop.remove_reader(1))
 
@@ -243,12 +243,12 @@ def test_remove_reader_read_write(self):
             self.loop._selector.modify.call_args[0])
 
     def test_remove_reader_unknown(self):
-        self.loop._selector.get_key.side_effect = KeyError
+        self.loop._selector.get_map.return_value = {}
         self.assertFalse(
             self.loop.remove_reader(1))
 
     def test_add_writer(self):
-        self.loop._selector.get_key.side_effect = KeyError
+        self.loop._selector.get_map.return_value = {}
         cb = lambda: True
         self.loop.add_writer(1, cb)
 
@@ -262,8 +262,8 @@ def test_add_writer(self):
     def test_add_writer_existing(self):
         reader = mock.Mock()
         writer = mock.Mock()
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
-            1, 1, selectors.EVENT_READ, (reader, writer))
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
+            1, 1, selectors.EVENT_READ, (reader, writer))}
         cb = lambda: True
         self.loop.add_writer(1, cb)
 
@@ -277,18 +277,18 @@ def test_add_writer_existing(self):
         self.assertEqual(cb, w._callback)
 
     def test_remove_writer(self):
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
-            1, 1, selectors.EVENT_WRITE, (None, None))
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
+            1, 1, selectors.EVENT_WRITE, (None, None))}
         self.assertFalse(self.loop.remove_writer(1))
 
         self.assertTrue(self.loop._selector.unregister.called)
 
     def test_remove_writer_read_write(self):
         reader = mock.Mock()
         writer = mock.Mock()
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
             1, 1, selectors.EVENT_READ | selectors.EVENT_WRITE,
-            (reader, writer))
+            (reader, writer))}
         self.assertTrue(
             self.loop.remove_writer(1))
 
@@ -298,7 +298,7 @@ def test_remove_writer_read_write(self):
             self.loop._selector.modify.call_args[0])
 
     def test_remove_writer_unknown(self):
-        self.loop._selector.get_key.side_effect = KeyError
+        self.loop._selector.get_map.return_value = {}
         self.assertFalse(
             self.loop.remove_writer(1))
 

Misc/NEWS.d/next/Library/2023-07-07-18-22-07.gh-issue-106527.spHQ0W.rst

@@ -0,0 +1 @@
+Reduce overhead to add and remove :mod:`asyncio` readers and writers.