Skip to content

gh-136234: Fix _SelectorSocketTransport.writelines to be robust to connection loss #136743

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 11 additions & 3 deletions Lib/asyncio/selector_events.py
Original file line number Diff line number Diff line change
Expand Up @@ -1048,6 +1048,11 @@ def _read_ready__on_eof(self):
else:
self.close()

def _write_after_conn_lost(self):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is only used at two places, I would prefer this to be inlined rather than as a new function.

if self._conn_lost >= constants.LOG_THRESHOLD_FOR_CONNLOST_WRITES:
logger.warning('socket.send() raised exception.')
self._conn_lost += 1

def write(self, data):
if not isinstance(data, (bytes, bytearray, memoryview)):
raise TypeError(f'data argument must be a bytes-like object, '
Expand All @@ -1060,9 +1065,7 @@ def write(self, data):
return

if self._conn_lost:
if self._conn_lost >= constants.LOG_THRESHOLD_FOR_CONNLOST_WRITES:
logger.warning('socket.send() raised exception.')
self._conn_lost += 1
self._write_after_conn_lost()
return

if not self._buffer:
Expand Down Expand Up @@ -1174,6 +1177,11 @@ def writelines(self, list_of_data):
raise RuntimeError('unable to writelines; sendfile is in progress')
if not list_of_data:
return

if self._conn_lost:
self._write_after_conn_lost()
return

self._buffer.extend([memoryview(data) for data in list_of_data])
self._write_ready()
# If the entire buffer couldn't be written, register a write handler
Expand Down
16 changes: 16 additions & 0 deletions Lib/test/test_asyncio/test_selector_events.py
Original file line number Diff line number Diff line change
Expand Up @@ -854,6 +854,22 @@ def test_writelines_pauses_protocol(self):
self.assertTrue(self.sock.send.called)
self.assertTrue(self.loop.writers)

def test_writelines_after_connection_lost(self):
# GH-136234
transport = self.socket_transport()
self.sock.send = mock.Mock()
self.sock.send.side_effect = ConnectionResetError
transport.write(b'data1') # Will fail immediately, causing connection lost

transport.writelines([b'data2'])
self.assertFalse(transport._buffer)
self.assertFalse(self.loop.writers)

test_utils.run_briefly(self.loop) # Allow _call_connection_lost to run
transport.writelines([b'data2'])
self.assertFalse(transport._buffer)
self.assertFalse(self.loop.writers)

@unittest.skipUnless(selector_events._HAS_SENDMSG, 'no sendmsg')
def test_write_sendmsg_full(self):
data = memoryview(b'data')
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fix :meth:`asyncio.WriteTransport.writelines` to be robust to connection
failure, by using the same behavior as :meth:`~asyncio.WriteTransport.write`.
Loading