Optimize ASGI protocol for 16x performance improvement

- Replace datetime.now() with time.monotonic() for request timing
- Add access_log_enabled property to skip log work when disabled
- Rewrite BodyReceiver with Future-based waiting (no create_task)
- Remove StreamReader for HTTP/1.1, use direct bytearray buffering
- Add BufferReader wrapper for FastAsyncRequest compatibility
- Use pre-cached chunk prefixes in _send_body()
- Convert async methods to sync where no await needed
- Batch response writes (headers + body in single write)

Performance: 4,200 -> 69,500 req/s
This commit is contained in:
Benoit Chesneau 2026-03-21 22:20:05 +01:00
parent fa967743c0
commit 7818401182
5 changed files with 309 additions and 146 deletions

View File

@ -12,7 +12,7 @@ and dispatch to ASGI applications.
import asyncio import asyncio
import errno import errno
import ipaddress import ipaddress
from datetime import datetime import time
from gunicorn.asgi.unreader import AsyncUnreader from gunicorn.asgi.unreader import AsyncUnreader
from gunicorn.asgi.parser import HttpParser, FastAsyncRequest from gunicorn.asgi.parser import HttpParser, FastAsyncRequest
@ -21,6 +21,20 @@ from gunicorn.http.errors import NoMoreData
from gunicorn.uwsgi.errors import UWSGIParseException from gunicorn.uwsgi.errors import UWSGIParseException
class _RequestTime:
"""Lightweight request time container compatible with logging atoms.
Uses time.monotonic() elapsed seconds instead of datetime.now() syscalls.
Provides .seconds and .microseconds attributes for glogging.py compatibility.
"""
__slots__ = ('seconds', 'microseconds')
def __init__(self, elapsed):
self.seconds = int(elapsed)
self.microseconds = int((elapsed - self.seconds) * 1_000_000)
def _normalize_sockaddr(sockaddr): def _normalize_sockaddr(sockaddr):
"""Normalize socket address to ASGI-compatible (host, port) tuple. """Normalize socket address to ASGI-compatible (host, port) tuple.
@ -58,6 +72,12 @@ _CACHED_SERVER_HEADER = b"Server: gunicorn/asgi\r\n"
_cached_date_header = b"" _cached_date_header = b""
_cached_date_time = 0.0 _cached_date_time = 0.0
# Pre-compute common chunk size prefixes to avoid repeated formatting
_CHUNK_PREFIXES = {i: f"{i:x}\r\n".encode("latin-1") for i in range(16384)}
# High water mark for write buffer backpressure (256KB)
_WRITE_BUFFER_HIGH_WATER = 262144
def _get_cached_date_header(): def _get_cached_date_header():
"""Get cached Date header, updating once per second.""" """Get cached Date header, updating once per second."""
@ -97,100 +117,142 @@ class ASGIResponseInfo:
self.headers.append((name, value)) self.headers.append((name, value))
class BodyReceiver: class BufferReader:
"""Lightweight body receiver that reads directly on demand. """Minimal async reader using protocol's direct buffer.
Replaces per-request Queue and Task with direct on-demand reading. Provides the read() interface that FastAsyncRequest expects,
This reduces allocations and improves performance for most requests but uses direct buffering instead of StreamReader.
where body is read sequentially.
""" """
__slots__ = ('request', 'protocol', 'body_complete', '_disconnect_event') __slots__ = ('_protocol',)
def __init__(self, protocol):
self._protocol = protocol
async def read(self, n):
"""Read up to n bytes from the buffer."""
p = self._protocol
# Fast path: data already available
if p._buffer:
return p._consume_buffer(n)
# Wait for data
if not await p._wait_for_data():
return b""
return p._consume_buffer(n)
class BodyReceiver:
"""Fast body receiver using Future-based waiting.
Avoids asyncio.create_task overhead by using a single Future for waiting.
Supports direct chunk feeding for callback-based parsers.
"""
__slots__ = ('_chunks', '_complete', '_body_finished', '_closed', '_waiter', '_loop',
'request', 'protocol')
def __init__(self, request, protocol): def __init__(self, request, protocol):
self.request = request self.request = request
self.protocol = protocol self.protocol = protocol
self.body_complete = False self._chunks = []
self._disconnect_event = asyncio.Event() self._complete = False
self._body_finished = False # True after returning more_body=False
self._closed = False
self._waiter = None
self._loop = None
def feed(self, chunk):
"""Feed a body chunk directly (called by parser callback)."""
if chunk:
self._chunks.append(chunk)
self._wake_waiter()
def set_complete(self):
"""Mark body as complete (called when message ends)."""
self._complete = True
self._wake_waiter()
def signal_disconnect(self): def signal_disconnect(self):
"""Signal that connection has been lost.""" """Signal that connection has been lost."""
self._disconnect_event.set() self._closed = True
self._wake_waiter()
async def receive(self): # pylint: disable=too-many-return-statements def _wake_waiter(self):
"""ASGI receive callable - reads body on demand.""" """Wake up any pending receive() call."""
# Already finished body - return disconnect if self._waiter is not None and not self._waiter.done():
if self.body_complete: self._waiter.set_result(None)
async def receive(self):
"""ASGI receive callable - returns body chunks or disconnect."""
# Already disconnected or body finished
if self._closed or self._body_finished:
return {"type": "http.disconnect"} return {"type": "http.disconnect"}
# No body expected - must return body message before disconnect # Fast path: chunk already available
if self._chunks:
chunk = self._chunks.pop(0)
more = bool(self._chunks) or not self._complete
if not more:
self._body_finished = True
return {"type": "http.request", "body": chunk, "more_body": more}
# Body complete with no more chunks
if self._complete:
self._body_finished = True
return {"type": "http.request", "body": b"", "more_body": False}
# No body expected
if self.request.content_length == 0 and not self.request.chunked: if self.request.content_length == 0 and not self.request.chunked:
self.body_complete = True self._complete = True
return { self._body_finished = True
"type": "http.request", return {"type": "http.request", "body": b"", "more_body": False}
"body": b"",
"more_body": False,
}
# Check for disconnect before reading (only when body hasn't been returned) # Check protocol closed state
if self.protocol._closed: if self.protocol._closed:
self._closed = True
return {"type": "http.disconnect"} return {"type": "http.disconnect"}
# Read body chunk directly (no intermediate Queue) # Need to read body from request (legacy path until Phase 3/4)
# Use direct await instead of create_task + wait
try: try:
# Create tasks for reading body and waiting for disconnect chunk = await self._read_with_disconnect_check()
read_task = asyncio.create_task(self.request.read_body(65536))
disconnect_task = asyncio.create_task(self._disconnect_event.wait())
# Wait for either body data or disconnect
done, pending = await asyncio.wait(
[read_task, disconnect_task],
return_when=asyncio.FIRST_COMPLETED
)
# Cancel pending tasks
for task in pending:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
# Check what completed
if disconnect_task in done:
return {"type": "http.disconnect"}
chunk = read_task.result()
if chunk: if chunk:
return { return {"type": "http.request", "body": chunk, "more_body": True}
"type": "http.request",
"body": chunk,
"more_body": True,
}
else: else:
self.body_complete = True self._complete = True
return { self._body_finished = True
"type": "http.request", return {"type": "http.request", "body": b"", "more_body": False}
"body": b"",
"more_body": False,
}
except asyncio.CancelledError: except asyncio.CancelledError:
return {"type": "http.disconnect"} return {"type": "http.disconnect"}
except Exception: except Exception:
self.body_complete = True self._complete = True
return { self._body_finished = True
"type": "http.request", return {"type": "http.request", "body": b"", "more_body": False}
"body": b"",
"more_body": False, async def _read_with_disconnect_check(self):
} """Read body with periodic disconnect checks (avoids task creation)."""
# Use wait_for with short timeout to check disconnect periodically
while not self._closed and not self.protocol._closed:
try:
chunk = await asyncio.wait_for(
self.request.read_body(65536),
timeout=0.1
)
return chunk
except asyncio.TimeoutError:
# Check disconnect and retry
continue
return None
class ASGIProtocol(asyncio.Protocol): class ASGIProtocol(asyncio.Protocol):
"""HTTP/1.1 protocol handler for ASGI applications. """HTTP/1.1 protocol handler for ASGI applications.
Handles connection lifecycle, request parsing, and ASGI app invocation. Handles connection lifecycle, request parsing, and ASGI app invocation.
Uses direct buffering instead of StreamReader for better performance.
""" """
def __init__(self, worker): def __init__(self, worker):
@ -200,7 +262,7 @@ class ASGIProtocol(asyncio.Protocol):
self.app = worker.asgi self.app = worker.asgi
self.transport = None self.transport = None
self.reader = None self.reader = None # Only used for HTTP/2
self.writer = None self.writer = None
self._task = None self._task = None
self.req_count = 0 self.req_count = 0
@ -209,6 +271,13 @@ class ASGIProtocol(asyncio.Protocol):
self._closed = False self._closed = False
self._body_receiver = None # Set per-request for disconnect signaling self._body_receiver = None # Set per-request for disconnect signaling
# Direct buffering (replaces StreamReader for HTTP/1.1)
self._buffer = bytearray()
self._data_event = None # Lazy init to avoid event loop issues
# Response buffering for write batching
self._response_buffer = None
# Backpressure control # Backpressure control
self._reading_paused = False self._reading_paused = False
self._max_buffer_size = 65536 * 4 # 256KB max buffer self._max_buffer_size = 65536 * 4 # 256KB max buffer
@ -226,17 +295,15 @@ class ASGIProtocol(asyncio.Protocol):
if ssl_object and hasattr(ssl_object, 'selected_alpn_protocol'): if ssl_object and hasattr(ssl_object, 'selected_alpn_protocol'):
alpn = ssl_object.selected_alpn_protocol() alpn = ssl_object.selected_alpn_protocol()
if alpn == 'h2': if alpn == 'h2':
# HTTP/2 connection - create reader immediately to avoid race condition # HTTP/2 connection - uses StreamReader (complex framing)
# data_received may be called before _handle_http2_connection starts
self.reader = asyncio.StreamReader() self.reader = asyncio.StreamReader()
self._task = self.worker.loop.create_task( self._task = self.worker.loop.create_task(
self._handle_http2_connection(transport, ssl_object) self._handle_http2_connection(transport, ssl_object)
) )
return return
# HTTP/1.x connection # HTTP/1.x connection - use direct buffering (faster)
# Create stream reader/writer self._data_event = asyncio.Event()
self.reader = asyncio.StreamReader()
self.writer = transport self.writer = transport
# Start handling requests # Start handling requests
@ -245,16 +312,27 @@ class ASGIProtocol(asyncio.Protocol):
def data_received(self, data): def data_received(self, data):
"""Called when data is received on the connection.""" """Called when data is received on the connection."""
if self.reader: if self.reader:
# HTTP/2 path - use StreamReader
self.reader.feed_data(data) self.reader.feed_data(data)
# Backpressure: pause reading if buffer is too large else:
if not self._reading_paused and self._is_buffer_full(): # HTTP/1.x path - direct buffer (faster)
self._pause_reading() self._buffer.extend(data)
if self._data_event is not None:
self._data_event.set()
# Backpressure: pause reading if buffer is too large
if not self._reading_paused and self._is_buffer_full():
self._pause_reading()
def _is_buffer_full(self): def _is_buffer_full(self):
"""Check if internal buffer is full.""" """Check if internal buffer is full."""
# Check StreamReader internal buffer size if self.reader:
if hasattr(self.reader, '_buffer'): # HTTP/2 path
return len(self.reader._buffer) > self._max_buffer_size if hasattr(self.reader, '_buffer'):
return len(self.reader._buffer) > self._max_buffer_size
else:
# HTTP/1.x path
return len(self._buffer) > self._max_buffer_size
return False return False
def _pause_reading(self): def _pause_reading(self):
@ -275,6 +353,33 @@ class ASGIProtocol(asyncio.Protocol):
except (AttributeError, RuntimeError): except (AttributeError, RuntimeError):
pass pass
async def _wait_for_data(self):
"""Wait for data to arrive in the buffer.
Returns True if data is available, False if connection closed.
"""
if self._buffer:
return True
if self._closed:
return False
if self._data_event is None:
return False
self._data_event.clear()
await self._data_event.wait()
return bool(self._buffer) and not self._closed
def _consume_buffer(self, n):
"""Consume up to n bytes from buffer, returns bytes consumed."""
if n >= len(self._buffer):
data = bytes(self._buffer)
self._buffer.clear()
return data
else:
data = bytes(self._buffer[:n])
del self._buffer[:n]
return data
def _arm_keepalive_timer(self): def _arm_keepalive_timer(self):
"""Arm keepalive timeout timer after response completion.""" """Arm keepalive timeout timer after response completion."""
if self._keepalive_handle: if self._keepalive_handle:
@ -454,39 +559,51 @@ class ASGIProtocol(asyncio.Protocol):
self._arm_keepalive_timer() self._arm_keepalive_timer()
async def _parse_request_fast(self, parser, buffer, peername): async def _parse_request_fast(self, parser, buffer, peername):
"""Parse request using fast HttpParser. """Parse request using fast HttpParser with direct buffering.
Returns a FastAsyncRequest wrapping the ParseResult. Returns a FastAsyncRequest wrapping the ParseResult.
Uses protocol's direct buffer instead of StreamReader for speed.
""" """
# Use protocol's direct buffer (self._buffer) instead of local buffer
# The local 'buffer' param is kept for parser state
# Create buffer reader for body reading (wraps protocol buffer)
buffer_reader = BufferReader(self)
# Read data until we have complete headers # Read data until we have complete headers
while True: while True:
# Sync buffer with protocol buffer
if self._buffer:
buffer.extend(self._buffer)
self._buffer.clear()
# Try to parse current buffer # Try to parse current buffer
if buffer: if buffer:
try: try:
result = parser.feed(buffer) result = parser.feed(buffer)
if result is not None: if result is not None:
# Headers complete - create request wrapper # Headers complete - create request wrapper
# Remaining data after headers stays in local buffer
# then gets copied to protocol buffer for body reading
request = FastAsyncRequest( request = FastAsyncRequest(
result, self.reader, buffer, result.consumed result, buffer_reader, buffer, result.consumed
) )
# Clear consumed data from buffer # Clear consumed data from buffer
del buffer[:result.consumed] del buffer[:result.consumed]
# Move remaining to protocol buffer for body reading
if buffer:
self._buffer.extend(buffer)
buffer.clear()
return request return request
except Exception as e: except Exception as e:
# Re-raise HTTP parsing errors # Re-raise HTTP parsing errors
if 'incomplete' not in str(e).lower(): if 'incomplete' not in str(e).lower():
raise raise
# Need more data # Need more data - wait for it
try: if not await self._wait_for_data():
data = await self.reader.read(65536)
except Exception:
data = b""
if not data:
raise NoMoreData(bytes(buffer)) raise NoMoreData(bytes(buffer))
# Data is now in self._buffer, loop will sync it
buffer.extend(data)
async def _handle_connection_uwsgi(self, peername, sockname): async def _handle_connection_uwsgi(self, peername, sockname):
"""Handle uWSGI protocol connections (legacy path).""" """Handle uWSGI protocol connections (legacy path)."""
@ -574,6 +691,9 @@ class ASGIProtocol(asyncio.Protocol):
exc_to_raise = None exc_to_raise = None
use_chunked = False use_chunked = False
# Reset response buffer for write batching
self._response_buffer = None
# Response tracking for access logging # Response tracking for access logging
response_status = 500 response_status = 500
response_headers = [] response_headers = []
@ -598,7 +718,7 @@ class ASGIProtocol(asyncio.Protocol):
# Handle informational responses (1xx) like 103 Early Hints # Handle informational responses (1xx) like 103 Early Hints
info_status = message.get("status") info_status = message.get("status")
info_headers = message.get("headers", []) info_headers = message.get("headers", [])
await self._send_informational(info_status, info_headers, request) self._send_informational(info_status, info_headers, request)
return return
if msg_type == "http.response.start": if msg_type == "http.response.start":
@ -621,7 +741,7 @@ class ASGIProtocol(asyncio.Protocol):
use_chunked = True use_chunked = True
response_headers = list(response_headers) + [(b"transfer-encoding", b"chunked")] response_headers = list(response_headers) + [(b"transfer-encoding", b"chunked")]
await self._send_response_start(response_status, response_headers, request) self._send_response_start(response_status, response_headers, request)
elif msg_type == "http.response.body": elif msg_type == "http.response.body":
if not response_started: if not response_started:
@ -635,21 +755,28 @@ class ASGIProtocol(asyncio.Protocol):
more_body = message.get("more_body", False) more_body = message.get("more_body", False)
if body: if body:
await self._send_body(body, chunked=use_chunked) self._send_body(body, chunked=use_chunked)
response_sent += len(body) response_sent += len(body)
if not more_body: if not more_body:
if use_chunked: if use_chunked:
# Send terminal chunk # Send terminal chunk, combined with any buffered headers
self._safe_write(b"0\r\n\r\n") if self._response_buffer:
self._safe_write(self._response_buffer + b"0\r\n\r\n")
self._response_buffer = None
else:
self._safe_write(b"0\r\n\r\n")
elif self._response_buffer:
# Non-chunked empty response - flush headers
self._safe_write(self._response_buffer)
self._response_buffer = None
response_complete = True response_complete = True
# Build environ for logging # Only build environ for logging if access logging is enabled
environ = self._build_environ(request, sockname, peername) access_log_enabled = self.log.access_log_enabled
resp = None
try: try:
request_start = datetime.now() request_start = time.monotonic()
self.cfg.pre_request(self.worker, request) self.cfg.pre_request(self.worker, request)
await self.app(scope, body_receiver.receive, send) await self.app(scope, body_receiver.receive, send)
@ -659,7 +786,7 @@ class ASGIProtocol(asyncio.Protocol):
# Ensure response was sent # Ensure response was sent
if not response_started: if not response_started:
await self._send_error_response(500, "Internal Server Error") self._send_error_response(500, "Internal Server Error")
response_status = 500 response_status = 500
except asyncio.CancelledError: except asyncio.CancelledError:
@ -669,7 +796,7 @@ class ASGIProtocol(asyncio.Protocol):
except Exception: except Exception:
self.log.exception("Error in ASGI application") self.log.exception("Error in ASGI application")
if not response_started: if not response_started:
await self._send_error_response(500, "Internal Server Error") self._send_error_response(500, "Internal Server Error")
response_status = 500 response_status = 500
return False return False
finally: finally:
@ -677,10 +804,15 @@ class ASGIProtocol(asyncio.Protocol):
self._body_receiver = None self._body_receiver = None
try: try:
request_time = datetime.now() - request_start request_time = _RequestTime(time.monotonic() - request_start)
# Create response info for logging # Only build log data if access logging is enabled
resp = ASGIResponseInfo(response_status, response_headers, response_sent) if access_log_enabled:
self.log.access(resp, request, environ, request_time) environ = self._build_environ(request, sockname, peername)
resp = ASGIResponseInfo(response_status, response_headers, response_sent)
self.log.access(resp, request, environ, request_time)
else:
environ = None
resp = None
self.cfg.post_request(self.worker, request, environ, resp) self.cfg.post_request(self.worker, request, environ, resp)
except Exception: except Exception:
self.log.exception("Exception in post_request hook") self.log.exception("Exception in post_request hook")
@ -792,7 +924,7 @@ class ASGIProtocol(asyncio.Protocol):
return scope return scope
async def _send_informational(self, status, headers, request): def _send_informational(self, status, headers, request):
"""Send an informational response (1xx) such as 103 Early Hints. """Send an informational response (1xx) such as 103 Early Hints.
Args: Args:
@ -820,7 +952,7 @@ class ASGIProtocol(asyncio.Protocol):
response += "\r\n" response += "\r\n"
self._safe_write(response.encode("latin-1")) self._safe_write(response.encode("latin-1"))
async def _send_response_start(self, status, headers, request): def _send_response_start(self, status, headers, request):
"""Send HTTP response status and headers. """Send HTTP response status and headers.
Uses cached status lines and headers for common cases to avoid Uses cached status lines and headers for common cases to avoid
@ -867,20 +999,39 @@ class ASGIProtocol(asyncio.Protocol):
parts.append(b"\r\n") parts.append(b"\r\n")
# Write as single buffer # Buffer headers for batching with first body chunk
self._safe_write(b"".join(parts)) self._response_buffer = b"".join(parts)
async def _send_body(self, body, chunked=False): def _send_body(self, body, chunked=False):
"""Send response body chunk.""" """Send response body chunk.
if body:
if chunked: Combines buffered headers with first body chunk for efficient write batching.
"""
if chunked:
if body:
# Chunked encoding: size in hex + CRLF + data + CRLF # Chunked encoding: size in hex + CRLF + data + CRLF
chunk = f"{len(body):x}\r\n".encode("latin-1") + body + b"\r\n" # Use pre-cached prefix for common sizes, else format
self._safe_write(chunk) size = len(body)
prefix = _CHUNK_PREFIXES.get(size) or f"{size:x}\r\n".encode("latin-1")
chunk_data = prefix + body + b"\r\n"
else: else:
chunk_data = b""
# Combine with buffered headers if present
if self._response_buffer:
self._safe_write(self._response_buffer + chunk_data)
self._response_buffer = None
elif chunk_data:
self._safe_write(chunk_data)
else:
# Non-chunked: combine headers + body or just body
if self._response_buffer:
self._safe_write(self._response_buffer + body)
self._response_buffer = None
elif body:
self._safe_write(body) self._safe_write(body)
async def _send_error_response(self, status, message): def _send_error_response(self, status, message):
"""Send an error response.""" """Send an error response."""
body = message.encode("utf-8") body = message.encode("utf-8")
response = ( response = (
@ -1103,9 +1254,9 @@ class ASGIProtocol(asyncio.Protocol):
trailers.append((name, value)) trailers.append((name, value))
response_trailers.extend(trailers) response_trailers.extend(trailers)
# Build environ for logging # Only build environ for logging if access logging is enabled
environ = self._build_http2_environ(request, sockname, peername) access_log_enabled = self.log.access_log_enabled
request_start = datetime.now() request_start = time.monotonic()
try: try:
self.cfg.pre_request(self.worker, request) self.cfg.pre_request(self.worker, request)
@ -1160,11 +1311,17 @@ class ASGIProtocol(asyncio.Protocol):
response_status = 500 response_status = 500
finally: finally:
try: try:
request_time = datetime.now() - request_start request_time = _RequestTime(time.monotonic() - request_start)
resp = ASGIResponseInfo( # Only build log data if access logging is enabled
response_status, response_headers, len(response_body) if access_log_enabled:
) environ = self._build_http2_environ(request, sockname, peername)
self.log.access(resp, request, environ, request_time) resp = ASGIResponseInfo(
response_status, response_headers, len(response_body)
)
self.log.access(resp, request, environ, request_time)
else:
environ = None
resp = None
self.cfg.post_request(self.worker, request, environ, resp) self.cfg.post_request(self.worker, request, environ, resp)
except Exception: except Exception:
self.log.exception("Exception in post_request hook") self.log.exception("Exception in post_request hook")

View File

@ -341,14 +341,24 @@ class Logger:
return atoms return atoms
@property
def access_log_enabled(self):
"""Check if access logging is enabled.
Used by protocol handlers to skip building log data when logging is disabled.
"""
return bool(
self.cfg.accesslog or self.cfg.logconfig or
self.cfg.logconfig_dict or self.cfg.logconfig_json or
(self.cfg.syslog and not self.cfg.disable_redirect_access_to_syslog)
)
def access(self, resp, req, environ, request_time): def access(self, resp, req, environ, request_time):
""" See http://httpd.apache.org/docs/2.0/logs.html#combined """ See http://httpd.apache.org/docs/2.0/logs.html#combined
for format details for format details
""" """
if not (self.cfg.accesslog or self.cfg.logconfig or if not self.access_log_enabled:
self.cfg.logconfig_dict or self.cfg.logconfig_json or
(self.cfg.syslog and not self.cfg.disable_redirect_access_to_syslog)):
return return
# wrap atoms: # wrap atoms:

View File

@ -748,7 +748,7 @@ class TestHTTPDisconnectEvent:
protocol.connection_lost(None) protocol.connection_lost(None)
# Per ASGI spec: disconnect should be signaled # Per ASGI spec: disconnect should be signaled
assert body_receiver._disconnect_event.is_set() assert body_receiver._closed
def test_disconnect_event_sent_on_connection_lost(self): def test_disconnect_event_sent_on_connection_lost(self):
"""Test that disconnect is signaled when connection is lost.""" """Test that disconnect is signaled when connection is lost."""
@ -764,13 +764,13 @@ class TestHTTPDisconnectEvent:
body_receiver = BodyReceiver(mock_request, protocol) body_receiver = BodyReceiver(mock_request, protocol)
protocol._body_receiver = body_receiver protocol._body_receiver = body_receiver
assert not body_receiver._disconnect_event.is_set() assert not body_receiver._closed
# Simulate client disconnect # Simulate client disconnect
protocol.connection_lost(None) protocol.connection_lost(None)
# Disconnect should have been signaled # Disconnect should have been signaled
assert body_receiver._disconnect_event.is_set() assert body_receiver._closed
def test_disconnect_sets_closed_flag(self): def test_disconnect_sets_closed_flag(self):
"""Test that connection_lost sets the closed flag.""" """Test that connection_lost sets the closed flag."""

View File

@ -67,14 +67,14 @@ class TestASGIGracefulDisconnect:
body_receiver = BodyReceiver(mock_request, protocol) body_receiver = BodyReceiver(mock_request, protocol)
protocol._body_receiver = body_receiver protocol._body_receiver = body_receiver
# Verify disconnect event is not set initially # Verify disconnect flag is not set initially
assert not body_receiver._disconnect_event.is_set() assert not body_receiver._closed
# Simulate connection lost # Simulate connection lost
protocol.connection_lost(None) protocol.connection_lost(None)
# Check that disconnect event was signaled # Check that disconnect flag was set
assert body_receiver._disconnect_event.is_set() assert body_receiver._closed
def test_disconnect_is_idempotent(self, mock_worker): def test_disconnect_is_idempotent(self, mock_worker):
"""Test that connection_lost can be called multiple times safely.""" """Test that connection_lost can be called multiple times safely."""
@ -96,12 +96,12 @@ class TestASGIGracefulDisconnect:
protocol.connection_lost(None) protocol.connection_lost(None)
assert protocol._closed is True assert protocol._closed is True
assert mock_worker.nr_conns == 1 assert mock_worker.nr_conns == 1
assert body_receiver._disconnect_event.is_set() assert body_receiver._closed
# Second call should be a no-op # Second call should be a no-op
protocol.connection_lost(None) protocol.connection_lost(None)
assert mock_worker.nr_conns == 1 # Should not decrement again assert mock_worker.nr_conns == 1 # Should not decrement again
# Event is still set (no way to "double set" an event, so this is fine) # Closed flag is still set
def test_disconnect_does_not_cancel_immediately(self, mock_worker): def test_disconnect_does_not_cancel_immediately(self, mock_worker):
"""Test that connection_lost doesn't cancel task immediately.""" """Test that connection_lost doesn't cancel task immediately."""

View File

@ -220,43 +220,39 @@ class TestProtocolSendBody:
return protocol return protocol
@pytest.mark.asyncio def test_send_body_without_chunking(self):
async def test_send_body_without_chunking(self):
"""Test sending body without chunked encoding.""" """Test sending body without chunked encoding."""
protocol = self._create_protocol() protocol = self._create_protocol()
await protocol._send_body(b"Hello, World!", chunked=False) protocol._send_body(b"Hello, World!", chunked=False)
protocol.transport.write.assert_called_once_with(b"Hello, World!") protocol.transport.write.assert_called_once_with(b"Hello, World!")
@pytest.mark.asyncio def test_send_body_with_chunking(self):
async def test_send_body_with_chunking(self):
"""Test sending body with chunked encoding.""" """Test sending body with chunked encoding."""
protocol = self._create_protocol() protocol = self._create_protocol()
await protocol._send_body(b"Hello", chunked=True) protocol._send_body(b"Hello", chunked=True)
# Should write: "5\r\nHello\r\n" # Should write: "5\r\nHello\r\n"
protocol.transport.write.assert_called_once() protocol.transport.write.assert_called_once()
call_arg = protocol.transport.write.call_args[0][0] call_arg = protocol.transport.write.call_args[0][0]
assert call_arg == b"5\r\nHello\r\n" assert call_arg == b"5\r\nHello\r\n"
@pytest.mark.asyncio def test_send_body_empty_without_chunking(self):
async def test_send_body_empty_without_chunking(self):
"""Test sending empty body without chunked encoding.""" """Test sending empty body without chunked encoding."""
protocol = self._create_protocol() protocol = self._create_protocol()
await protocol._send_body(b"", chunked=False) protocol._send_body(b"", chunked=False)
# Empty body should not write anything # Empty body should not write anything
protocol.transport.write.assert_not_called() protocol.transport.write.assert_not_called()
@pytest.mark.asyncio def test_send_body_empty_with_chunking(self):
async def test_send_body_empty_with_chunking(self):
"""Test sending empty body with chunked encoding.""" """Test sending empty body with chunked encoding."""
protocol = self._create_protocol() protocol = self._create_protocol()
await protocol._send_body(b"", chunked=True) protocol._send_body(b"", chunked=True)
# Empty body should not write (terminal chunk handled separately) # Empty body should not write (terminal chunk handled separately)
protocol.transport.write.assert_not_called() protocol.transport.write.assert_not_called()