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

alisaifee · aisk · commit 44b69f67de2e · 2023-04-18T16:24:26.000+08:00
diff --git a/Lib/asyncio/selector_events.py b/Lib/asyncio/selector_events.py
@@ -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
diff --git a/Lib/test/test_asyncio/test_selector_events.py b/Lib/test/test_asyncio/test_selector_events.py
@@ -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')
diff --git a/Misc/NEWS.d/next/Library/2023-04-12-06-00-02.gh-issue-103462.w6yBlM.rst b/Misc/NEWS.d/next/Library/2023-04-12-06-00-02.gh-issue-103462.w6yBlM.rst
@@ -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.
