Skip to content

Commit 1417737

Browse files
authored
gh-146507: cache the buffer size for SelectorTransport.get_write_buffer_size() (#146518)
1 parent c8ee196 commit 1417737

File tree

2 files changed

+16
-3
lines changed

2 files changed

+16
-3
lines changed

Lib/asyncio/selector_events.py

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -793,6 +793,7 @@ def __init__(self, loop, sock, protocol, extra=None, server=None, context=None):
793793

794794
self._server = server
795795
self._buffer = collections.deque()
796+
self._buffer_size = 0
796797
self._conn_lost = 0 # Set when call to connection_lost scheduled.
797798
self._closing = False # Set when close() called.
798799
self._paused = False # Set when pause_reading() called
@@ -897,6 +898,7 @@ def _force_close(self, exc):
897898
return
898899
if self._buffer:
899900
self._buffer.clear()
901+
self._buffer_size = 0
900902
self._loop._remove_writer(self._sock_fd)
901903
if not self._closing:
902904
self._closing = True
@@ -919,7 +921,7 @@ def _call_connection_lost(self, exc):
919921
self._server = None
920922

921923
def get_write_buffer_size(self):
922-
return sum(map(len, self._buffer))
924+
return self._buffer_size
923925

924926
def _add_reader(self, fd, callback, *args):
925927
if not self.is_reading():
@@ -1090,6 +1092,7 @@ def write(self, data):
10901092

10911093
# Add it to the buffer.
10921094
self._buffer.append(data)
1095+
self._buffer_size += len(data)
10931096
self._maybe_pause_protocol()
10941097

10951098
def _get_sendmsg_buffer(self):
@@ -1109,6 +1112,7 @@ def _write_sendmsg(self):
11091112
except BaseException as exc:
11101113
self._loop._remove_writer(self._sock_fd)
11111114
self._buffer.clear()
1115+
self._buffer_size = 0
11121116
self._fatal_error(exc, 'Fatal write error on socket transport')
11131117
if self._empty_waiter is not None:
11141118
self._empty_waiter.set_exception(exc)
@@ -1124,6 +1128,7 @@ def _write_sendmsg(self):
11241128
self._sock.shutdown(socket.SHUT_WR)
11251129

11261130
def _adjust_leftover_buffer(self, nbytes: int) -> None:
1131+
self._buffer_size -= nbytes
11271132
buffer = self._buffer
11281133
while nbytes:
11291134
b = buffer.popleft()
@@ -1144,13 +1149,16 @@ def _write_send(self):
11441149
if n != len(buffer):
11451150
# Not all data was written
11461151
self._buffer.appendleft(buffer[n:])
1152+
self._buffer_size -= n
11471153
except (BlockingIOError, InterruptedError):
1148-
pass
1154+
self._buffer.appendleft(buffer)
1155+
return
11491156
except (SystemExit, KeyboardInterrupt):
11501157
raise
11511158
except BaseException as exc:
11521159
self._loop._remove_writer(self._sock_fd)
11531160
self._buffer.clear()
1161+
self._buffer_size = 0
11541162
self._fatal_error(exc, 'Fatal write error on socket transport')
11551163
if self._empty_waiter is not None:
11561164
self._empty_waiter.set_exception(exc)
@@ -1186,7 +1194,9 @@ def writelines(self, list_of_data):
11861194
self._conn_lost += 1
11871195
return
11881196

1189-
self._buffer.extend([memoryview(data) for data in list_of_data])
1197+
for data in list_of_data:
1198+
self._buffer.append(memoryview(data))
1199+
self._buffer_size += len(data)
11901200
self._write_ready()
11911201
# If the entire buffer couldn't be written, register a write handler
11921202
if self._buffer:
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Make :meth:`asyncio.SelectorEventLoop` stream transport's
2+
:meth:`~asyncio.WriteTransport.get_write_buffer_size` O(1) by maintaining a
3+
running byte counter instead of iterating the buffer on every call.

0 commit comments

Comments
 (0)