Skip to content

Commit

Permalink
gh-103462: Ensure SelectorSocketTransport.writelines registers a writ…
Browse files Browse the repository at this point in the history
…er when data is still pending (#103463)
  • Loading branch information
alisaifee committed Apr 13, 2023
1 parent 9e67740 commit 19d2639
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 0 deletions.
3 changes: 3 additions & 0 deletions Lib/asyncio/selector_events.py
Original file line number Diff line number Diff line change
Expand Up @@ -1176,6 +1176,9 @@ def writelines(self, list_of_data):
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
if self._buffer:
self._loop._add_writer(self._sock_fd, self._write_ready)

def can_write_eof(self):
return True
Expand Down
42 changes: 42 additions & 0 deletions Lib/test/test_asyncio/test_selector_events.py
Original file line number Diff line number Diff line change
Expand Up @@ -747,6 +747,48 @@ def test_write_sendmsg_no_data(self):
self.assertFalse(self.sock.sendmsg.called)
self.assertEqual(list_to_buffer([b'data']), transport._buffer)

@unittest.skipUnless(selector_events._HAS_SENDMSG, 'no sendmsg')
def test_writelines_sendmsg_full(self):
data = memoryview(b'data')
self.sock.sendmsg = mock.Mock()
self.sock.sendmsg.return_value = len(data)

transport = self.socket_transport(sendmsg=True)
transport.writelines([data])
self.assertTrue(self.sock.sendmsg.called)
self.assertFalse(self.loop.writers)

@unittest.skipUnless(selector_events._HAS_SENDMSG, 'no sendmsg')
def test_writelines_sendmsg_partial(self):
data = memoryview(b'data')
self.sock.sendmsg = mock.Mock()
self.sock.sendmsg.return_value = 2

transport = self.socket_transport(sendmsg=True)
transport.writelines([data])
self.assertTrue(self.sock.sendmsg.called)
self.assertTrue(self.loop.writers)

def test_writelines_send_full(self):
data = memoryview(b'data')
self.sock.send.return_value = len(data)
self.sock.send.fileno.return_value = 7

transport = self.socket_transport()
transport.writelines([data])
self.assertTrue(self.sock.send.called)
self.assertFalse(self.loop.writers)

def test_writelines_send_partial(self):
data = memoryview(b'data')
self.sock.send.return_value = 2
self.sock.send.fileno.return_value = 7

transport = self.socket_transport()
transport.writelines([data])
self.assertTrue(self.sock.send.called)
self.assertTrue(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,4 @@
Fixed an issue with using :meth:`~asyncio.WriteTransport.writelines` in :mod:`asyncio` to send very
large payloads that exceed the amount of data that can be written in one
call to :meth:`socket.socket.send` or :meth:`socket.socket.sendmsg`,
resulting in the remaining buffer being left unwritten.

0 comments on commit 19d2639

Please sign in to comment.