From 3fa281625fe2d9e251b567f5e58844c195f8efd0 Mon Sep 17 00:00:00 2001 From: Flowseal Date: Tue, 6 Oct 2026 23:21:03 +0300 Subject: [PATCH] Optimizations --- proxy/bridge.py | 7 +++++-- proxy/h2_transport.py | 30 ++++++++++++++++++++++-------- tests/test_h2_transport.py | 4 ++-- 3 files changed, 29 insertions(+), 12 deletions(-) diff --git a/proxy/bridge.py b/proxy/bridge.py index c2db7f6..de64dc7 100644 --- a/proxy/bridge.py +++ b/proxy/bridge.py @@ -72,11 +72,14 @@ class MsgSplitter: if packet_len is None: break if packet_len <= 0: - parts.append(bytes(self._cipher_buf[offset:])) + parts.append(bytes(self._cipher_buf) if offset == 0 else bytes(self._cipher_buf[offset:])) offset = buf_len self._disabled = True break - parts.append(bytes(self._cipher_buf[offset:offset + packet_len])) + if offset == 0 and packet_len == buf_len: + parts.append(bytes(self._cipher_buf)) + else: + parts.append(bytes(self._cipher_buf[offset:offset + packet_len])) offset += packet_len if offset: diff --git a/proxy/h2_transport.py b/proxy/h2_transport.py index 7bcfa50..130d44e 100644 --- a/proxy/h2_transport.py +++ b/proxy/h2_transport.py @@ -20,6 +20,7 @@ from .utils import create_ssl_context STREAM_RECEIVE_WINDOW = 256 * 1024 CONNECTION_RECEIVE_WINDOW = 4 * 1024 * 1024 +IO_BATCH_SIZE = 64 * 1024 @lru_cache(maxsize=1) @@ -36,6 +37,7 @@ class _ResponseStream(httpx.AsyncByteStream): self.read_timeout = read_timeout self.changed = asyncio.Event() self.chunks = deque() + self.consumed_credit = 0 self.headers = None self.response_started = None self.error = None @@ -100,6 +102,8 @@ class _Connection: data = self.h2.data_to_send() if data: self.writer.write(data) + return True + return False async def drain(self): try: @@ -127,11 +131,12 @@ class _Connection: async def read_loop(self): try: while self.error is None: - data = await self.reader.read(64 * 1024) + data = await self.reader.read(IO_BATCH_SIZE) if not data: raise httpx.ReadError('H2 peer closed the connection') events = self.h2.receive_data(data) terminated = any(isinstance(e, ConnectionTerminated) for e in events) + connection_credit = 0 for event in events: if isinstance(event, RemoteSettingsChanged): self.settings_received = True @@ -145,7 +150,7 @@ class _Connection: stream = self.streams.get(getattr(event, 'stream_id', None)) if isinstance(event, DataReceived): if event.flow_controlled_length and not terminated: - self.h2.increment_flow_control_window(event.flow_controlled_length) + connection_credit += event.flow_controlled_length if stream is not None: if event.data: stream.chunks.append((event.data, event.flow_controlled_length)) @@ -164,8 +169,10 @@ class _Connection: self.changed.set() if stream is not None: stream.changed.set() - self.flush() - await self.drain() + if connection_credit: + self.h2.increment_flow_control_window(connection_credit) + if self.flush(): + await self.drain() except asyncio.CancelledError: self.fail(httpx.ReadError('H2 reader stopped')) raise @@ -202,7 +209,7 @@ class _Connection: self.changed.clear() self.check_error() stream.check_error() - count = min(len(data) - offset, self.h2.max_outbound_frame_size, + count = min(len(data) - offset, IO_BATCH_SIZE, self.h2.local_flow_control_window(stream.stream_id)) if count <= 0: try: @@ -210,8 +217,11 @@ class _Connection: except asyncio.TimeoutError as exc: raise httpx.WriteTimeout('H2 upload flow control timed out') from exc continue - self.h2.send_data(stream.stream_id, data[offset:offset + count]) - offset += count + batch_end = offset + count + while offset < batch_end: + frame_end = min(batch_end, offset + self.h2.max_outbound_frame_size) + self.h2.send_data(stream.stream_id, data[offset:frame_end]) + offset = frame_end self.flush() await self.drain() self.h2.end_stream(stream.stream_id) @@ -259,8 +269,12 @@ class _Connection: def consumed(self, stream, credit): if credit and not stream.ended and not stream.closed and self.error is None: + stream.consumed_credit += credit + if stream.chunks and stream.consumed_credit < IO_BATCH_SIZE: + return try: - self.h2.increment_flow_control_window(credit, stream.stream_id) + self.h2.increment_flow_control_window(stream.consumed_credit, stream.stream_id) + stream.consumed_credit = 0 self.flush() except StreamClosedError: pass diff --git a/tests/test_h2_transport.py b/tests/test_h2_transport.py index 96e616f..41742e2 100644 --- a/tests/test_h2_transport.py +++ b/tests/test_h2_transport.py @@ -1016,8 +1016,8 @@ class H2WireTest(unittest.IsolatedAsyncioTestCase): self.respond(request, b's' * (1024 * 1024)) response = await opening fast, fast_request = await self.post() - self.respond(fast_request, b'f' * (128 * 1024)) - self.assertEqual(await asyncio.wait_for(fast, .5), b'f' * (128 * 1024)) + self.respond(fast_request, b'f' * (4 * 1024 * 1024)) + self.assertEqual(await asyncio.wait_for(fast, 2), b'f' * (4 * 1024 * 1024)) stream = self.transport.connection.streams[request[1]] self.assertLessEqual(sum(len(data) for data, _ in stream.chunks), STREAM_RECEIVE_WINDOW) await response.aclose()