WIP sendmsg

kumaraditya303 · Dec 24, 2022 · Mar 14, 2022 · Mar 15, 2022 · Mar 15, 2022 · Mar 15, 2022
commit abd2dc31d2a04a525aca5534730aa0017cc17f04
@@ -899,7 +899,10 @@ def __init__(self, loop, sock, protocol, waiter=None,
         self._eof = False
         self._paused = False
         self._empty_waiter = None
-
+        if hasattr(socket.socket, 'sendmsg'):
+            self._write_ready = self._write_sendmsg
+        else:
+            self._write_ready = self._write_send
         # Disable the Nagle algorithm -- small writes will be
         # sent without waiting for the TCP ACK.  This generally
         # decreases the latency (in some cases significantly.)
@@ -1067,7 +1070,46 @@ def write(self, data):
         self._buffer.append(data)
         self._maybe_pause_protocol()
 
-    def _write_ready(self):
+    def _write_sendmsg(self):
+        assert self._buffer, 'Data should not be empty'
+        if self._conn_lost:
+            return
+        try:
+            n = self._sock.sendmsg(self._buffer)
+            self._adjust_leftover_buffer(n)
+        except (BlockingIOError, InterruptedError):
+            pass
+        except (SystemExit, KeyboardInterrupt):
+            raise
+        except BaseException as exc:
+            self._loop._remove_writer(self._sock_fd)
+            self._buffer.clear()
+            self._fatal_error(exc, 'Fatal write error on socket transport')
+            if self._empty_waiter is not None:
+                self._empty_waiter.set_exception(exc)
+        else:
+            self._maybe_resume_protocol()  # May append to buffer.
+            if not self._buffer:
+                self._loop._remove_writer(self._sock_fd)
+                if self._empty_waiter is not None:
+                    self._empty_waiter.set_result(None)
+                if self._closing:
+                    self._call_connection_lost(None)
+                elif self._eof:
+                    self._sock.shutdown(socket.SHUT_WR)
+
+    def _adjust_leftover_buffer(self, n: int, /) -> None:
+        buffer = self._buffer
+        while n:
+            b = buffer.popleft()
+            b_len = len(b)
+            if b_len <= n:
+                n -= b_len
+            else:
+                buffer.appendleft(b[n:])
+                break
+
+    def _write_send(self):
         assert self._buffer, 'Data should not be empty'
         if self._conn_lost:
             return
@@ -1105,6 +1147,15 @@ def write_eof(self):
         if not self._buffer:
             self._sock.shutdown(socket.SHUT_WR)
 
+    def writelines(self, list_of_data):
+        hasbuffer = len(self._buffer)
+        self._buffer.extend([memoryview(i) for i in list_of_data])
+        if not hasbuffer:
+            # Optimization: try to send now
+            self._write_ready()
+            return
+        self._maybe_pause_protocol()
+
     def can_write_eof(self):
         return True
 

@@ -497,9 +497,13 @@ def setUp(self):
         self.sock = mock.Mock(socket.socket)
         self.sock_fd = self.sock.fileno.return_value = 7
 
-    def socket_transport(self, waiter=None):
+    def socket_transport(self, waiter=None, sendmsg=False):
         transport = _SelectorSocketTransport(self.loop, self.sock,
                                              self.protocol, waiter=waiter)
+        if sendmsg:
+            transport._write_ready = transport._write_sendmsg
+        else:
+            transport._write_ready = transport._write_send
         self.addCleanup(close_transport, transport)
         return transport
 
@@ -733,6 +737,40 @@ def test_write_tryagain(self):
         self.loop.assert_writer(7, transport._write_ready)
         self.assertEqual(list_to_buffer([b'data']), transport._buffer)
 
+    def test_write_sendmsg_no_data(self):
+        self.sock.sendmsg = mock.Mock()
+        self.sock.sendmsg.return_value = 0
+        transport = self.socket_transport(sendmsg=True)
+        transport._buffer.append(memoryview(b'data'))
+        transport.write(b'')
+        self.assertFalse(self.sock.sendmsg.called)
+        self.assertEqual(list_to_buffer([b'data']), transport._buffer)
+
+    def test_write_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._buffer.append(data)
+        self.loop._add_writer(7, transport._write_ready)
+        transport._write_ready()
+        self.assertTrue(self.sock.sendmsg.called)
+        self.assertFalse(self.loop.writers)
+
+    def test_write_sendmsg_partial(self):
+        data = memoryview(b'data')
+        self.sock.sendmsg = mock.Mock()
+        # Sent partial data
+        self.sock.sendmsg.return_value = len(data) // 2
+
+        transport = self.socket_transport(sendmsg=True)
+        transport._buffer.append(data)
+        self.loop._add_writer(7, transport._write_ready)
+        transport._write_ready()
+        self.assertTrue(self.sock.sendmsg.called)
+        self.assertTrue(self.loop.writers)
+
     @mock.patch('asyncio.selector_events.logger')
     def test_write_exception(self, m_log):
         err = self.sock.send.side_effect = OSError()