Skip to content

Commit

Permalink
pythongh-106527: asyncio: optimize to add/remove readers and writers (p…
Browse files Browse the repository at this point in the history 
…ython#106528)
  • Loading branch information
@bdraco @jtcave
bdraco authored and jtcave committed Jul 23, 2023
1 parent 5b02245 commit 636f995
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 53 deletions.
64 changes: 29 additions & 35 deletions Lib/asyncio/selector_events.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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."""
Expand Down
36 changes: 18 additions & 18 deletions Lib/test/test_asyncio/test_selector_events.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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)

Expand All @@ -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)

Expand All @@ -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))

Expand All @@ -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)

Expand All @@ -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)

Expand All @@ -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))

Expand All @@ -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))

Expand Down
1 change: 1 addition & 0 deletions Misc/NEWS.d/next/Library/2023-07-07-18-22-07.gh-issue-106527.spHQ0W.rst
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Reduce overhead to add and remove :mod:`asyncio` readers and writers.

0 comments on commit 636f995

Please sign in to comment.