gh-103462: Ensure SelectorSocketTransport.writelines registers a writer when data is still pending (#103463)

2 files changed, 45 insertions, 0 deletions

diff --git a/Lib/asyncio/selector_events.py b/Lib/asyncio/selector_events.py
index de5076a..3a69712 100644
--- a/Lib/asyncio/selector_events.py
+++ b/Lib/asyncio/selector_events.py
@@ -1176,6 +1176,9 @@ class _SelectorSocketTransport(_SelectorTransport):
             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
diff --git a/Lib/test/test_asyncio/test_selector_events.py b/Lib/test/test_asyncio/test_selector_events.py
index 921c98a..e41341f 100644
--- a/Lib/test/test_asyncio/test_selector_events.py
+++ b/Lib/test/test_asyncio/test_selector_events.py
@@ -748,6 +748,48 @@ class SelectorSocketTransportTests(test_utils.TestCase):
         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')
         self.sock.sendmsg = mock.Mock()