diff --git a/CHANGES.rst b/CHANGES.rst index 8b06d08c674..e27829785ee 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -472,8 +472,7 @@ Features -- Added :attr:`~aiohttp.ClientResponse.output_size` and - :attr:`~aiohttp.ClientResponse.upload_complete` -- by :user:`Dreamsorcerer`. +- Added ``ClientResponse.output_size`` and ``ClientResponse.upload_complete`` -- by :user:`Dreamsorcerer`. *Related issues and pull requests on GitHub:* diff --git a/CHANGES/13579.bugfix.rst b/CHANGES/13579.bugfix.rst new file mode 100644 index 00000000000..f4b41121d6e --- /dev/null +++ b/CHANGES/13579.bugfix.rst @@ -0,0 +1,4 @@ +Fixed a connection being eligible for reuse after its request was cancelled +or failed while waiting for a ``100 Continue`` response or finalizing the +body; the request headers were already sent, so reusing the connection +corrupted the next request on it -- by :user:`Dreamsorcerer`. diff --git a/CHANGES/13579.deprecation.rst b/CHANGES/13579.deprecation.rst new file mode 100644 index 00000000000..1b5ebddb220 --- /dev/null +++ b/CHANGES/13579.deprecation.rst @@ -0,0 +1,2 @@ +Deprecated ``ClientResponse.output_size`` and ``ClientResponse.upload_complete``; +use ``aiohttp.UploadTracker`` instead -- by :user:`Dreamsorcerer`. diff --git a/CHANGES/13579.feature.rst b/CHANGES/13579.feature.rst new file mode 100644 index 00000000000..74c1463c3e8 --- /dev/null +++ b/CHANGES/13579.feature.rst @@ -0,0 +1 @@ +Added :class:`aiohttp.UploadTracker` for observing a client request's upload progress -- by :user:`Dreamsorcerer`. diff --git a/THREAT_MODEL.md b/THREAT_MODEL.md index 41d0ef51c7b..d16088274ca 100644 --- a/THREAT_MODEL.md +++ b/THREAT_MODEL.md @@ -31,7 +31,7 @@ Key public APIs (non-exhaustive): | Surface | Entry points | | --- | --- | | Server | `aiohttp.web.Application`, `web.RouteTableDef`, `web.run_app`, `web.AppRunner`, `web.WebSocketResponse`, `web.FileResponse` | -| Client | `aiohttp.ClientSession`, `aiohttp.TCPConnector`, `aiohttp.ClientResponse`, `aiohttp.WSMessage`, `aiohttp.BasicAuth` | +| Client | `aiohttp.ClientSession`, `aiohttp.TCPConnector`, `aiohttp.ClientResponse`, `aiohttp.UploadTracker`, `aiohttp.WSMessage`, `aiohttp.BasicAuth` | | Shared | `aiohttp.MultipartReader`/`MultipartWriter`, `aiohttp.CookieJar`, `aiohttp.TraceConfig`, `aiohttp.resolver.AsyncResolver` | --- diff --git a/aiohttp/__init__.py b/aiohttp/__init__.py index 64bc3ced0a1..34ff7b71c86 100644 --- a/aiohttp/__init__.py +++ b/aiohttp/__init__.py @@ -43,6 +43,8 @@ TCPConnector, TooManyRedirects, UnixConnector, + UploadAbortedError, + UploadTracker, WSMessageTypeError, WSServerHandshakeError, request, @@ -156,6 +158,8 @@ "TCPConnector", "TooManyRedirects", "UnixConnector", + "UploadAbortedError", + "UploadTracker", "NamedPipeConnector", "WSServerHandshakeError", "request", diff --git a/aiohttp/abc.py b/aiohttp/abc.py index c90f9f459ee..a21c56d5329 100644 --- a/aiohttp/abc.py +++ b/aiohttp/abc.py @@ -204,6 +204,12 @@ class AbstractStreamWriter(ABC): buffer_size: int = 0 output_size: int = 0 length: int | None = 0 + # Called with each accepted body chunk's byte count (before any + # transport-level transformation such as compression or chunked + # framing). Assigned by the client request machinery for upload + # progress tracking; write()/write_eof() implementations should + # invoke it for every body chunk they accept. + on_body_write: Callable[[int], None] | None = None @abstractmethod async def write( diff --git a/aiohttp/client.py b/aiohttp/client.py index f91b81e5712..672a12449c0 100644 --- a/aiohttp/client.py +++ b/aiohttp/client.py @@ -65,6 +65,7 @@ ServerTimeoutError, SocketTimeoutError, TooManyRedirects, + UploadAbortedError, WSMessageTypeError, WSServerHandshakeError, ) @@ -77,6 +78,7 @@ Fingerprint, RequestInfo, ResponseParams, + UploadTracker, ) from .client_ws import ( DEFAULT_WS_CLIENT_TIMEOUT, @@ -137,12 +139,14 @@ "ServerTimeoutError", "SocketTimeoutError", "TooManyRedirects", + "UploadAbortedError", "WSServerHandshakeError", # client_reqrep "ClientRequest", "ClientResponse", "Fingerprint", "RequestInfo", + "UploadTracker", # connector "BaseConnector", "TCPConnector", @@ -194,6 +198,7 @@ class _RequestOptions(TypedDict, total=False): max_field_size: int | None max_headers: int | None middlewares: Sequence[ClientMiddlewareType] | None + upload_tracker: UploadTracker | None class _WSConnectOptions(TypedDict, total=False): @@ -234,6 +239,8 @@ class _WSConnectOptions(TypedDict, total=False): async def _connect_and_send_request(req: ClientRequest) -> ClientResponse: connector = req._session._connector assert connector is not None + if (tracker := req._upload_tracker) is not None: + req._upload_gen = tracker._attempt_started() try: conn = await connector.connect(req, traces=req._traces, timeout=req._timeout) except asyncio.TimeoutError as exc: @@ -494,112 +501,147 @@ async def _request( max_field_size: int | None = None, max_headers: int | None = None, middlewares: Sequence[ClientMiddlewareType] | None = None, + upload_tracker: UploadTracker | None = None, ) -> ClientResponse: # NOTE: timeout clamps existing connect and read timeouts. We cannot # set the default to None because we need to detect if the user wants # to use the existing timeouts by setting timeout to None. - if self.closed: - raise RuntimeError("Session is closed") + # Bound outside the settle-guaranteeing try block: a rebind error + # must not settle a tracker owned by another request. + if upload_tracker is not None: + upload_tracker._bind() - method = method.upper() + tm: TimeoutHandle | None = None + handle: asyncio.TimerHandle | None = None + traces: list[Trace] = [] + traces_started = 0 + trace_url: URL | None = None + trace_headers: "CIMultiDict[str] | None" = None - if ssl is sentinel: - ssl = self._default_ssl - if not isinstance(ssl, SSL_ALLOWED_TYPES): - raise TypeError( - "ssl should be SSLContext, Fingerprint, or bool, " - f"got {ssl!r} instead." - ) - - if data is not None and json is not None: - raise ValueError( - "data and json parameters can not be used at the same time" - ) - elif json is not None: - if self._json_serialize_bytes is not None: - data = payload.JsonBytesPayload(json, dumps=self._json_serialize_bytes) - else: - data = payload.JsonPayload(json, dumps=self._json_serialize) + try: + if self.closed: + raise RuntimeError("Session is closed") - redirects = 0 - history: list[ClientResponse] = [] - version = self._version - params = params or {} + method = method.upper() - # Merge with default headers and transform to CIMultiDict - headers = self._prepare_headers(headers) + if ssl is sentinel: + ssl = self._default_ssl + if not isinstance(ssl, SSL_ALLOWED_TYPES): + raise TypeError( + "ssl should be SSLContext, Fingerprint, or bool, " + f"got {ssl!r} instead." + ) - try: - url = self._build_url(str_or_url) - except ValueError as e: - raise InvalidUrlClientError(str_or_url) from e + if data is not None and json is not None: + raise ValueError( + "data and json parameters can not be used at the same time" + ) + elif json is not None: + if self._json_serialize_bytes is not None: + data = payload.JsonBytesPayload( + json, dumps=self._json_serialize_bytes + ) + else: + data = payload.JsonPayload(json, dumps=self._json_serialize) - assert self._connector is not None - if url.scheme not in self._connector.allowed_protocol_schema_set: - raise NonHttpUrlClientError(url) + real_timeout = ( + self._timeout if timeout is sentinel or timeout is None else timeout + ) + # timeout is cumulative for all request operations + # (request, redirects, responses, data consuming) + tm = TimeoutHandle( + self._loop, + real_timeout.total, + ceil_threshold=real_timeout.ceil_threshold, + ) - skip_headers: Iterable[istr] | None - if skip_auto_headers is not None: - skip_headers = { - istr(i) for i in skip_auto_headers - } | self._skip_auto_headers - elif self._skip_auto_headers: - skip_headers = self._skip_auto_headers - else: - skip_headers = None + redirects = 0 + history: list[ClientResponse] = [] + version = self._version + params = params or {} - if proxy is None: - proxy = self._default_proxy + # Merge with default headers and transform to CIMultiDict + headers = self._prepare_headers(headers) - resolved_proxy_headers: CIMultiDict[str] | None - if proxy is None: - resolved_proxy_headers = None - else: - resolved_proxy_headers = self._prepare_headers(proxy_headers) try: - proxy = URL(proxy) + url = self._build_url(str_or_url) except ValueError as e: - raise InvalidURL(proxy) from e + raise InvalidUrlClientError(str_or_url) from e + + assert self._connector is not None + if url.scheme not in self._connector.allowed_protocol_schema_set: + raise NonHttpUrlClientError(url) + + skip_headers: Iterable[istr] | None + if skip_auto_headers is not None: + skip_headers = { + istr(i) for i in skip_auto_headers + } | self._skip_auto_headers + elif self._skip_auto_headers: + skip_headers = self._skip_auto_headers + else: + skip_headers = None - if timeout is sentinel or timeout is None: - real_timeout: ClientTimeout = self._timeout - else: - real_timeout = timeout - # timeout is cumulative for all request operations - # (request, redirects, responses, data consuming) - tm = TimeoutHandle( - self._loop, real_timeout.total, ceil_threshold=real_timeout.ceil_threshold - ) - handle = tm.start() + if proxy is None: + proxy = self._default_proxy - if read_bufsize is None: - read_bufsize = self._read_bufsize + resolved_proxy_headers: CIMultiDict[str] | None + if proxy is None: + resolved_proxy_headers = None + else: + resolved_proxy_headers = self._prepare_headers(proxy_headers) + try: + proxy = URL(proxy) + except ValueError as e: + raise InvalidURL(proxy) from e - if auto_decompress is None: - auto_decompress = self._auto_decompress + handle = tm.start() - if max_line_size is None: - max_line_size = self._max_line_size + if read_bufsize is None: + read_bufsize = self._read_bufsize - if max_field_size is None: - max_field_size = self._max_field_size + if auto_decompress is None: + auto_decompress = self._auto_decompress - if max_headers is None: - max_headers = self._max_headers + if max_line_size is None: + max_line_size = self._max_line_size - traces = [ - Trace( - self, - trace_config, - trace_config.trace_config_ctx(trace_request_ctx=trace_request_ctx), - ) - for trace_config in self._trace_configs - ] + if max_field_size is None: + max_field_size = self._max_field_size - for trace in traces: - await trace.send_request_start(method, url.update_query(params), headers) + if max_headers is None: + max_headers = self._max_headers + traces = [ + Trace( + self, + trace_config, + trace_config.trace_config_ctx(trace_request_ctx=trace_request_ctx), + ) + for trace_config in self._trace_configs + ] + + trace_url = url.update_query(params) + trace_headers = headers + for trace in traces: + await trace.send_request_start(method, trace_url, trace_headers) + traces_started += 1 + except BaseException as e: + if tm is not None: + tm.close() + if handle is not None: + handle.cancel() + if upload_tracker is not None: + upload_tracker._finalize() + # Traces that saw send_request_start must also see a terminal + # event; traces whose start never ran get neither. + for trace in traces[:traces_started]: + assert trace_url is not None and trace_headers is not None + await trace.send_request_exception(method, trace_url, trace_headers, e) + raise + + assert tm is not None timer = tm.timer() req: ClientRequest | None = None try: @@ -709,6 +751,7 @@ async def _request( traces=traces, trust_env=self.trust_env, ) + req._upload_tracker = upload_tracker # Apply middleware (if any) - per-request middleware overrides session middleware effective_middlewares = ( @@ -883,6 +926,8 @@ async def _request( await trace.send_request_end( method, url.update_query(params), headers, resp ) + if upload_tracker is not None: + upload_tracker._finalize() return resp except BaseException as e: @@ -890,7 +935,9 @@ async def _request( tm.close() if handle: handle.cancel() - handle = None + + if upload_tracker is not None: + upload_tracker._finalize() if req is not None and req._body is not None: await req._body.close() diff --git a/aiohttp/client_exceptions.py b/aiohttp/client_exceptions.py index 826533af81a..dd41c2315de 100644 --- a/aiohttp/client_exceptions.py +++ b/aiohttp/client_exceptions.py @@ -41,6 +41,7 @@ "WSServerHandshakeError", "ContentTypeError", "ClientPayloadError", + "UploadAbortedError", "InvalidURL", "InvalidUrlClientError", "RedirectClientError", @@ -261,6 +262,10 @@ class ClientPayloadError(ClientError): """Response payload error.""" +class UploadAbortedError(ClientError): + """The request body was never fully sent.""" + + class InvalidURL(ClientError, ValueError): """Invalid URL. diff --git a/aiohttp/client_reqrep.py b/aiohttp/client_reqrep.py index 516a4cdb17f..f617bfb5105 100644 --- a/aiohttp/client_reqrep.py +++ b/aiohttp/client_reqrep.py @@ -9,6 +9,7 @@ import warnings from asyncio.base_events import BaseEventLoop from collections.abc import Callable, Iterable, Sequence +from enum import Enum, auto from hashlib import md5, sha1, sha256 from http.cookies import BaseCookie, SimpleCookie from types import MappingProxyType, TracebackType @@ -32,6 +33,7 @@ ContentTypeError, InvalidURL, ServerFingerprintMismatch, + UploadAbortedError, ) from .compression_utils import HAS_BROTLI, HAS_ZSTD from .formdata import FormData @@ -245,6 +247,114 @@ class ResponseParams(TypedDict): max_headers: int +class _UploadState(Enum): + PENDING = auto() + ACTIVE = auto() + FINISHED = auto() + FAILED = auto() + CANCELLED = auto() + + +class UploadTracker: + """Tracks upload progress of a single client request. + + Pass a fresh instance via the request's ``upload_tracker`` argument and + read the plain attributes (``bytes_written``, ``attempts``, + ``upload_complete``) while the request runs. + + Must be created inside a running event loop. A tracker observes exactly + one request: passing it to a second request raises :exc:`RuntimeError`. + """ + + def __init__(self) -> None: + self.bytes_written = 0 + self.attempts = 0 + self.upload_complete: asyncio.Future[None] = ( + asyncio.get_running_loop().create_future() + ) + self._bound = False + # State of the latest upload attempt. + self._state = _UploadState.PENDING + self._exc: BaseException | None = None + # Set once the request reaches a terminal point (returned or raised); + # no further attempts can start after that. + self._final = False + + def _bind(self) -> None: + if self._bound: + raise RuntimeError("UploadTracker is already bound to a request") + self._bound = True + + def _attempt_started(self) -> int: + """A new upload attempt is dispatched; returns its generation token. + + Called before connecting, so a resend failing early settles as + aborted instead of reporting the superseded attempt's outcome. + """ + self.attempts += 1 + self.bytes_written = 0 + self._state = _UploadState.PENDING + self._exc = None + return self.attempts + + def _attempt_writing(self, gen: int) -> None: + """The attempt's body write has begun; settling now defers to it.""" + if gen == self.attempts: + self._state = _UploadState.ACTIVE + + def _add_bytes(self, gen: int, size: int) -> None: + # A stale writer of a superseded attempt (e.g. still being torn + # down while a redirect resends the body) must not corrupt the counter. + if gen == self.attempts: + self.bytes_written += size + + def _attempt_finished(self, gen: int) -> None: + if gen == self.attempts: + self._state = _UploadState.FINISHED + if self._final: + self._settle() + + def _attempt_failed(self, gen: int, exc: BaseException) -> None: + if gen == self.attempts: + self._state = _UploadState.FAILED + self._exc = exc + if self._final: + self._settle() + + def _attempt_cancelled(self, gen: int) -> None: + if gen == self.attempts: + self._state = _UploadState.CANCELLED + if self._final: + self._settle() + + def _finalize(self) -> None: + """Mark the request terminal: no further attempts will start. + + Settles upload_complete unless the final attempt is still writing, + in which case the attempt's own terminal event settles it. + """ + self._final = True + if self._state is not _UploadState.ACTIVE: + self._settle() + + def _settle(self) -> None: + fut = self.upload_complete + if fut.done(): + return + if self._state is _UploadState.FINISHED: + fut.set_result(None) + return + if self._state is _UploadState.FAILED: + assert self._exc is not None + fut.set_exception(self._exc) + else: + # PENDING (the body was never sent) or CANCELLED (the attempt + # was cut short). + fut.set_exception(UploadAbortedError("The request body was not fully sent")) + # Avoid 'exception was never retrieved' when doing a .done() poll. + fut.exception() + + class ClientResponse(HeadersMixin): # Some of these attributes are None when created, # but will be set by the start() method. @@ -363,7 +473,17 @@ def _writer(self, writer: asyncio.Task[None] | None) -> None: @property def output_size(self) -> int: - """Number of bytes sent for this request.""" + """Number of bytes sent for this request. + + .. deprecated:: 3.14.4 + Use :class:`UploadTracker` instead. + """ + warnings.warn( + "ClientResponse.output_size is deprecated, " + "use aiohttp.UploadTracker instead", + DeprecationWarning, + stacklevel=2, + ) if self._stream_writer is not None: return self._stream_writer.output_size return self._output_size @@ -372,8 +492,15 @@ def output_size(self) -> int: def upload_complete(self) -> "asyncio.Future[None]": """Future set when the request body has been fully sent. - Already done when the request had no body or was written eagerly. + .. deprecated:: 3.14.4 + Use :class:`UploadTracker` instead. """ + warnings.warn( + "ClientResponse.upload_complete is deprecated, " + "use aiohttp.UploadTracker instead", + DeprecationWarning, + stacklevel=2, + ) if self._upload_complete is None: self._upload_complete = self._loop.create_future() if self._stream_writer is None: # upload already finished @@ -808,6 +935,9 @@ class ClientRequestBase: _writer_task: asyncio.Task[None] | None = None # async task for streaming data + _upload_tracker: UploadTracker | None = None + _upload_gen = 0 + _skip_auto_headers: "CIMultiDict[None] | None" = None # N.B. @@ -1010,6 +1140,10 @@ async def _send(self, conn: "Connection") -> ClientResponse: protocol.start_timeout() writer.set_eof() task = None + if (tracker := self._upload_tracker) is not None: + # The request went out without a body-writer task: the + # dispatched attempt is trivially complete with zero bytes. + tracker._attempt_finished(self._upload_gen) self._response = self._create_response(task, stream_writer=writer) return self._response @@ -1481,6 +1615,7 @@ async def _write_bytes( - Content length constraints for chunked encoding - Error handling for network issues, cancellation, and other exceptions - Signaling EOF and timeout management + - Upload progress reporting to the request's UploadTracker Raises: ClientOSError: When there's an OS-level error writing the body @@ -1488,48 +1623,69 @@ async def _write_bytes( asyncio.CancelledError: When the operation is cancelled """ - # 100 response - if self._continue is not None: - # Force headers to be sent before waiting for 100-continue - writer.send_headers() - await writer.drain() - await self._continue - - protocol = conn.protocol - assert protocol is not None + tracker = self._upload_tracker + gen = self._upload_gen + if tracker is not None: + writer.on_body_write = functools.partial(tracker._add_bytes, gen) + tracker._attempt_writing(gen) try: - await self._body.write_with_length(writer, content_length) - except OSError as underlying_exc: - reraised_exc = underlying_exc + # 100 response + if self._continue is not None: + # Force headers to be sent before waiting for 100-continue + writer.send_headers() + await writer.drain() + await self._continue + + protocol = conn.protocol + assert protocol is not None + try: + await self._body.write_with_length(writer, content_length) + except OSError as underlying_exc: + reraised_exc = underlying_exc - # Distinguish between timeout and other OS errors for better error reporting - exc_is_not_timeout = underlying_exc.errno is not None or not isinstance( - underlying_exc, asyncio.TimeoutError - ) - if exc_is_not_timeout: - reraised_exc = ClientOSError( - underlying_exc.errno, - f"Can not write request body for {self.url !s}", + # Distinguish between timeout and other OS errors for better error reporting + exc_is_not_timeout = underlying_exc.errno is not None or not isinstance( + underlying_exc, asyncio.TimeoutError ) - - set_exception(protocol, reraised_exc, underlying_exc) + if exc_is_not_timeout: + reraised_exc = ClientOSError( + underlying_exc.errno, + f"Can not write request body for {self.url !s}", + ) + + set_exception(protocol, reraised_exc, underlying_exc) + if tracker is not None: + tracker._attempt_failed(gen, reraised_exc) + except Exception as underlying_exc: + wrapped_exc = ClientConnectionError( + "Failed to send bytes into the underlying connection " + f"{conn !s}: {underlying_exc!r}", + ) + set_exception(protocol, wrapped_exc, underlying_exc) + if tracker is not None: + tracker._attempt_failed(gen, wrapped_exc) + else: + # Successfully wrote the body, signal EOF and start response timeout + await writer.write_eof() + if tracker is not None: + tracker._attempt_finished(gen) + protocol.start_timeout() except asyncio.CancelledError: - # Body hasn't been fully sent, so connection can't be reused + # Record the attempt first: the tracker transitions cannot + # fail, while conn.close() conceivably can. + if tracker is not None: + tracker._attempt_cancelled(gen) + # Body hasn't been fully sent, so the connection can't be reused. + conn.close() + raise + except BaseException as underlying_exc: + # Failures escaping the inner handlers (100-continue preamble, + # write_eof) leave the body unsent: the connection can't be + # reused, and the tracker must still be notified. + if tracker is not None: + tracker._attempt_failed(gen, underlying_exc) conn.close() raise - except Exception as underlying_exc: - set_exception( - protocol, - ClientConnectionError( - "Failed to send bytes into the underlying connection " - f"{conn !s}: {underlying_exc!r}", - ), - underlying_exc, - ) - else: - # Successfully wrote the body, signal EOF and start response timeout - await writer.write_eof() - protocol.start_timeout() async def _close(self) -> None: if self._writer_task is not None: diff --git a/aiohttp/http_writer.py b/aiohttp/http_writer.py index a1168cfdebb..19c31529e34 100644 --- a/aiohttp/http_writer.py +++ b/aiohttp/http_writer.py @@ -182,14 +182,20 @@ async def write( if self._on_chunk_sent is not None: await self._on_chunk_sent(chunk) + notify = self.on_body_write if isinstance(chunk, memoryview): - if chunk.nbytes != len(chunk): + body_size = chunk.nbytes + if body_size != len(chunk): # just reshape it chunk = chunk.cast("c") + else: + body_size = len(chunk) if self._compress is not None: chunk = await self._compress.compress(chunk) if not chunk: + if notify is not None and body_size: + notify(body_size) return if self.length is not None: @@ -197,6 +203,8 @@ async def write( if self.length >= chunk_len: self.length = self.length - chunk_len else: + # Bytes clipped at the declared length are discarded, not sent. + body_size = self.length chunk = chunk[: self.length] self.length = 0 if not chunk: @@ -205,6 +213,8 @@ async def write( # Handle buffered headers for small payload optimization if self._headers_buf and not self._headers_written: self._send_headers_with_payload(chunk, False) + if notify is not None and body_size: + notify(body_size) if drain and self.buffer_size > LIMIT: self.buffer_size = 0 await self.drain() @@ -215,6 +225,8 @@ async def write( self._write_chunked_payload(chunk) else: self._write(chunk) + if notify is not None and body_size: + notify(body_size) if drain and self.buffer_size > LIMIT: self.buffer_size = 0 @@ -280,6 +292,9 @@ async def write_eof(self, chunk: bytes = b"") -> None: if chunk and self._on_chunk_sent is not None: await self._on_chunk_sent(chunk) + notify = self.on_body_write + body_size = len(chunk) + # Handle body/compression if self._compress: chunks: list[bytes] = [] @@ -308,6 +323,8 @@ async def write_eof(self, chunk: bytes = b"") -> None: else: # Coalesce headers with compressed data self._writelines((headers_buf, *chunks)) + if notify is not None and body_size: + notify(body_size) await self.drain() self._eof = True return @@ -320,6 +337,8 @@ async def write_eof(self, chunk: bytes = b"") -> None: self._writelines(chunks) else: self._write(chunks[0]) + if notify is not None and body_size: + notify(body_size) await self.drain() self._eof = True return @@ -328,6 +347,8 @@ async def write_eof(self, chunk: bytes = b"") -> None: if self._headers_buf and not self._headers_written: # Use helper to send headers with payload self._send_headers_with_payload(chunk, True) + if notify is not None and body_size: + notify(body_size) await self.drain() self._eof = True return @@ -341,12 +362,16 @@ async def write_eof(self, chunk: bytes = b"") -> None: ) else: self._write(b"0\r\n\r\n") + if notify is not None and body_size: + notify(body_size) await self.drain() self._eof = True return if chunk: self._write(chunk) + if notify is not None and body_size: + notify(body_size) await self.drain() self._eof = True diff --git a/docs/client_reference.rst b/docs/client_reference.rst index 72a68670a5c..fc460f6a9ea 100644 --- a/docs/client_reference.rst +++ b/docs/client_reference.rst @@ -564,6 +564,12 @@ The client session supports the context manager protocol for self closing. .. versionadded:: 3.12 + :param upload_tracker: An :class:`UploadTracker` instance observing + this request's body upload. + ``None`` by default (no tracking). + + .. versionadded:: 3.14.4 + :param int read_bufsize: Size of the read buffer (:attr:`ClientResponse.content`). ``None`` by default, it means that the session global value is used. @@ -1565,30 +1571,6 @@ Response object .. versionadded:: 3.2 - .. attribute:: output_size - - Number of bytes sent for this request. - - Pair with :attr:`upload_complete` to display upload progress:: - - async with session.post(url, data=mpwriter) as resp: - while not resp.upload_complete.done(): - print(f"uploaded {resp.output_size} bytes") - await asyncio.sleep(0.5) - print(f"upload complete: {resp.output_size} bytes") - - .. versionadded:: 3.14 - - .. attribute:: upload_complete - - An :class:`asyncio.Future` set when the request body has been fully sent. - - Use ``await resp.upload_complete`` to block until the upload finishes, or - ``resp.upload_complete.done()`` to poll from a progress-sampling loop - (see :attr:`output_size`). - - .. versionadded:: 3.14 - .. attribute:: content_type Read-only property with *content* part of *Content-Type* header. @@ -2368,6 +2350,67 @@ Utilities .. versionadded:: 3.14 +.. class:: UploadTracker() + :canonical: aiohttp.client_reqrep.UploadTracker + + Tracks upload progress of a single client request. + + Create a tracker inside a running event loop and pass it to a request + via the ``upload_tracker`` argument; read its attributes while the + request runs, e.g. from a progress-reporting task:: + + tracker = aiohttp.UploadTracker() + + async def report_progress() -> None: + while not tracker.upload_complete.done(): + print(f"uploaded {tracker.bytes_written} bytes") + await asyncio.sleep(0.5) + + progress = asyncio.create_task(report_progress()) + async with session.post(url, data=data, upload_tracker=tracker) as resp: + ... + await progress + # Raises if the upload failed, even when the request itself + # succeeded because the server answered early. + await tracker.upload_complete + + A tracker observes exactly one request: passing it to a second request + raises :exc:`RuntimeError`. Two requests uploading the same payload + each get their own tracker. + + .. attribute:: bytes_written + + Body bytes of the current upload attempt handed to the connection, + before transport-level transformations such as compression or + chunked framing. Reset to ``0`` when the request moves on to a new + attempt. + + .. attribute:: attempts + + Number of upload attempts dispatched: ``1`` for a plain request, + incremented each time the request is resent, e.g. following a + redirect or a middleware retry. A redirect that drops the body + (such as a *303 See Other*) still counts a zero-byte attempt. + + .. attribute:: upload_complete + + An :class:`asyncio.Future` settled once no more of the body will + be sent: with ``None`` when the final attempt wrote the body fully + (even if the server responded first, the future settles when + writing finishes); with the upload error when the final attempt + failed (the request raises the same error unless the server + already responded successfully); or with :exc:`UploadAbortedError` + when the body was never fully sent, e.g. the request failed before + writing started, was cancelled, or a middleware answered without + sending the request. + + When the server answers before reading the whole body, the request + succeeds without raising; await the future (or check its + :meth:`~asyncio.Future.exception`) to observe such upload failures. + + .. versionadded:: 3.14.4 + + .. class:: DigestAuthMiddleware(login, password, *, preemptive=True) :canonical: aiohttp.client_middleware_digest_auth.DigestAuthMiddleware @@ -2728,6 +2771,16 @@ All exceptions are available as members of *aiohttp* module. Derived from :exc:`ClientError` +.. class:: UploadAbortedError + :canonical: aiohttp.client_exceptions.UploadAbortedError + + Set on :attr:`UploadTracker.upload_complete` when the request body + was never fully sent. + + Derived from :exc:`ClientError` + + .. versionadded:: 3.14.4 + .. exception:: InvalidURL :canonical: aiohttp.client_exceptions.InvalidURL @@ -3010,6 +3063,8 @@ Hierarchy of exceptions * :exc:`ClientPayloadError` + * :exc:`UploadAbortedError` + * :exc:`ClientResponseError` * :exc:`~aiohttp.ClientHttpProxyError` diff --git a/tests/test_client_functional.py b/tests/test_client_functional.py index f0dbaa5da5b..107426c0538 100644 --- a/tests/test_client_functional.py +++ b/tests/test_client_functional.py @@ -17,6 +17,7 @@ import zlib from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import suppress +from types import SimpleNamespace from typing import TYPE_CHECKING, Any, NoReturn from unittest import mock @@ -1871,6 +1872,32 @@ async def handler(request: web.Request) -> web.Response: assert content == {"some": "data"} +async def test_upload_tracker_compressed_body(aiohttp_client: AiohttpClient) -> None: + """bytes_written reports pre-compression payload bytes, not wire bytes.""" + encodings: list[str | None] = [] + + async def handler(request: web.Request) -> web.Response: + encodings.append(request.headers.get(hdrs.CONTENT_ENCODING)) + # The server transparently inflates the body back to the original. + assert await request.read() == body + return web.Response() + + app = web.Application() + app.router.add_post("/", handler) + client = await aiohttp_client(app) + + tracker = aiohttp.UploadTracker() + body = b"x" * 8192 + async with client.post( + "/", data=body, compress="deflate", upload_tracker=tracker + ) as resp: + assert resp.status == 200 + + assert encodings == ["deflate"] + assert tracker.bytes_written == len(body) + assert tracker.upload_complete.result() is None + + async def test_POST_FILES(aiohttp_client: AiohttpClient, fname: pathlib.Path) -> None: content1 = fname.read_bytes() @@ -6006,7 +6033,7 @@ async def handler(request: web.Request) -> web.Response: assert resp.content.total_raw_bytes == int(resp.headers["Content-Length"]) -async def test_output_size_bytes(aiohttp_client: AiohttpClient) -> None: +async def test_upload_tracker_bytes(aiohttp_client: AiohttpClient) -> None: async def handler(request: web.Request) -> web.Response: await request.read() return web.Response() @@ -6015,12 +6042,17 @@ async def handler(request: web.Request) -> web.Response: app.router.add_post("/", handler) client = await aiohttp_client(app) + tracker = aiohttp.UploadTracker() body = b"x" * 1024 - async with client.post("/", data=body) as resp: - assert resp.output_size >= len(body) + async with client.post("/", data=body, upload_tracker=tracker) as resp: + assert resp.status == 200 + + assert tracker.attempts == 1 + assert tracker.bytes_written == len(body) + assert tracker.upload_complete.result() is None -async def test_output_size_multipart(aiohttp_client: AiohttpClient) -> None: +async def test_upload_tracker_multipart(aiohttp_client: AiohttpClient) -> None: async def handler(request: web.Request) -> web.Response: await request.read() return web.Response() @@ -6035,40 +6067,16 @@ async def handler(request: web.Request) -> web.Response: expected_body_size = mpwriter.size assert expected_body_size is not None - async with client.post("/", data=mpwriter) as resp: - assert resp.output_size >= expected_body_size - - -async def test_output_size_keepalive_isolated( - aiohttp_client: AiohttpClient, -) -> None: - """Each request on a keep-alive connection has its own counter.""" - transports: set[object] = set() - - async def handler(request: web.Request) -> web.Response: - transports.add(request.transport) - await request.read() - return web.Response() - - app = web.Application() - app.router.add_post("/", handler) - connector = aiohttp.TCPConnector(limit=1, force_close=False) - client = await aiohttp_client(app, connector=connector) - body = b"x" * 65536 - - async with client.post("/", data=body) as resp1: - size1 = resp1.output_size - - async with client.post("/", data=body) as resp2: - size2 = resp2.output_size + tracker = aiohttp.UploadTracker() + async with client.post("/", data=mpwriter, upload_tracker=tracker) as resp: + assert resp.status == 200 - assert len(transports) == 1 # Check keep-alive worked. - assert size1 >= len(body) - assert size1 == size2 + assert tracker.bytes_written == expected_body_size + assert tracker.upload_complete.result() is None -async def test_output_size_progress(aiohttp_client: AiohttpClient) -> None: - """output_size advances by exactly one chunk per yield.""" +async def test_upload_tracker_progress(aiohttp_client: AiohttpClient) -> None: + """bytes_written advances by exactly one chunk per yield.""" async def handler(request: web.Request) -> web.StreamResponse: response = web.StreamResponse() @@ -6096,26 +6104,24 @@ async def gated_body() -> AsyncIterator[bytes]: next_chunk.set() await sample_taken.wait() - async with client.post("/", data=gated_body()) as resp: + tracker = aiohttp.UploadTracker() + async with client.post("/", data=gated_body(), upload_tracker=tracker) as resp: samples: list[int] = [] for _ in range(num_chunks): await next_chunk.wait() next_chunk.clear() - samples.append(resp.output_size) - assert not resp.upload_complete.done() + samples.append(tracker.bytes_written) + assert not tracker.upload_complete.done() sample_taken.set() - await resp.upload_complete - assert resp.upload_complete.done() + await tracker.upload_complete await resp.read() - # Each sample after the first reflects exactly one more chunk on the wire. - chunked_framing = len(f"{chunk_size:x}".encode()) + 4 - deltas = [samples[i] - samples[i - 1] for i in range(1, len(samples))] - assert deltas == [chunk_size + chunked_framing] * (num_chunks - 1) + # Each sample reflects exactly one more payload chunk, without framing. + assert samples == [chunk_size * (i + 1) for i in range(num_chunks)] -async def test_output_size_get_request(aiohttp_client: AiohttpClient) -> None: - """GET request with no body still reports the request header byte count.""" +async def test_upload_tracker_no_body(aiohttp_client: AiohttpClient) -> None: + """A bodyless request settles the tracker with a zero-byte attempt.""" async def handler(request: web.Request) -> web.Response: return web.Response() @@ -6124,13 +6130,16 @@ async def handler(request: web.Request) -> web.Response: app.router.add_get("/", handler) client = await aiohttp_client(app) - async with client.get("/") as resp: - assert resp.output_size >= 0 + tracker = aiohttp.UploadTracker() + async with client.get("/", upload_tracker=tracker) as resp: + assert resp.status == 200 + assert tracker.attempts == 1 + assert tracker.bytes_written == 0 + assert tracker.upload_complete.result() is None -async def test_output_size_writer_released(aiohttp_client: AiohttpClient) -> None: - """Writer is dropped once body upload completes; output_size survives.""" +async def test_upload_tracker_expect100(aiohttp_client: AiohttpClient) -> None: async def handler(request: web.Request) -> web.Response: await request.read() return web.Response() @@ -6139,14 +6148,195 @@ async def handler(request: web.Request) -> web.Response: app.router.add_post("/", handler) client = await aiohttp_client(app) - body = b"x" * 1024 - async with client.post("/", data=body) as resp: + tracker = aiohttp.UploadTracker() + body = b"e" * 256 + async with client.post( + "/", data=body, expect100=True, upload_tracker=tracker + ) as resp: + assert resp.status == 200 + + assert tracker.attempts == 1 + assert tracker.bytes_written == len(body) + assert tracker.upload_complete.result() is None + + +async def test_upload_tracker_redirect_resends_body( + aiohttp_client: AiohttpClient, +) -> None: + """A 307 redirect resends the body as a new attempt of the same tracker.""" + + async def redirect(request: web.Request) -> NoReturn: + await request.read() + raise web.HTTPTemporaryRedirect("/final") + + async def final(request: web.Request) -> web.Response: + await request.read() + return web.Response() + + app = web.Application() + app.router.add_post("/", redirect) + app.router.add_post("/final", final) + client = await aiohttp_client(app) + + tracker = aiohttp.UploadTracker() + body = b"r" * 2048 + async with client.post("/", data=body, upload_tracker=tracker) as resp: + assert resp.status == 200 + + assert tracker.attempts == 2 + assert tracker.bytes_written == len(body) + assert tracker.upload_complete.result() is None + + +async def test_upload_tracker_redirect_drops_body( + aiohttp_client: AiohttpClient, +) -> None: + """A 303 redirect turns POST into a bodyless GET; the final attempt sent 0 bytes.""" + + async def redirect(request: web.Request) -> NoReturn: + await request.read() + raise web.HTTPSeeOther("/final") + + async def final(request: web.Request) -> web.Response: + return web.Response() + + app = web.Application() + app.router.add_post("/", redirect) + app.router.add_get("/final", final) + client = await aiohttp_client(app) + + tracker = aiohttp.UploadTracker() + async with client.post("/", data=b"d" * 512, upload_tracker=tracker) as resp: + assert resp.status == 200 + + assert tracker.attempts == 2 + assert tracker.bytes_written == 0 + assert tracker.upload_complete.result() is None + + +async def test_upload_tracker_redirect_to_unreachable_host( + aiohttp_client: AiohttpClient, +) -> None: + """A resend failing before its body write must not report success.""" + + async def redirect(request: web.Request) -> NoReturn: + await request.read() + raise web.HTTPTemporaryRedirect("http://127.0.0.1:1/") + + app = web.Application() + app.router.add_post("/", redirect) + client = await aiohttp_client(app) + + tracker = aiohttp.UploadTracker() + with pytest.raises(aiohttp.ClientConnectorError): + await client.post("/", data=b"r" * 2048, upload_tracker=tracker) + + # The first hop's completed upload is history; the dispatched resend + # is the outcome. + assert tracker.attempts == 2 + assert tracker.bytes_written == 0 + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) + + +async def test_upload_tracker_settles_after_early_response( + aiohttp_client: AiohttpClient, +) -> None: + """A server can respond before the body is sent; the future settles later.""" + body_unblocked = asyncio.Event() + + async def handler(request: web.Request) -> web.StreamResponse: + response = web.StreamResponse() + await response.prepare(request) + await response.write(b"x") + await request.read() + return response + + app = web.Application() + app.router.add_post("/", handler) + client = await aiohttp_client(app) + + chunk = b"z" * 4096 + + async def gated_body() -> AsyncIterator[bytes]: + yield chunk + await body_unblocked.wait() + yield chunk + + tracker = aiohttp.UploadTracker() + async with client.post("/", data=gated_body(), upload_tracker=tracker) as resp: + # Response headers arrived while the body is still being written. + assert not tracker.upload_complete.done() + body_unblocked.set() + await tracker.upload_complete + assert tracker.bytes_written == 2 * len(chunk) await resp.read() - assert resp._stream_writer is None - assert resp.output_size >= len(body) -async def test_upload_complete_no_body(aiohttp_client: AiohttpClient) -> None: +async def test_upload_tracker_error_after_response( + aiohttp_client: AiohttpClient, +) -> None: + """Upload failure after the request succeeded stays on the future.""" + body_unblocked = asyncio.Event() + + async def handler(request: web.Request) -> web.StreamResponse: + response = web.StreamResponse() + await response.prepare(request) + # Flush headers so the request succeeds while the body is still + # being written; keep the connection open by reading the body. + await response.write(b"x") + with suppress(Exception): + await request.read() + assert False + + app = web.Application() + app.router.add_post("/", handler) + client = await aiohttp_client(app) + + async def failing_body() -> AsyncIterator[bytes]: + yield b"a" * 100 + await body_unblocked.wait() + raise ValueError("boom") + + tracker = aiohttp.UploadTracker() + async with client.post("/", data=failing_body(), upload_tracker=tracker) as resp: + assert resp.status == 200 + assert not tracker.upload_complete.done() + body_unblocked.set() + with pytest.raises(aiohttp.ClientConnectionError): + await tracker.upload_complete + + assert tracker.attempts == 1 + + +async def test_upload_tracker_upload_error_propagated_to_caller( + aiohttp_client: AiohttpClient, +) -> None: + """The upload failure the request raises is the one on the future.""" + + async def handler(request: web.Request) -> web.Response: + with suppress(Exception): + await request.read() + assert False + + app = web.Application() + app.router.add_post("/", handler) + client = await aiohttp_client(app) + + async def failing_body() -> AsyncIterator[bytes]: + yield b"a" * 100 + raise ValueError("boom") + + tracker = aiohttp.UploadTracker() + with pytest.raises(aiohttp.ClientConnectionError) as exc_info: + await client.post("/", data=failing_body(), upload_tracker=tracker) + + assert tracker.attempts == 1 + assert tracker.upload_complete.exception() is exc_info.value + + +async def test_upload_tracker_connect_error(aiohttp_client: AiohttpClient) -> None: + """A request failing before the body write leaves the future cancelled.""" + async def handler(request: web.Request) -> web.Response: return web.Response() @@ -6154,12 +6344,325 @@ async def handler(request: web.Request) -> web.Response: app.router.add_get("/", handler) client = await aiohttp_client(app) - async with client.get("/") as resp: - assert resp.upload_complete.done() + tracker = aiohttp.UploadTracker() + with pytest.raises(aiohttp.ClientConnectorError): + await client.session.post( + "http://127.0.0.1:1/", data=b"x", upload_tracker=tracker + ) + assert tracker.attempts == 1 + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) -async def test_upload_complete_late_access(aiohttp_client: AiohttpClient) -> None: - """Accessing upload_complete after the upload finished returns a done future.""" + +async def test_upload_tracker_request_cancelled( + aiohttp_client: AiohttpClient, +) -> None: + async def handler(request: web.Request) -> web.Response: + await request.read() + assert False + + app = web.Application() + app.router.add_post("/", handler) + client = await aiohttp_client(app) + + first_chunk_sent = asyncio.Event() + + async def gated_body() -> AsyncIterator[bytes]: + yield b"a" * 100 + first_chunk_sent.set() + await asyncio.Event().wait() # Block until cancelled. + + tracker = aiohttp.UploadTracker() + task = asyncio.create_task( + client.post("/", data=gated_body(), upload_tracker=tracker) + ) + await first_chunk_sent.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + with pytest.raises(aiohttp.UploadAbortedError): + await tracker.upload_complete + assert tracker.attempts == 1 + + +async def test_upload_tracker_middleware_short_circuit( + aiohttp_client: AiohttpClient, +) -> None: + """A middleware answering without sending the request aborts the upload.""" + + async def handler(request: web.Request) -> web.Response: + await request.read() + return web.Response() + + app = web.Application() + app.router.add_route("*", "/", handler) + client = await aiohttp_client(app) + + cached = await client.get("/") + await cached.read() + cached.release() + + async def cache_mw( + req: aiohttp.ClientRequest, handler_: aiohttp.ClientHandlerType + ) -> aiohttp.ClientResponse: + return cached + + tracker = aiohttp.UploadTracker() + async with client.post( + "/", data=b"x" * 128, middlewares=(cache_mw,), upload_tracker=tracker + ) as resp: + assert resp is cached + + assert tracker.attempts == 0 + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) + + +async def test_upload_tracker_middleware_resend( + aiohttp_client: AiohttpClient, +) -> None: + """A middleware resending the request counts a new attempt.""" + server_hits = 0 + + async def handler(request: web.Request) -> web.Response: + nonlocal server_hits + await request.read() + server_hits += 1 + return web.Response(status=401 if server_hits == 1 else 200) + + app = web.Application() + app.router.add_post("/", handler) + client = await aiohttp_client(app) + + async def retry_mw( + req: aiohttp.ClientRequest, handler_: aiohttp.ClientHandlerType + ) -> aiohttp.ClientResponse: + resp = await handler_(req) + assert resp.status == 401 + resp.release() + resp = await handler_(req) + return resp + + tracker = aiohttp.UploadTracker() + body = b"b" * 512 + async with client.post( + "/", data=body, middlewares=(retry_mw,), upload_tracker=tracker + ) as resp: + assert resp.status == 200 + + assert tracker.attempts == 2 + assert tracker.bytes_written == len(body) + assert tracker.upload_complete.result() is None + + +async def test_upload_tracker_rebind_raises(aiohttp_client: AiohttpClient) -> None: + async def handler(request: web.Request) -> web.Response: + return web.Response() + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app) + + tracker = aiohttp.UploadTracker() + async with client.get("/", upload_tracker=tracker): + pass + + with pytest.raises(RuntimeError, match="already bound"): + await client.get("/", upload_tracker=tracker) + + +async def test_upload_tracker_non_http_scheme(aiohttp_client: AiohttpClient) -> None: + """A request rejected before being built settles the tracker.""" + + async def handler(request: web.Request) -> web.Response: + assert False + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app) + + tracker = aiohttp.UploadTracker() + with pytest.raises(aiohttp.NonHttpUrlClientError): + await client.session.get("ftp://example.com/", upload_tracker=tracker) + + assert tracker.attempts == 0 + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) + + +async def test_upload_tracker_trace_start_failure( + aiohttp_client: AiohttpClient, +) -> None: + """A failing on_request_start trace disarms the timeout and settles the tracker.""" + + async def on_request_start( + session: aiohttp.ClientSession, + ctx: SimpleNamespace, + params: aiohttp.TraceRequestStartParams, + ) -> None: + raise RuntimeError("trace boom") + + trace_config = aiohttp.TraceConfig() + trace_config.on_request_start.append(on_request_start) + + async def handler(request: web.Request) -> web.Response: + assert False + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app, trace_configs=[trace_config]) + + tracker = aiohttp.UploadTracker() + with pytest.raises(RuntimeError, match="trace boom"): + await client.get("/", upload_tracker=tracker) + + assert tracker.attempts == 0 + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) + + +async def test_trace_start_failure_notifies_started_traces( + aiohttp_client: AiohttpClient, +) -> None: + """Traces whose on_request_start already ran get the terminal event.""" + seen_exceptions: list[BaseException] = [] + + async def good_start( + session: aiohttp.ClientSession, + ctx: SimpleNamespace, + params: aiohttp.TraceRequestStartParams, + ) -> None: + """Complete normally so this config expects a terminal event.""" + + async def good_exception( + session: aiohttp.ClientSession, + ctx: SimpleNamespace, + params: aiohttp.TraceRequestExceptionParams, + ) -> None: + seen_exceptions.append(params.exception) + + good = aiohttp.TraceConfig() + good.on_request_start.append(good_start) + good.on_request_exception.append(good_exception) + + bad_exception_calls: list[BaseException] = [] + + async def bad_start( + session: aiohttp.ClientSession, + ctx: SimpleNamespace, + params: aiohttp.TraceRequestStartParams, + ) -> None: + raise RuntimeError("trace boom") + + async def bad_exception( + session: aiohttp.ClientSession, + ctx: SimpleNamespace, + params: aiohttp.TraceRequestExceptionParams, + ) -> None: + bad_exception_calls.append(params.exception) # pragma: no cover + + bad = aiohttp.TraceConfig() + bad.on_request_start.append(bad_start) + bad.on_request_exception.append(bad_exception) + + async def handler(request: web.Request) -> web.Response: + assert False + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app, trace_configs=[good, bad]) + + with pytest.raises(RuntimeError, match="trace boom"): + await client.get("/") + + assert len(seen_exceptions) == 1 + assert isinstance(seen_exceptions[0], RuntimeError) + # The config whose start never completed gets no terminal event. + assert not bad_exception_calls + + +@pytest.mark.parametrize( + "bad_kwargs", + ( + {"data": b"x", "json": {"a": 1}}, + {"json": {"a": object()}}, + {"ssl": object()}, + ), + ids=("data-and-json", "unserializable-json", "bad-ssl-type"), +) +async def test_upload_tracker_validation_failure( + aiohttp_client: AiohttpClient, bad_kwargs: dict[str, object] +) -> None: + """Argument validation failures settle the tracker instead of hanging it.""" + + async def handler(request: web.Request) -> web.Response: + assert False + + app = web.Application() + app.router.add_post("/", handler) + client = await aiohttp_client(app) + + tracker = aiohttp.UploadTracker() + with pytest.raises((TypeError, ValueError)): + await client.post("/", upload_tracker=tracker, **bad_kwargs) # type: ignore[arg-type] + + assert tracker.attempts == 0 + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) + + +async def test_upload_tracker_invalid_timeout_type( + aiohttp_client: AiohttpClient, +) -> None: + """A wrong-typed timeout argument still settles the tracker.""" + + async def handler(request: web.Request) -> web.Response: + assert False + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app) + + tracker = aiohttp.UploadTracker() + with pytest.raises(AttributeError): + await client.get("/", timeout=5, upload_tracker=tracker) # type: ignore[arg-type] + + assert tracker.attempts == 0 + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) + + +async def test_upload_tracker_closed_session() -> None: + """A request on a closed session settles the tracker.""" + session = aiohttp.ClientSession() + await session.close() + + tracker = aiohttp.UploadTracker() + with pytest.raises(RuntimeError, match="Session is closed"): + await session.get("http://example.com/", upload_tracker=tracker) + + assert tracker.attempts == 0 + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) + + +async def test_upload_tracker_invalid_proxy(aiohttp_client: AiohttpClient) -> None: + """An invalid proxy URL settles the tracker before the request is built.""" + + async def handler(request: web.Request) -> web.Response: + assert False + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app) + + tracker = aiohttp.UploadTracker() + with pytest.raises(InvalidURL): + await client.get("/", proxy="http://[invalid", upload_tracker=tracker) + + assert tracker.attempts == 0 + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) + + +async def test_upload_complete_deprecated_late_access( + aiohttp_client: AiohttpClient, +) -> None: + """First access after the upload finished returns an already-done future.""" async def handler(request: web.Request) -> web.Response: await request.read() @@ -6171,6 +6674,46 @@ async def handler(request: web.Request) -> web.Response: async with client.post("/", data=b"x" * 1024) as resp: await resp.read() - # Writer task is done; future is created lazily on this first access. - assert resp._upload_complete is None - assert resp.upload_complete.done() + with pytest.warns(DeprecationWarning, match="upload_complete is deprecated"): + fut = resp.upload_complete + assert fut.done() + + +async def test_deprecated_upload_attributes(aiohttp_client: AiohttpClient) -> None: + """The deprecated ClientResponse attributes still work in both writer states.""" + body_unblocked = asyncio.Event() + + async def handler(request: web.Request) -> web.StreamResponse: + response = web.StreamResponse() + await response.prepare(request) + # Flush headers so the client sees the response mid-upload. + await response.write(b"x") + await request.read() + return response + + app = web.Application() + app.router.add_post("/", handler) + client = await aiohttp_client(app) + + async def gated_body() -> AsyncIterator[bytes]: + yield b"x" * 512 + await body_unblocked.wait() + yield b"y" * 512 + + async with client.post("/", data=gated_body()) as resp: + # The writer is still active on first access. + with pytest.warns(DeprecationWarning, match="output_size is deprecated"): + assert resp.output_size >= 0 + with pytest.warns(DeprecationWarning, match="upload_complete is deprecated"): + fut = resp.upload_complete + assert not fut.done() + + body_unblocked.set() + await fut + + # The writer has been released; the same future is returned. + with pytest.warns(DeprecationWarning, match="upload_complete is deprecated"): + assert resp.upload_complete is fut + with pytest.warns(DeprecationWarning, match="output_size is deprecated"): + assert resp.output_size >= 1024 + await resp.read() diff --git a/tests/test_client_request.py b/tests/test_client_request.py index 78c36d26eca..4827c14f37e 100644 --- a/tests/test_client_request.py +++ b/tests/test_client_request.py @@ -23,6 +23,7 @@ ClientResponse, ClientTimeout, Fingerprint, + UploadTracker, _gen_default_accept_encoding, ) from aiohttp.compression_utils import ZLibBackend @@ -2323,3 +2324,223 @@ async def test_empty_body_isolation_after_update( assert ClientRequest._EMPTY_BODY.consumed is False assert ClientRequest._EMPTY_BODY.size == 0 + + +async def test_upload_tracker_stale_attempt_events_ignored() -> None: + """Events from a superseded attempt must not affect the current one.""" + tracker = UploadTracker() + stale = tracker._attempt_started() + gen = tracker._attempt_started() + tracker._attempt_writing(gen) + + tracker._add_bytes(gen, 10) + tracker._add_bytes(stale, 100) + assert tracker.bytes_written == 10 + + tracker._attempt_failed(stale, RuntimeError("stale")) + tracker._attempt_cancelled(stale) + tracker._attempt_finished(stale) + assert not tracker.upload_complete.done() + + tracker._finalize() + # The final attempt is still writing; its own completion settles. + assert not tracker.upload_complete.done() + tracker._attempt_finished(gen) + assert tracker.upload_complete.result() is None + assert tracker.attempts == 2 + + +async def test_upload_tracker_failed_attempt_superseded_by_resend() -> None: + """Only the final attempt's outcome is reported after a resend.""" + tracker = UploadTracker() + gen = tracker._attempt_started() + tracker._attempt_failed(gen, RuntimeError("first try")) + assert not tracker.upload_complete.done() + + gen = tracker._attempt_started() + tracker._attempt_finished(gen) + tracker._finalize() + assert tracker.upload_complete.result() is None + + +async def test_upload_tracker_cancelled_attempt_superseded_by_resend() -> None: + """A cancelled non-final attempt leaves the future pending for a resend.""" + tracker = UploadTracker() + gen = tracker._attempt_started() + tracker._attempt_cancelled(gen) + assert not tracker.upload_complete.done() + + gen = tracker._attempt_started() + tracker._attempt_finished(gen) + tracker._finalize() + assert tracker.upload_complete.result() is None + + +async def test_upload_tracker_externally_cancelled_future() -> None: + """User code cancelling upload_complete must not break settling.""" + tracker = UploadTracker() + tracker.upload_complete.cancel() + gen = tracker._attempt_started() + tracker._attempt_finished(gen) + tracker._finalize() + assert tracker.upload_complete.cancelled() + + +async def test_oserror_on_write_bytes_with_upload_tracker( + conn: mock.Mock, make_client_request: _RequestMaker +) -> None: + """A write failure is recorded on the request's UploadTracker.""" + loop = asyncio.get_running_loop() + req = make_client_request("POST", URL("http://python.org/"), loop=loop) + await req.update_body(b"test data") + tracker = UploadTracker() + req._upload_tracker = tracker + req._upload_gen = tracker._attempt_started() + + writer = WriterMock() + writer.write.side_effect = OSError + + await req._write_bytes(writer, conn, None) + + tracker._finalize() + assert tracker.attempts == 1 + assert isinstance(tracker.upload_complete.exception(), aiohttp.ClientOSError) + + +@pytest.mark.parametrize("with_tracker", (True, False)) +async def test_preamble_failure_reported_to_upload_tracker( + conn: mock.Mock, make_client_request: _RequestMaker, with_tracker: bool +) -> None: + """A failure before the body write (100-continue preamble) is recorded.""" + loop = asyncio.get_running_loop() + req = make_client_request( + "POST", URL("http://python.org/"), data=b"test data", expect100=True, loop=loop + ) + tracker = UploadTracker() if with_tracker else None + req._upload_tracker = tracker + if tracker is not None: + req._upload_gen = tracker._attempt_started() + + writer = WriterMock() + writer.send_headers = mock.Mock() + writer.drain.side_effect = RuntimeError("preamble boom") + + with pytest.raises(RuntimeError, match="preamble boom"): + await req._write_bytes(writer, conn, None) + + # The body was never sent on a connection with headers on the wire. + assert conn.close.called + if tracker is not None: + tracker._finalize() + assert tracker.attempts == 1 + assert isinstance(tracker.upload_complete.exception(), RuntimeError) + + +@pytest.mark.skipif( + sys.version_info < (3, 11), reason="TimeoutError is OSError only on 3.11+" +) +async def test_timeout_on_write_bytes_not_wrapped( + conn: mock.Mock, make_client_request: _RequestMaker +) -> None: + """An asyncio.TimeoutError from the write is not wrapped in ClientOSError.""" + loop = asyncio.get_running_loop() + req = make_client_request("POST", URL("http://python.org/"), loop=loop) + await req.update_body(b"test data") + + writer = WriterMock() + writer.write.side_effect = asyncio.TimeoutError + + await req._write_bytes(writer, conn, None) + + assert conn.protocol.set_exception.called + exc = conn.protocol.set_exception.call_args[0][0] + assert type(exc) is asyncio.TimeoutError + + +async def test_upload_tracker_unretrieved_error_not_logged() -> None: + """A settled upload error must not log 'Future exception was never retrieved'.""" + tracker = UploadTracker() + gen = tracker._attempt_started() + tracker._attempt_failed(gen, RuntimeError("boom")) + tracker._finalize() + + # The documented pattern only polls done(), so the settle itself must + # have consumed the log; the error stays retrievable. + assert tracker.upload_complete._log_traceback is False + assert isinstance(tracker.upload_complete.exception(), RuntimeError) + + +async def test_cancel_during_expect100_preamble_closes_connection( + conn: mock.Mock, make_client_request: _RequestMaker +) -> None: + """A writer cancelled while waiting for 100-continue closes the connection. + + The request headers are already on the wire at that point, so releasing + the connection for reuse would corrupt the next request on it. + """ + loop = asyncio.get_running_loop() + req = make_client_request( + "POST", URL("http://python.org/"), data=b"test data", expect100=True, loop=loop + ) + + writer = WriterMock() + writer.send_headers = mock.Mock() + + task = asyncio.create_task(req._write_bytes(writer, conn, None)) + # Let the writer park on the 100-continue waiter. + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert conn.close.called + + +async def test_conn_close_failure_still_settles_upload_tracker( + conn: mock.Mock, make_client_request: _RequestMaker +) -> None: + """The attempt is recorded before conn.close(), which conceivably raises.""" + loop = asyncio.get_running_loop() + req = make_client_request( + "POST", URL("http://python.org/"), data=b"test data", expect100=True, loop=loop + ) + tracker = UploadTracker() + req._upload_tracker = tracker + req._upload_gen = tracker._attempt_started() + conn.close.side_effect = RuntimeError("close boom") + + writer = WriterMock() + writer.send_headers = mock.Mock() + + task = asyncio.create_task(req._write_bytes(writer, conn, None)) + # Let the writer park on the 100-continue waiter. + await asyncio.sleep(0) + task.cancel() + with pytest.raises(RuntimeError, match="close boom"): + await task + + tracker._finalize() + assert tracker.upload_complete.done() + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) + + +async def test_upload_tracker_dispatched_resend_never_writing() -> None: + """A dispatched resend that never starts writing settles as aborted.""" + tracker = UploadTracker() + gen = tracker._attempt_started() + tracker._attempt_writing(gen) + tracker._add_bytes(gen, 2048) + tracker._attempt_finished(gen) + + tracker._attempt_started() # Resend dispatched; never writes. + assert tracker.bytes_written == 0 + # Events from the superseded attempt are ignored; a stale writing + # transition must not make _finalize() defer to a dead writer. + tracker._attempt_writing(gen) + tracker._attempt_finished(gen) + tracker._add_bytes(gen, 100) + assert tracker.bytes_written == 0 + + tracker._finalize() + assert tracker.attempts == 2 + assert isinstance(tracker.upload_complete.exception(), aiohttp.UploadAbortedError) diff --git a/tests/test_http_writer.py b/tests/test_http_writer.py index baab763c297..f81c08f85fc 100644 --- a/tests/test_http_writer.py +++ b/tests/test_http_writer.py @@ -1595,3 +1595,92 @@ async def test_send_headers_with_payload_chunked_eof_no_data( assert b"GET /test HTTP/1.1\r\n" in buf assert b"Transfer-Encoding: chunked\r\n" in buf assert buf.endswith(b"0\r\n\r\n") + + +async def test_on_body_write_counts_payload_bytes( + buf: bytearray, protocol: BaseProtocol, transport: asyncio.Transport +) -> None: + """The on_body_write callback reports accepted body bytes.""" + sizes: list[int] = [] + msg = http.StreamWriter(protocol, asyncio.get_running_loop()) + msg.on_body_write = sizes.append + + await msg.write(b"a" * 4) + await msg.write(b"") + await msg.write_eof(b"b" * 2) + + assert sizes == [4, 2] + + +async def test_on_body_write_clipped_chunk( + buf: bytearray, protocol: BaseProtocol, transport: asyncio.Transport +) -> None: + """Bytes clipped at the declared length are not reported.""" + sizes: list[int] = [] + msg = http.StreamWriter(protocol, asyncio.get_running_loop()) + msg.on_body_write = sizes.append + msg.length = 2 + + await msg.write(b"a" * 4) + await msg.write(b"b" * 4) + + assert buf == b"aa" + assert sizes == [2] + + +async def test_on_body_write_compressed_chunk( + buf: bytearray, protocol: BaseProtocol, transport: asyncio.Transport +) -> None: + """With compression the pre-compression payload size is reported.""" + sizes: list[int] = [] + msg = http.StreamWriter(protocol, asyncio.get_running_loop()) + msg.on_body_write = sizes.append + msg.enable_compression("deflate") + + # The compressor emits only its stream header for the first small + # chunk and buffers the second one entirely; both writes send no + # payload yet but are still reported as accepted. + buffered = b"y" * 16 + await msg.write(buffered) + await msg.write(buffered) + + chunk = b"x" * 8192 + await msg.write(chunk) + await msg.write_eof() + + assert sizes == [len(buffered), len(buffered), len(chunk)] + # The wire carried the compressed form, smaller than what was reported. + assert 0 < len(buf) < len(chunk) + + +@pytest.mark.parametrize( + ("compress", "chunked", "buffer_headers"), + ( + (True, False, True), + (True, False, False), + (False, False, True), + (False, True, False), + ), +) +async def test_on_body_write_write_eof_final_chunk( + buf: bytearray, + protocol: BaseProtocol, + transport: asyncio.Transport, + compress: bool, + chunked: bool, + buffer_headers: bool, +) -> None: + """A final chunk is reported across the framing/compression paths.""" + sizes: list[int] = [] + msg = http.StreamWriter(protocol, asyncio.get_running_loop()) + msg.on_body_write = sizes.append + if compress: + msg.enable_compression("deflate") + if chunked: + msg.enable_chunking() + if buffer_headers: + await msg.write_headers("POST / HTTP/1.1", CIMultiDict()) + + await msg.write_eof(b"x" * 64) + + assert sizes == [64]