gh-146507: cache the buffer size for SelectorTransport.get_write_buffer_size() (#146518)

moktamd · web-flow · commit 1417737810e2 · 2026-03-27T21:43:15.000+05:30
diff --git a/Lib/asyncio/selector_events.py b/Lib/asyncio/selector_events.py
@@ -793,6 +793,7 @@ def __init__(self, loop, sock, protocol, extra=None, server=None, context=None):
 
         self._server = server
         self._buffer = collections.deque()
+        self._buffer_size = 0
         self._conn_lost = 0  # Set when call to connection_lost scheduled.
         self._closing = False  # Set when close() called.
         self._paused = False  # Set when pause_reading() called
@@ -897,6 +898,7 @@ def _force_close(self, exc):
             return
         if self._buffer:
             self._buffer.clear()
+            self._buffer_size = 0
             self._loop._remove_writer(self._sock_fd)
         if not self._closing:
             self._closing = True
@@ -919,7 +921,7 @@ def _call_connection_lost(self, exc):
                 self._server = None
 
     def get_write_buffer_size(self):
-        return sum(map(len, self._buffer))
+        return self._buffer_size
 
     def _add_reader(self, fd, callback, *args):
         if not self.is_reading():
@@ -1090,6 +1092,7 @@ def write(self, data):
 
         # Add it to the buffer.
         self._buffer.append(data)
+        self._buffer_size += len(data)
         self._maybe_pause_protocol()
 
     def _get_sendmsg_buffer(self):
@@ -1109,6 +1112,7 @@ def _write_sendmsg(self):
         except BaseException as exc:
             self._loop._remove_writer(self._sock_fd)
             self._buffer.clear()
+            self._buffer_size = 0
             self._fatal_error(exc, 'Fatal write error on socket transport')
             if self._empty_waiter is not None:
                 self._empty_waiter.set_exception(exc)
@@ -1124,6 +1128,7 @@ def _write_sendmsg(self):
                     self._sock.shutdown(socket.SHUT_WR)
 
     def _adjust_leftover_buffer(self, nbytes: int) -> None:
+        self._buffer_size -= nbytes
         buffer = self._buffer
         while nbytes:
             b = buffer.popleft()
@@ -1144,13 +1149,16 @@ def _write_send(self):
             if n != len(buffer):
                 # Not all data was written
                 self._buffer.appendleft(buffer[n:])
+            self._buffer_size -= n
         except (BlockingIOError, InterruptedError):
-            pass
+            self._buffer.appendleft(buffer)
+            return
         except (SystemExit, KeyboardInterrupt):
             raise
         except BaseException as exc:
             self._loop._remove_writer(self._sock_fd)
             self._buffer.clear()
+            self._buffer_size = 0
             self._fatal_error(exc, 'Fatal write error on socket transport')
             if self._empty_waiter is not None:
                 self._empty_waiter.set_exception(exc)
@@ -1186,7 +1194,9 @@ def writelines(self, list_of_data):
             self._conn_lost += 1
             return
 
-        self._buffer.extend([memoryview(data) for data in list_of_data])
+        for data in list_of_data:
+            self._buffer.append(memoryview(data))
+            self._buffer_size += len(data)
         self._write_ready()
         # If the entire buffer couldn't be written, register a write handler
         if self._buffer:
diff --git a/Misc/NEWS.d/next/Library/2026-03-27-12-00-00.gh-issue-146507.1D95A7.rst b/Misc/NEWS.d/next/Library/2026-03-27-12-00-00.gh-issue-146507.1D95A7.rst
@@ -0,0 +1,3 @@
+Make :meth:`asyncio.SelectorEventLoop` stream transport's
+:meth:`~asyncio.WriteTransport.get_write_buffer_size` O(1) by maintaining a
+running byte counter instead of iterating the buffer on every call.

Original file line number	Diff line number	Diff 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.`