Skip to content

Commit 8ea34d2

Browse files
committed
bpo-40007: Make asyncio.transport.writelines on selector use sendmsg
1 parent 56bfdeb commit 8ea34d2

File tree

2 files changed

+100
-3
lines changed

2 files changed

+100
-3
lines changed

Lib/asyncio/selector_events.py

Lines changed: 99 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@
2929
from .log import logger
3030

3131

32+
sendmsg = getattr(socket.socket, "sendmsg", False)
33+
34+
3235
def _test_selector_event(selector, fd, event):
3336
# Test if the selector is monitoring 'event' events
3437
# for the file descriptor 'fd'.
@@ -746,6 +749,7 @@ def _add_reader(self, fd, callback, *args):
746749

747750
class _SelectorSocketTransport(_SelectorTransport):
748751

752+
_buffer_factory = list
749753
_start_tls_compatible = True
750754
_sendfile_compatible = constants._SendfileMode.TRY_NATIVE
751755

@@ -921,16 +925,76 @@ def write(self, data):
921925
self._loop._add_writer(self._sock_fd, self._write_ready)
922926

923927
# Add it to the buffer.
924-
self._buffer.extend(data)
928+
self._buffer.append(data)
929+
self._maybe_pause_protocol()
930+
931+
@staticmethod
932+
def _calculate_leftovers(n, items):
933+
leftovers = []
934+
whole = False
935+
for item in items:
936+
if whole:
937+
leftovers.append(item)
938+
continue
939+
n -= len(item)
940+
if n >= 0:
941+
continue
942+
leftovers.append(memoryview(item)[n:])
943+
whole = True
944+
return leftovers
945+
946+
def writelines(self, lines):
947+
if not sendmsg:
948+
return self.write(b''.join(lines))
949+
if self._eof:
950+
raise RuntimeError('Cannot call write() after write_eof()')
951+
if self._empty_waiter is not None:
952+
raise RuntimeError('unable to write; sendfile is in progress')
953+
if not lines:
954+
return
955+
956+
if self._conn_lost:
957+
if self._conn_lost >= constants.LOG_THRESHOLD_FOR_CONNLOST_WRITES:
958+
logger.warning('socket.send() raised exception.')
959+
self._conn_lost += 1
960+
return
961+
962+
if not self._buffer:
963+
# Optimization: try to send now.
964+
try:
965+
n = self._sock.sendmsg(lines)
966+
except OSError:
967+
return self.write(b''.join(lines))
968+
except (BlockingIOError, InterruptedError):
969+
pass
970+
except (SystemExit, KeyboardInterrupt):
971+
raise
972+
except BaseException as exc:
973+
self._fatal_error(exc, 'Fatal write error on socket transport')
974+
return
975+
else:
976+
lines = self._calculate_leftovers(n, lines)
977+
if not lines:
978+
return
979+
# Not all was written; register write handler.
980+
self._loop._add_writer(self._sock_fd, self._write_ready)
981+
982+
# Add it to the buffer.
983+
self._buffer.extend(lines)
925984
self._maybe_pause_protocol()
926985

927986
def _write_ready(self):
928987
assert self._buffer, 'Data should not be empty'
929988

930989
if self._conn_lost:
931990
return
991+
992+
if sendmsg:
993+
return self._write_vectored_self()
994+
932995
try:
933-
n = self._sock.send(self._buffer)
996+
tmp = b''.join(self._buffer)
997+
n = self._sock.send(tmp)
934998
except (BlockingIOError, InterruptedError):
935999
pass
9361000
except (SystemExit, KeyboardInterrupt):
@@ -943,7 +1007,36 @@ def _write_ready(self):
9431007
self._empty_waiter.set_exception(exc)
9441008
else:
9451009
if n:
946-
del self._buffer[:n]
1010+
self._buffer = [tmp[:n]]
1011+
self._maybe_resume_protocol() # May append to buffer.
1012+
if not self._buffer:
1013+
self._loop._remove_writer(self._sock_fd)
1014+
if self._empty_waiter is not None:
1015+
self._empty_waiter.set_result(None)
1016+
if self._closing:
1017+
self._call_connection_lost(None)
1018+
elif self._eof:
1019+
self._sock.shutdown(socket.SHUT_WR)
1020+
1021+
def _write_vectored_self(self):
1022+
try:
1023+
try:
1024+
n = self._sock.sendmsg(self._buffer)
1025+
except OSError:
1026+
self._buffer = [b''.join(self._buffer)]
1027+
n = self._sock.sendmsg(self._buffer)
1028+
except (BlockingIOError, InterruptedError):
1029+
pass
1030+
except (SystemExit, KeyboardInterrupt):
1031+
raise
1032+
except BaseException as exc:
1033+
self._loop._remove_writer(self._sock_fd)
1034+
self._buffer.clear()
1035+
self._fatal_error(exc, 'Fatal write error on socket transport')
1036+
if self._empty_waiter is not None:
1037+
self._empty_waiter.set_exception(exc)
1038+
else:
1039+
self._buffer = self._calculate_leftovers(n, self._buffer)
9471040
self._maybe_resume_protocol() # May append to buffer.
9481041
if not self._buffer:
9491042
self._loop._remove_writer(self._sock_fd)
@@ -954,6 +1047,9 @@ def _write_ready(self):
9541047
elif self._eof:
9551048
self._sock.shutdown(socket.SHUT_WR)
9561049

1050+
def get_write_buffer_size(self):
1051+
return sum(len(data) for data in self._buffer)
1052+
9571053
def write_eof(self):
9581054
if self._closing or self._eof:
9591055
return
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Make asyncio.transport.writelines on selector use sendmsg

0 commit comments

Comments
 (0)