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

This commit is contained in:
J. Nick Koston 2023-07-22 16:07:40 -05:00 committed by GitHub
parent ee15844db8
commit b7dc795dfd
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 48 additions and 53 deletions

View file

@ -274,9 +274,8 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
def _add_reader(self, fd, callback, *args): def _add_reader(self, fd, callback, *args):
self._check_closed() self._check_closed()
handle = events.Handle(callback, args, self, None) handle = events.Handle(callback, args, self, None)
try: key = self._selector.get_map().get(fd)
key = self._selector.get_key(fd) if key is None:
except KeyError:
self._selector.register(fd, selectors.EVENT_READ, self._selector.register(fd, selectors.EVENT_READ,
(handle, None)) (handle, None))
else: else:
@ -290,11 +289,9 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
def _remove_reader(self, fd): def _remove_reader(self, fd):
if self.is_closed(): if self.is_closed():
return False return False
try: key = self._selector.get_map().get(fd)
key = self._selector.get_key(fd) if key is None:
except KeyError:
return False return False
else:
mask, (reader, writer) = key.events, key.data mask, (reader, writer) = key.events, key.data
mask &= ~selectors.EVENT_READ mask &= ~selectors.EVENT_READ
if not mask: if not mask:
@ -311,9 +308,8 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
def _add_writer(self, fd, callback, *args): def _add_writer(self, fd, callback, *args):
self._check_closed() self._check_closed()
handle = events.Handle(callback, args, self, None) handle = events.Handle(callback, args, self, None)
try: key = self._selector.get_map().get(fd)
key = self._selector.get_key(fd) if key is None:
except KeyError:
self._selector.register(fd, selectors.EVENT_WRITE, self._selector.register(fd, selectors.EVENT_WRITE,
(None, handle)) (None, handle))
else: else:
@ -328,11 +324,9 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
"""Remove a writer callback.""" """Remove a writer callback."""
if self.is_closed(): if self.is_closed():
return False return False
try: key = self._selector.get_map().get(fd)
key = self._selector.get_key(fd) if key is None:
except KeyError:
return False return False
else:
mask, (reader, writer) = key.events, key.data mask, (reader, writer) = key.events, key.data
# Remove both writer and connector. # Remove both writer and connector.
mask &= ~selectors.EVENT_WRITE mask &= ~selectors.EVENT_WRITE

View file

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

View file

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