diff --git a/clients/python/README.md b/clients/python/README.md index 318367e2..3fea83bb 100644 --- a/clients/python/README.md +++ b/clients/python/README.md @@ -120,6 +120,127 @@ Arbitrary key-value pairs can be attached to objects and retrieved on download. session.put(b"payload", metadata={"source": "upload-service"}) ``` +### Many API (batch operations) + +`session.many()` executes a number of operations with as few requests as +possible. Those within the batch protocol's per-part limit of 1 MB are grouped +into requests to Objectstore's batch endpoint which cuts network overhead +considerably. Inserts too large for that, of unknown size, or compressed with +gzip by the client are sent as individual requests instead. + +Pass any iterable of `Get`, `Put`, `Delete`, and `Head` operations, which live in +the `many` module. Results come back as `GetResult`, `PutResult`, `DeleteResult`, +and `HeadResult` objects, each carrying the object's `key` and an `error` that is +`None` when the operation succeeded: + +```python +from objectstore_client import Client, Usecase, many + +client = Client("http://localhost:8888") +session = client.session(Usecase("attachments"), org=42, project=1337) + +results = session.many( + [ + many.Put(b"file1 contents", key="file1"), + many.Put(b"file2 contents", key="file2"), + many.Get("file3"), + many.Delete("file4"), + many.Head("file5"), + ] +) + +for result in results: + if result.error is not None: + ... # this operation failed + elif isinstance(result, many.GetResult): + # `response` is None if the object does not exist. + payload = result.response.payload if result.response else None +``` + +`session.many()` returns an `OperationResults` object which is a lazy iterator. +As the iterator is consumed, it assembles batch requests and sends them to +Objectstore. As responses come in, the operation results are yielded. Abandoning +the iterator without fully consuming it will cancel whatever has not been +dispatched yet. + +If successful results don't need to be processed or inspected, callers can call +`raise_for_failures()` to drain the results and raise an `ExceptionGroup` with +all per-operation errors, or `failures()` which returns the failed results as a +list: + +```python +session.many([many.Delete("file1"), many.Delete("file2")]).raise_for_failures() + +for failure in session.many([many.Delete("file3")]).failures(): + print(failure.key, failure.error) +``` + +#### Concurrency + +`concurrency` caps how many requests are in flight, and defaults to `3`. +Requests run on a thread pool created by `session.many()`, which is shut down +when the results are exhausted or abandoned. When `concurrency` is set to `1`, +requests are run serially on the caller thread instead, with no thread pool. + +```python +client = Client("http://localhost:8888") +session = client.session(Usecase("attachments"), org=42, project=1337) + +for result in session.many(operations, concurrency=8): + ... +``` + +Note: when a `Client` is built with custom `connection_kwargs` that include +`"block": True`, the `concurrency` argument is clamped to the connection pool's +configured size. An illustrative example: + +```python +# Client created with a pool size of 4 and block=True +client = Client( + "http://localhost:8888", + connection_kwargs={"maxsize": 4, "block": True}, +) +session = client.session(Usecase("attachments"), org=42, project=1337) + +# `concurrency` is clamped to `4` because `block=True` was set +for failure in session.many(operations, concurrency=16).failures(): + print(failure.key, failure.error) +``` + +Results are yielded as responses are received, and the order isn't necessarily +the same order that operations were given in. Each result carries an `index` +field that corresponds to the index of the `Get` / `Put` / `Delete` / `Head` +operation in the operation iterable passed into `session.many()`. This `index` +allows a keyless `Put` operation to be linked with its result to learn the key +that was assigned. + +```python +uploads = [many.Put(b"first"), many.Put(b"second")] + +for result in session.many(uploads): + print(f"{uploads[result.index].contents!r} was stored as {result.key}") +``` + +A response part that cannot be attributed to any operation at all is reported +as an `ErrorResult`, whose `index` is always `None`. Every other result carries +the index of the operation it belongs to. + +Within a single batch, the Objectstore server processes individual operations +concurrently and each operation's relative order is undefined. Two operations on +the same key therefore race, and `session.many()` does nothing to prevent that. + +#### Metrics + +When a metrics backend is configured, `session.many()` emits some metrics: +- `storage.batch.latency`: a timer recording a batch request's execution time, + tagged with a (bucketed) number of operations included in the batch +- `stoarge.batch.operations`: a simple counter of individual operations, tagged + with each operation's kind (i.e. `PUT`/`GET`/`DELETE`). + +An operation that doesn't qualify for batching will be sent through the +`session`'s regular single-operation API for that operation and will emit +single-object metrics on that path rather than batch metrics here. + ### Authentication If your Objectstore instance enforces authorization, you must configure authentication diff --git a/clients/python/docs/objectstore_client.rst b/clients/python/docs/objectstore_client.rst index 061b2131..fc43e200 100644 --- a/clients/python/docs/objectstore_client.rst +++ b/clients/python/docs/objectstore_client.rst @@ -33,6 +33,14 @@ objectstore\_client.errors module :show-inheritance: :undoc-members: +objectstore\_client.many module +------------------------------- + +.. automodule:: objectstore_client.many + :members: + :show-inheritance: + :undoc-members: + objectstore\_client.metadata module ----------------------------------- diff --git a/clients/python/src/objectstore_client/__init__.py b/clients/python/src/objectstore_client/__init__.py index 8cd93138..ca655025 100644 --- a/clients/python/src/objectstore_client/__init__.py +++ b/clients/python/src/objectstore_client/__init__.py @@ -1,3 +1,4 @@ +from objectstore_client import many from objectstore_client.auth import ( Permission, SecretKey, @@ -34,6 +35,7 @@ "RequestError", "ObjectNotFound", "ExpiryExtensionRejected", + "many", "Compression", "ExpirationPolicy", "Metadata", diff --git a/clients/python/src/objectstore_client/client.py b/clients/python/src/objectstore_client/client.py index a18d6ddb..98338a84 100644 --- a/clients/python/src/objectstore_client/client.py +++ b/clients/python/src/objectstore_client/client.py @@ -1,14 +1,14 @@ from __future__ import annotations import math -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from contextlib import nullcontext from dataclasses import asdict, dataclass from datetime import UTC, datetime, timedelta from importlib.metadata import version from io import BytesIO from os import PathLike -from typing import IO, Any, Literal, NamedTuple, cast +from typing import IO, TYPE_CHECKING, Any, Literal, NamedTuple, cast from urllib.parse import urlparse import urllib3 @@ -47,6 +47,9 @@ # Identifies this library to the server, mirroring the Rust client's user agent. USER_AGENT = f"objectstore-client/{version('objectstore-client')}" +if TYPE_CHECKING: + from objectstore_client.many import Operation, OperationResults + def _response_options_query( disposition: Literal["inline", "attachment"], @@ -382,6 +385,69 @@ def _make_url(self, key: str | None, full: bool = False) -> str: return f"{self._base_url()}{path}" return path + def _make_batch_url(self) -> str: + relative_path = f"/v1/objects:batch/{self._usecase.name}/{self._scope}/" + return utils.encode_path(self._base_path.rstrip("/") + relative_path) + + def many( + self, + operations: Iterable[Operation], + *, + concurrency: int | None = None, + ) -> OperationResults: + """ + Executes multiple operations, batching them where possible. + + Operations that satisfy the batch protocol's per-part size limit of 1MB + are grouped into batch requests to reduce network overhead. Inserts with + larger sizes (or unknown sizes) and inserts the client compresses with + gzip are sent as individual requests instead. + + Operations run concurrently and in no particular order, so two + operations on the same key race. Sequence them with separate calls. + + Args: + operations: The operations to execute, as + :class:`~objectstore_client.many.Get`, + :class:`~objectstore_client.many.Put`, + :class:`~objectstore_client.many.Delete`, and + :class:`~objectstore_client.many.Head` instances. Any iterable + works. + concurrency: The maximum number of requests in flight. Defaults to + ``3``. Pass ``1`` to run everything sequentially on the calling + thread without a thread pool. A client created with + ``connection_kwargs`` that include ``"block": True`` caps the + size of its connection pool explicitly so, in that case, + ``concurrency`` will be clamped to its ``maxsize``. + + Returns: + An :class:`~objectstore_client.many.OperationResults` iterator over + the results. Results are yielded as responses come in, in no + particular order; each carries the ``index`` of the operation it + belongs to, which is the only handle on a keyless + :class:`~objectstore_client.many.Put`. This iterator is lazy, and if + it's abandoned without being fully consumed then operations that + haven't been dispatched are cancelled. + + Raises: + ValueError: If ``concurrency`` is less than ``1``. + + Example:: + + from objectstore_client import many + + results = session.many([many.Put(b"hello", key="k1"), many.Get("k2")]) + for result in results: + if result.error is not None: + ... # this operation failed + elif isinstance(result, many.GetResult): + ... # `result.response` is None if the object does not exist + """ + # Imported lazily to avoid a circular import at module load time. + from objectstore_client.many import execute_many + + return execute_many(self, operations, concurrency=concurrency) + def put( self, contents: bytes | IO[bytes], diff --git a/clients/python/src/objectstore_client/formdata.py b/clients/python/src/objectstore_client/formdata.py new file mode 100644 index 00000000..31a266d5 --- /dev/null +++ b/clients/python/src/objectstore_client/formdata.py @@ -0,0 +1,191 @@ +""" +Streaming ``multipart/form-data`` codec for batch requests. + +`urllib3`'s `encode_multipart_formdata` doesn't support per-part custom headers, +but we need them for operation-kind/key/index headers. +""" + +from __future__ import annotations + +import re +import secrets +from collections.abc import Iterable, Iterator, Sequence +from dataclasses import dataclass, field +from typing import IO + +# 64 KiB matches urllib3's default chunk size. +CHUNK_SIZE = 64 * 1024 + +# Header name/value regular expressions adapted from `h11`'s `_abnf.py` (MIT). +_VALID_HEADER_NAME = re.compile(r"[-!#$%&'*+.^_`|~0-9a-zA-Z]+") +_VISIBLE_CHARS = r"[\x21-\x7e]" +_VALID_HEADER_VALUE = re.compile(rf"({_VISIBLE_CHARS}+(?:[ \t]+{_VISIBLE_CHARS}+)*)?") + +_BOUNDARY_RE = re.compile(r'boundary="?([^";]+)"?') + +PartBody = bytes | IO[bytes] + + +@dataclass +class RequestPart: + """A single part of a multipart request.""" + + headers: dict[str, str] + body: PartBody = b"" + + def validate(self) -> None: + """ + Raises ``ValueError`` if a header cannot be represented on the wire. + """ + for name, value in self.headers.items(): + _validate_header(name, value) + + +@dataclass +class ResponsePart: + """A single part parsed from a multipart response.""" + + headers: dict[str, str] + """Part headers, with names lowercased.""" + + body: bytes = field(repr=False) + + +class MalformedMultipart(ValueError): + """Raised when a multipart body cannot be parsed.""" + + +class MultipartBody: + """ + A lazily encoded ``multipart/form-data`` request body. Pass the instance as + the ``body`` param in a urllib3 request and :attr:`content_type` as the + request's ``Content-Type``. + + :attr:`resendable` indicates whether the instance can be iterated over + multiple times and, consequently, whether urllib3 can retry requests. True + when every part's body is ``bytes``. + """ + + def __init__(self, parts: Sequence[RequestPart]): + for part in parts: + part.validate() + self._parts = parts + self._boundary = f"os-boundary-{secrets.token_hex(16)}" + self.resendable = all(isinstance(part.body, bytes) for part in parts) + self.content_type = f'multipart/form-data; boundary="{self._boundary}"' + + def __iter__(self) -> Iterator[bytes]: + for part in self._parts: + yield self._preamble(part) + if isinstance(part.body, bytes): + if part.body: + yield part.body + else: + while chunk := part.body.read(CHUNK_SIZE): + yield chunk + yield b"\r\n" + yield f"--{self._boundary}--\r\n".encode() + + def _preamble(self, part: RequestPart) -> bytes: + lines = [ + f"--{self._boundary}", + "content-disposition: form-data; name=part", + ] + lines += [f"{name}: {value}" for name, value in part.headers.items()] + return ("\r\n".join(lines) + "\r\n\r\n").encode() + + +def _validate_header(name: str, value: str) -> None: + if not _VALID_HEADER_NAME.fullmatch(name): + raise ValueError(f"invalid multipart header name: {name!r}") + if not _VALID_HEADER_VALUE.fullmatch(value): + raise ValueError(f"invalid multipart header value for {name!r}") + + +def _extract_boundary(content_type: str) -> str: + """ + Returns the ``boundary`` parameter of a ``multipart/*`` content type. + + Raises :class:`MalformedMultipart` if the content type carries no boundary. + """ + match = _BOUNDARY_RE.search(content_type) + if not match: + raise MalformedMultipart( + f"no multipart boundary in Content-Type: {content_type!r}" + ) + return match.group(1) + + +def iter_multipart( + content_type: str, chunks: Iterable[bytes] +) -> Iterator[ResponsePart]: + """ + Parses a multipart response body, yielding parts as they come in. + + Accepts an iterable of byte chunks (e.g. from urllib's `response.stream()`). + Parts are yielded as soon as they arrive in full. + + Raises :class:`MalformedMultipart` if the body's ``boundary`` isn't valid + and consistent. + """ + boundary = _extract_boundary(content_type) + opening = f"--{boundary}\r\n".encode() + delimiter = f"\r\n--{boundary}".encode() + + buffer = bytearray() + started = False + + for chunk in chunks: + buffer.extend(chunk) + + if not started: + position = buffer.find(opening) + if position == -1: + # Discard everything except the last len(opening)-1 bytes - + # only that suffix could form a partial match across the next chunk. + del buffer[: -(len(opening) - 1)] + continue + del buffer[: position + len(opening)] + started = True + + while True: + position = buffer.find(delimiter) + if position == -1: + break + trailer = position + len(delimiter) + # Need at least 2 bytes after the delimiter to distinguish + # \r\n (next part follows) from -- (closing boundary). + if len(buffer) < trailer + 2: + break + yield _parse_part(bytes(buffer[:position])) + suffix = bytes(buffer[trailer : trailer + 2]) + if suffix == b"--": + return + if suffix != b"\r\n": + raise MalformedMultipart( + f"unexpected bytes {suffix!r} after multipart delimiter" + ) + del buffer[: trailer + 2] # consume delimiter + \r\n + + if started: + raise MalformedMultipart("multipart body ended without a closing delimiter") + + +def _parse_part(data: bytes) -> ResponsePart: + """ + Parses a single part from its raw bytes (headers, blank line, body). + """ + header_blob, separator, body = data.partition(b"\r\n\r\n") + if not separator: + raise MalformedMultipart("multipart part has no header separator") + + headers: dict[str, str] = {} + for line in header_blob.split(b"\r\n"): + name, separator, value = line.partition(b":") + if not separator: + continue + utf8_name = name.decode("utf-8", "replace").strip().lower() + utf8_value = value.decode("utf-8", "replace").strip() + headers[utf8_name] = utf8_value + + return ResponsePart(headers=headers, body=body) diff --git a/clients/python/src/objectstore_client/many.py b/clients/python/src/objectstore_client/many.py new file mode 100644 index 00000000..6639c484 --- /dev/null +++ b/clients/python/src/objectstore_client/many.py @@ -0,0 +1,1006 @@ +""" +Batch operations API for executing multiple operations. + +:meth:`Session.many ` takes any iterable +of :class:`Get`, :class:`Put`, :class:`Delete`, and :class:`Head` operations and +executes them with as few requests as possible. Operations within the batch +protocol's per-part size limit are grouped into multipart requests to the +``objects:batch`` endpoint. Oversized and unsized inserts fall back to individual +requests. + +Everything is streamed, both on the send and receive side. +""" + +from __future__ import annotations + +import atexit +import queue +import threading +from collections import Counter +from collections.abc import Generator, Iterable, Iterator, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from io import BytesIO +from typing import IO, TYPE_CHECKING, cast + +import urllib3 +import zstandard +from zstandard import ZstdCompressionReader + +from objectstore_client import _compression +from objectstore_client.client import GetResponse +from objectstore_client.errors import RequestError, raise_for_status +from objectstore_client.formdata import ( + CHUNK_SIZE, + MultipartBody, + PartBody, + RequestPart, + ResponsePart, + iter_multipart, +) +from objectstore_client.metadata import ( + HEADER_EXPIRATION, + HEADER_FILENAME, + HEADER_META_PREFIX, + HEADER_ORIGIN, + Compression, + ExpirationPolicy, + Metadata, + format_expiration, +) +from objectstore_client.metrics import ( + batch_size_bucket, + count_batch_operations, + measure_storage_operation, +) +from objectstore_client.utils import decode_header_value, encode_header_value + +if TYPE_CHECKING: + from objectstore_client.client import Session + +# Maximum number of operations to send in a single batch request. +MAX_BATCH_OPS = 1000 + +# Maximum (post-compression, estimated) size of a single part's body in a batch +# request. Inserts above this are sent as individual requests instead. +MAX_BATCH_PART_SIZE = 1024 * 1024 # 1 MB + +# Maximum total (post-compression, estimated) body size of a single batch request. +MAX_BATCH_BODY_SIZE = 100 * 1024 * 1024 # 100 MB + +# Default maximum number of concurrent requests. +DEFAULT_CONCURRENCY = 3 + +# How long a worker waits on a full result queue before re-checking whether the +# caller has abandoned the results. +_CANCEL_POLL_INTERVAL = 0.2 + +# Reported as the key of an operation that failed before it had one, which can +# only happen for a keyless insert. +_UNKNOWN_KEY = "" + +# Set while the interpreter shuts down, to release a worker waiting to hand over +# a result. Results abandoned without being closed would otherwise leave the +# thread pool's join waiting on workers that nothing is going to drain. +# +# `ThreadPoolExecutor` joins its threads from a `threading._register_atexit` +# callback rather than an `atexit` one, because `atexit` runs after those joins. +# Those callbacks run in reverse registration order, so registering ours here +# (below the import that registers the pool's) releases the workers just before +# the join. +_SHUTTING_DOWN = threading.Event() +_register_before_join = getattr(threading, "_register_atexit", atexit.register) +_register_before_join(_SHUTTING_DOWN.set) + +HEADER_BATCH_OPERATION_KIND = "x-sn-batch-operation-kind" +HEADER_BATCH_OPERATION_KEY = "x-sn-batch-operation-key" +HEADER_BATCH_OPERATION_INDEX = "x-sn-batch-operation-index" +HEADER_BATCH_OPERATION_STATUS = "x-sn-batch-operation-status" + + +@dataclass(frozen=True) +class Get: + """ + Fetches an object. See :meth:`Session.get `. + """ + + key: str + + decompress: bool = True + """Whether to transparently decompress the payload.""" + + accept_encoding: Sequence[str] | None = None + """Compression algorithms to pass through compressed instead of decompressing.""" + + +@dataclass(frozen=True) +class Put: + """ + Uploads an object. See :meth:`Session.put `. + + Stream ``contents`` must be seekable to be batched, as its size is otherwise + unknown; unseekable streams are sent as individual requests. + """ + + contents: bytes | IO[bytes] + key: str | None = None + compress: Compression | None = None + precompressed: Compression | None = None + content_type: str | None = None + metadata: dict[str, str] | None = None + expiration_policy: ExpirationPolicy | None = None + origin: str | None = None + filename: str | None = None + + def __post_init__(self) -> None: + if self.compress is not None and self.precompressed is not None: + raise ValueError("Cannot pass both `compress` and `precompressed`") + _compression.validate(self.compress) + _compression.validate(self.precompressed, allow_none=False) + + +@dataclass(frozen=True) +class Delete: + """ + Deletes an object. + See :meth:`Session.delete `. + """ + + key: str + + +@dataclass(frozen=True) +class Head: + """ + Fetches an object's metadata. + See :meth:`Session.head `. + """ + + key: str + + +Operation = Get | Put | Delete | Head +"""An operation that can be passed to :meth:`Session.many`.""" + + +@dataclass +class GetResult: + """The result of a :class:`Get` operation.""" + + index: int + """Position of the operation in the sequence given to :meth:`Session.many`.""" + + key: str + + response: GetResponse | None + """The fetched object, or ``None`` if it does not exist.""" + + error: Exception | None + """The error that occurred, or ``None`` if the operation succeeded.""" + + +@dataclass +class PutResult: + """The result of a :class:`Put` operation.""" + + index: int + """Position of the operation in the sequence given to :meth:`Session.many`.""" + + key: str + """The object key, as assigned by the server for a keyless :class:`Put`.""" + + error: Exception | None + """The error that occurred, or ``None`` if the operation succeeded.""" + + +@dataclass +class DeleteResult: + """The result of a :class:`Delete` operation.""" + + index: int + """Position of the operation in the sequence given to :meth:`Session.many`.""" + + key: str + + error: Exception | None + """The error that occurred, or ``None`` if the operation succeeded.""" + + +@dataclass +class HeadResult: + """The result of a :class:`Head` operation.""" + + index: int + """Position of the operation in the sequence given to :meth:`Session.many`.""" + + key: str + + metadata: Metadata | None + """The object's metadata, or ``None`` if it does not exist.""" + + error: Exception | None + """The error that occurred, or ``None`` if the operation succeeded.""" + + +@dataclass +class ErrorResult: + """ + An error that cannot be attributed to a specific operation. + """ + + index: None + """ + Always ``None``: an error this class carries belongs to no operation. The + field exists so that every result in a batch can be asked for its index. + """ + + error: Exception + + +OperationResult = GetResult | PutResult | DeleteResult | HeadResult | ErrorResult +"""The result of a single operation in a :meth:`Session.many` call.""" + + +class OperationResults: + """ + A lazy iterator over the results of a :meth:`Session.many` call. + """ + + def __init__(self, results: Generator[OperationResult, None, None]): + self._results = results + + def __iter__(self) -> Iterator[OperationResult]: + return self + + def __next__(self) -> OperationResult: + return next(self._results) + + def __enter__(self) -> OperationResults: + return self + + def __exit__(self, *_exc_info: object) -> None: + self.close() + + def close(self) -> None: + """Abandons all operations that have not completed yet.""" + self._results.close() + + def failures(self) -> list[OperationResult]: + """ + Drains the results and returns those that carry an error. + """ + return [result for result in self if result.error is not None] + + def raise_for_failures(self) -> None: + """ + Drains the results and raises an ``ExceptionGroup`` of all errors, if any. + """ + errors = [result.error for result in self if result.error is not None] + if errors: + raise ExceptionGroup("one or more batch operations failed", errors) + + +def execute_many( + session: Session, + operations: Iterable[Operation], + *, + concurrency: int | None = None, +) -> OperationResults: + """ + Executes ``operations``, batching them where possible. + + This is the implementation of :meth:`Session.many + `; see there for documentation. + """ + concurrency = DEFAULT_CONCURRENCY if concurrency is None else concurrency + if concurrency < 1: + raise ValueError(f"concurrency must be at least 1, got {concurrency}") + return OperationResults(_execute(session, operations, concurrency)) + + +def _classify(session: Session, op: Operation) -> int | None: + """ + Returns the body size an operation contributes to a batch request. + + ``None`` marks it unbatchable. + """ + if isinstance(op, (Get, Delete, Head)): + return 0 + if not isinstance(op, Put): + raise TypeError(f"not an objectstore operation: {op!r}") + + size = _body_size(op.contents) + if size is None: + return None + + # Compression happens while the body is streamed, so its size on the wire is + # only known up to the worst case. Precompressed bodies are sent verbatim. + # Only zstd has a bound here, so gzip bodies go through `Session.put`. + compress_with = _compress_with(session, op) + if compress_with == "zstd": + size = _zstd_compress_bound(size) + elif compress_with != "none": + return None + + return size if size <= MAX_BATCH_PART_SIZE else None + + +def _body_size(contents: bytes | IO[bytes]) -> int | None: + if isinstance(contents, bytes): + return len(contents) + try: + if not contents.seekable(): + return None + position = contents.tell() + end = contents.seek(0, 2) # Seek to end + contents.seek(position) + return end - position + except (OSError, ValueError): + return None + + +def _zstd_compress_bound(size: int) -> int: + """ + Returns the worst-case size of ``size`` bytes compressed as a single zstd frame. + + A port of the ``ZSTD_COMPRESSBOUND`` macro definition in ``zstd.h``. + """ + margin = ((128 << 10) - size) >> 11 if size < (128 << 10) else 0 + return size + (size >> 8) + margin + + +def _encoding(session: Session, op: Put) -> Compression: + """Returns the ``Content-Encoding`` the object is stored with.""" + return op.precompressed or op.compress or session._usecase._compression + + +def _compress_with(session: Session, op: Put) -> Compression: + """Returns the compression the client applies to the body, if any.""" + if op.precompressed is not None: + return "none" + return _encoding(session, op) + + +@dataclass +class _Batch: + """A group of operations sent as a single batch request.""" + + ops: list[tuple[int, Operation]] + + +@dataclass +class _Individual: + """An operation sent as an individual (non-batch) request.""" + + index: int + op: Operation + + +_Work = _Batch | _Individual + + +def _iter_work( + classified: Iterable[tuple[int, Operation, int | None]], +) -> Iterator[_Work]: + """ + Groups classified operations into work items, lazily. + + Each operation travels with the index it had in the sequence the caller gave, + which is what its result reports back. + """ + batch: list[tuple[int, Operation]] = [] + batch_size = 0 + + for index, op, size in classified: + # Ops with unknown size must be sent as individual ops + if size is None: + # If we've been accumulating a batch, cut if off here and yield it + if batch: + yield _Batch(batch) + batch, batch_size = [], 0 + # Now yield this individual op + yield _Individual(index, op) + continue + + # If we've hit a batch size/length limit, cut the batch off here. + if batch and ( + len(batch) >= MAX_BATCH_OPS or batch_size + size > MAX_BATCH_BODY_SIZE + ): + yield _Batch(batch) + batch, batch_size = [], 0 + + # Add to current batch + batch.append((index, op)) + batch_size += size + + # Yield whatever is left over + if batch: + yield _Batch(batch) + + +def _op_key(op: Operation) -> str | None: + if isinstance(op, Put): + return op.key or None + return op.key + + +def _execute( + session: Session, + operations: Iterable[Operation], + concurrency: int, +) -> Generator[OperationResult, None, None]: + concurrency = _fit_to_connection_pool(session, concurrency) + work = _iter_work( + (index, op, _classify(session, op)) for index, op in enumerate(operations) + ) + + if concurrency == 1: + for item in work: + yield from _run_work(session, item) + else: + yield from _execute_concurrent(session, work, concurrency) + + +def _fit_to_connection_pool(session: Session, concurrency: int) -> int: + """ + Returns how many requests to keep in flight, sizing the pool to suit. + + This function reconciles the ``concurrency`` argument that the caller passed + in to ``session.many()`` with the ``maxsize`` and ``block`` options on the + ``urllib3`` connection pool. + + ``maxsize`` controls how many free connections are retained in the pool. + ``block`` controls whether a request must wait for a free connection before + being sent or if it may open a new connection. + + When ``block=False``(default), a connection pool may send more than + ``maxsize`` concurrent requests, but they will not reuse connections + effectively and ``urllib3`` emits warnings about it. So, this function + overwrites ``maxsize`` if the caller's desired ``concurrency`` is higher. + + When ``block=True``, ``maxsize`` is meant to be a hard cap on the number of + concurrent requests. This function leaves ``maxsize`` alone in this case + and clamps the caller's desired ``concurrency`` to be no larger than + ``maxsize``. + """ + pool = getattr(session, "_pool", None) + connections = getattr(pool, "pool", None) + if connections is None: + return concurrency + + if getattr(pool, "block", False): + return max(1, min(concurrency, connections.maxsize)) + + if connections.maxsize < concurrency: + connections.maxsize = concurrency + return concurrency + + +class _Done: + """ + Sentinel value that workers push to the results queue when their work is + finished. + + This signals when every submitted item is accounted for and when a new item + may be submitted. + """ + + +_DONE = _Done() + + +def _execute_concurrent( + session: Session, + work: Iterator[_Work], + concurrency: int, +) -> Iterator[OperationResult]: + """ + Runs work items on a thread pool, yielding results in completion order. + + At most ``concurrency`` items are in flight. Workers hand results over + through a bounded queue. + """ + pool = ThreadPoolExecutor( + max_workers=concurrency, thread_name_prefix="objectstore-many" + ) + results: queue.Queue[OperationResult | _Done] = queue.Queue(maxsize=concurrency * 2) + cancelled = threading.Event() + + def put(item: OperationResult | _Done) -> bool: + """ + Passes a result through the results queue out to the caller. Returns + ``False`` if we were told to exit. + """ + while not (cancelled.is_set() or _SHUTTING_DOWN.is_set()): + try: + results.put(item, timeout=_CANCEL_POLL_INTERVAL) + return True + except queue.Full: + continue + return False + + def run(item: _Work) -> None: + """Work item runner that runs on the thread pool.""" + if cancelled.is_set() or _SHUTTING_DOWN.is_set(): + # The caller abandoned the results between this item's submission + # and a worker picking it up, so it must not be sent. + return + + try: + work_results = _run_work(session, item) + try: + for result in work_results: + if not put(result): + return + finally: + # Releases the connection if the consumer abandoned us + # part-way through the batch response. + work_results.close() + except Exception as error: + put(ErrorResult(None, error)) + finally: + put(_DONE) + + submitted = 0 + finished = 0 + exhausted = False + try: + while True: + # Submit work items until we finish or hit `concurrency`. + while not exhausted and submitted - finished < concurrency: + item = next(work, None) + if item is None: + exhausted = True + break + pool.submit(run, item) + submitted += 1 + + # Normal completion condition has triggered + if exhausted and submitted == finished: + return + + # Pop a result off the queue. If it's `_Done` then the whole work + # item has completed. Otherwise it's an individual operation result. + # + # Polled rather than blocking: a worker releases itself at + # interpreter exit without handing over its `_DONE`, and nothing + # would ever wake this up again. + try: + result = results.get(timeout=_CANCEL_POLL_INTERVAL) + except queue.Empty: + if _SHUTTING_DOWN.is_set(): + return + continue + + if isinstance(result, _Done): + finished += 1 + else: + yield result + finally: + # We've exited the runloop but we haven't hit the normal completion + # case. Set `cancelled.set()` so `pool` workers will give up and we can + # clean everything up. + cancelled.set() + pool.shutdown(wait=False, cancel_futures=True) + + +def _run_work(session: Session, item: _Work) -> Generator[OperationResult, None, None]: + if isinstance(item, _Batch): + yield from _stream_batch(session, item.ops) + else: + yield _execute_individual(session, item.index, item.op) + + +def _execute_individual(session: Session, index: int, op: Operation) -> OperationResult: + """Executes a single operation through the ordinary (non-batch) endpoints.""" + if isinstance(op, Get): + try: + response = session.get( + op.key, + decompress=op.decompress, + accept_encoding=op.accept_encoding, + ) + return GetResult(index, op.key, response, None) + except Exception as error: + return GetResult(index, op.key, None, error) + elif isinstance(op, Put): + try: + key = session.put( + op.contents, + key=op.key, + compress=op.compress, + precompressed=op.precompressed, + content_type=op.content_type, + metadata=op.metadata, + expiration_policy=op.expiration_policy, + origin=op.origin, + filename=op.filename, + ) + return PutResult(index, key, None) + except Exception as error: + return PutResult(index, op.key or _UNKNOWN_KEY, error) + elif isinstance(op, Delete): + try: + session.delete(op.key) + return DeleteResult(index, op.key, None) + except Exception as error: + return DeleteResult(index, op.key, error) + else: + try: + return HeadResult(index, op.key, session.head(op.key), None) + except Exception as error: + return HeadResult(index, op.key, None, error) + + +def _stream_batch( + session: Session, ops: list[tuple[int, Operation]] +) -> Iterator[OperationResult]: + """ + Sends ``ops`` as one batch request, streaming out results as parts arrive. + + Whether the request fails outright or part-way through the response, the + error goes to every operation that has not produced a result yet — so none is + reported twice, and none is dropped. + """ + # Actual parts to be sent in the batch request. Filters out malformed parts. + parts: list[RequestPart] = [] + # Associates each operation with its original index in the batch. + sent: list[tuple[int, Operation]] = [] + for index, op in ops: + try: + request_part = _build_part(session, op) + request_part.validate() + except Exception as invalid: + yield _error_result(index, op, invalid) + continue + parts.append(request_part) + sent.append((index, op)) + + if not sent: + return + + completed = False + seen: set[int] = set() + response: urllib3.BaseHTTPResponse | None = None + try: + body = MultipartBody(parts) + headers = session._make_headers() + headers["Content-Type"] = body.content_type + + usecase = session._usecase.name + # Emit counters for each op we're about to send in this batch + count_batch_operations( + session._metrics_backend, usecase, Counter(_op_kind(op) for _, op in sent) + ) + + # Measures batch processing latency. Note that this includes the time + # the caller spends consuming the results iterator and not just + # client/server latency, since that is what paces reading the response. + with measure_storage_operation( + session._metrics_backend, + "batch", + usecase, + tags={"operations": batch_size_bucket(len(sent))}, + ): + response = session._pool.request( + "POST", + session._make_batch_url(), + body=body, + headers=headers, + preload_content=False, + decode_content=True, + # `None` means defer to the pool, not "no retries" + retries=None if body.resendable else _single_send_retries(session), + ) + raise_for_status(response) + + content_type = response.headers.get("content-type", "") + for part in iter_multipart(content_type, response.stream(CHUNK_SIZE)): + position, result = _result_from_part(sent, part) + if position is not None: + seen.add(position) + yield result + completed = True + + # Yield an error for every operation that didn't have a + # corresponding result part. + yield from _unanswered(sent, seen, None) + except Exception as error: + # Yield an error for every operation that we haven't yielded a + # response part for yet. + yield from _unanswered(sent, seen, error) + finally: + if response is None: + pass # The request never made it out; there is nothing to release. + elif completed: + response.drain_conn() + response.release_conn() + else: + # Not read to the end, so the connection cannot be reused and + # draining it could mean reading an unbounded body. + response.close() + + +def _unanswered( + sent: list[tuple[int, Operation]], seen: set[int], error: Exception | None +) -> Iterator[OperationResult]: + """ + Reports every operation of a batch that no response part accounted for. + """ + for position, (index, op) in enumerate(sent): + if position in seen: + continue + yield _error_result( + index, + op, + error + or RequestError( + f"server did not return a response for the operation at index {index}", + status=0, + response="", + ), + ) + + +def _single_send_retries(session: Session) -> urllib3.Retry: + """ + Returns a policy that lets urllib3 retry a request so long as the retry + doesn't re-send the body. + + A body that reads from a stream can only be sent once, so the only retry we + can try one that happens before the body is touched: a connection failure. + """ + retries = urllib3.Retry.from_int(session._pool.retries) + return retries.new(read=0, other=0, redirect=0) + + +def _op_kind(op: Operation) -> str: + """Returns the batch protocol's name for this kind of operation.""" + if isinstance(op, Get): + return "get" + if isinstance(op, Put): + return "insert" + if isinstance(op, Delete): + return "delete" + return "head" + + +def _build_part(session: Session, op: Operation) -> RequestPart: + if isinstance(op, Put): + return _insert_part(session, op) + return _keyed_part(_op_kind(op), op.key) + + +def _keyed_part(kind: str, key: str) -> RequestPart: + return RequestPart( + headers={ + HEADER_BATCH_OPERATION_KIND: kind, + HEADER_BATCH_OPERATION_KEY: encode_header_value(key), + } + ) + + +def _insert_part(session: Session, op: Put) -> RequestPart: + headers = {HEADER_BATCH_OPERATION_KIND: _op_kind(op)} + key = op.key or None + if key is not None: + headers[HEADER_BATCH_OPERATION_KEY] = encode_header_value(key) + + encoding = _encoding(session, op) + if encoding != "none": + headers["Content-Encoding"] = encoding + if op.content_type: + headers["Content-Type"] = op.content_type + + expiration_policy = op.expiration_policy or session._usecase._expiration_policy + if expiration_policy: + headers[HEADER_EXPIRATION] = format_expiration(expiration_policy) + if op.origin: + headers[HEADER_ORIGIN] = encode_header_value(op.origin) + if op.filename is not None: + headers[HEADER_FILENAME] = encode_header_value(op.filename) + if op.metadata: + for name, value in op.metadata.items(): + headers[f"{HEADER_META_PREFIX}{name}"] = encode_header_value(value) + + return RequestPart(headers, _part_body(op.contents, _compress_with(session, op))) + + +def _part_body(contents: bytes | IO[bytes], compress_with: Compression) -> PartBody: + """ + Returns the part's payload, deferring reads and compression to send time. + + An uncompressed payload is handed over as it is. A compressing one becomes a + :class:`_ZstdBody` that the writer reads as the request goes out. + """ + if compress_with != "zstd": + return contents + + source = BytesIO(contents) if isinstance(contents, bytes) else contents + return cast(IO[bytes], _ZstdBody(source)) + + +class _ZstdBody: + """ + A part payload that compresses a source stream as it is read. File-like, + non-rewindable stream. + + The zstd compressor is somewhat heavy so it is created only once needed and + dropped as soon as the stream is exhausted. + """ + + def __init__(self, original_stream: IO[bytes]): + # Only set until `self._compressed_stream` is lazily created + self._original_stream: IO[bytes] | None = original_stream + # Once set, unsets `self._original_stream` + self._compressed_stream: ZstdCompressionReader | None = None + self._closed = False + + @property + def _stream(self) -> ZstdCompressionReader | None: + """ + The compressor, created on first access. ``None`` once exhausted. + """ + if self._original_stream is not None: + compressor = zstandard.ZstdCompressor() + self._compressed_stream = compressor.stream_reader( + self._original_stream, + closefd=False, # Not our stream to close + ) + self._original_stream = None + return self._compressed_stream + + @_stream.setter + def _stream(self, value: None) -> None: + self._original_stream = None + self._compressed_stream = None + + def read(self, size: int = -1, /) -> bytes: + if self._closed: + raise ValueError("read of closed stream") + + stream = self._stream + # Not closed, but source is exhausted. EOF + if stream is None: + return b"" + + chunk: bytes = stream.read(size) + # When the stream is exhausted, drop the compressor + if not chunk: + self._stream = None + + return chunk + + def close(self) -> None: + # `_compressed_stream` and not `_stream`: reading the latter would build + # a compressor for a body that was never read, only to close it again. + if self._compressed_stream is not None: + self._compressed_stream.close() + self._closed = True + self._stream = None + + @property + def closed(self) -> bool: + return self._closed + + +def _result_from_part( + sent: list[tuple[int, Operation]], part: ResponsePart +) -> tuple[int | None, OperationResult]: + """ + Turns one response part into a result for the operation it belongs to. + + Returns the operation's position within the batch (server-side, so not + including malformed operations that were dropped before sending) or ``None`` + when the part can't be linked to an operation. + """ + raw_position = part.headers.get(HEADER_BATCH_OPERATION_INDEX) + if raw_position is None or not raw_position.isdigit(): + return None, ErrorResult( + None, + _malformed(f"missing or invalid {HEADER_BATCH_OPERATION_INDEX} header"), + ) + position = int(raw_position) + + if position >= len(sent): + return None, ErrorResult( + None, _malformed(f"response references unknown operation index {position}") + ) + index, op = sent[position] + + try: + return position, _interpret_part(index, op, part) + except Exception as error: + return position, _error_result(index, op, error) + + +def _interpret_part(index: int, op: Operation, part: ResponsePart) -> OperationResult: + """Turns the response part for a known operation into that operation's result.""" + status = _parse_status(part.headers.get(HEADER_BATCH_OPERATION_STATUS)) + if status is None: + return _error_result( + index, + op, + _malformed(f"missing or invalid {HEADER_BATCH_OPERATION_STATUS} header"), + ) + + # Prefer the server's key as the source of truth. + raw_key = part.headers.get(HEADER_BATCH_OPERATION_KEY) + key = decode_header_value(raw_key) if raw_key else _op_key(op) + + # A missing object is a successful "not found" for reads, an error otherwise. + is_error = status >= 400 and not (isinstance(op, (Get, Head)) and status == 404) + + if key is None: + # Any successful operation must have a key, so this should never happen. + if not is_error: + return _error_result( + index, op, _malformed(f"missing {HEADER_BATCH_OPERATION_KEY} header") + ) + key = _UNKNOWN_KEY + + if is_error: + error = RequestError( + f"batch operation failed with status {status}", + status=status, + response=part.body.decode("utf-8", "replace"), + ) + return _error_result(index, op, error, key=key) + + if isinstance(op, Get): + if status == 404: + return GetResult(index, key, None, None) + metadata = Metadata.from_headers(part.headers) + payload = _maybe_decompress(part.body, metadata, op) + return GetResult(index, key, GetResponse(metadata, payload), None) + if isinstance(op, Put): + return PutResult(index, key, None) + if isinstance(op, Delete): + return DeleteResult(index, key, None) + if status == 404: + return HeadResult(index, key, None, None) + return HeadResult(index, key, Metadata.from_headers(part.headers), None) + + +def _maybe_decompress(payload: bytes, metadata: Metadata, op: Get) -> IO[bytes]: + """Applies transparent decompression, mirroring :meth:`Session.get`.""" + accept_encoding = op.accept_encoding + encoding_accepted = accept_encoding is not None and ( + "*" in accept_encoding or metadata.compression in accept_encoding + ) + if metadata.compression and op.decompress and not encoding_accepted: + stream = _compression.decompressing_reader( + BytesIO(payload), metadata.compression + ) + metadata.compression = None + return stream + return BytesIO(payload) + + +def _parse_status(value: str | None) -> int | None: + """Parses a status header like ``"200 OK"`` into its numeric code.""" + if not value: + return None + code = value.split(" ", 1)[0] + try: + return int(code) + except ValueError: + return None + + +def _malformed(message: str) -> RequestError: + return RequestError(f"malformed batch response: {message}", status=0, response="") + + +def _error_result( + index: int, op: Operation, error: Exception, *, key: str | None = None +) -> OperationResult: + """Creates a result of the right kind for ``op``, carrying ``error``.""" + key = key or _op_key(op) or _UNKNOWN_KEY + if isinstance(op, Get): + return GetResult(index, key, None, error) + if isinstance(op, Put): + return PutResult(index, key, error) + if isinstance(op, Delete): + return DeleteResult(index, key, error) + return HeadResult(index, key, None, error) diff --git a/clients/python/src/objectstore_client/metrics.py b/clients/python/src/objectstore_client/metrics.py index 6c7e59d0..5289e1e3 100644 --- a/clients/python/src/objectstore_client/metrics.py +++ b/clients/python/src/objectstore_client/metrics.py @@ -75,10 +75,17 @@ def distribution( class StorageMetricEmitter: - def __init__(self, backend: MetricsBackend, operation: str, usecase: str): + def __init__( + self, + backend: MetricsBackend, + operation: str, + usecase: str, + tags: Tags | None = None, + ): self.backend = backend self.operation = operation - self.usecase = usecase + self.tags: Tags = {"usecase": usecase, **(tags or {})} + """Tags common to every metric of this operation.""" # These may be set during or after the enclosed operation self.start: int | None = None @@ -89,28 +96,28 @@ def __init__(self, backend: MetricsBackend, operation: str, usecase: str): self.compression: str = "unknown" def record_latency(self, elapsed: float) -> None: - tags = {"usecase": self.usecase} + tags = dict(self.tags) self.backend.distribution( f"storage.{self.operation}.latency", elapsed, tags=tags ) self.elapsed = elapsed def record_uncompressed_size(self, value: int) -> None: - tags = {"usecase": self.usecase, "compression": "none"} + tags = {**self.tags, "compression": "none"} self.backend.distribution( f"storage.{self.operation}.size", value, tags=tags, unit="byte" ) self.uncompressed_size = value def record_size(self, value: int) -> None: - tags = {"usecase": self.usecase} + tags = dict(self.tags) self.backend.distribution( f"storage.{self.operation}.size", value, tags=tags, unit="byte" ) self.size = value def record_compressed_size(self, value: int, compression: str = "unknown") -> None: - tags = {"usecase": self.usecase, "compression": compression} + tags = {**self.tags, "compression": compression} self.backend.distribution( f"storage.{self.operation}.size", value, tags=tags, unit="byte" ) @@ -121,7 +128,7 @@ def maybe_record_compression_ratio(self) -> None: if not self.uncompressed_size or not self.compressed_size: return None - tags = {"usecase": self.usecase, "compression": self.compression} + tags = {**self.tags, "compression": self.compression} self.backend.distribution( f"storage.{self.operation}.compression_ratio", self.compressed_size / self.uncompressed_size, @@ -141,7 +148,7 @@ def maybe_record_throughputs(self) -> None: sizes.append((self.compressed_size, self.compression)) for size, compression in sizes: - tags: dict[str, str] = {"usecase": self.usecase} + tags: dict[str, str] = dict(self.tags) if compression is not None: tags["compression"] = compression self.backend.distribution( @@ -154,6 +161,31 @@ def maybe_record_throughputs(self) -> None: ) +def batch_size_bucket(count: int) -> str: + """ + Buckets a count of batch operations into buckets between powers of 2. Used + to tag batch latency metrics with a rough indication of the batch size. + """ + if count <= 1: + return str(count) + low = 1 << (count.bit_length() - 1) + return f"{low}-{2 * low - 1}" + + +def count_batch_operations( + backend: MetricsBackend, usecase: str, kinds: Mapping[str, int] +) -> None: + """ + Counts the operations of one batch request, given as a count per kind. + """ + for kind, count in kinds.items(): + backend.increment( + "storage.batch.operations", + count, + tags={"usecase": usecase, "operation": kind}, + ) + + @contextmanager def measure_storage_operation( backend: MetricsBackend, @@ -162,6 +194,7 @@ def measure_storage_operation( uncompressed_size: int | None = None, compressed_size: int | None = None, compression: str = "unknown", + tags: Tags | None = None, ) -> Generator[StorageMetricEmitter]: """ Context manager which records the latency of the enclosed storage operation. @@ -170,8 +203,10 @@ def measure_storage_operation( Yields a `StorageMetricEmitter` because for some operations (GET) the size is not known until the inside of the enclosed block. + + Any `tags` are added to every metric recorded for this operation. """ - emitter = StorageMetricEmitter(backend, operation, usecase) + emitter = StorageMetricEmitter(backend, operation, usecase, tags) if uncompressed_size: emitter.record_uncompressed_size(uncompressed_size) diff --git a/clients/python/tests/test_e2e.py b/clients/python/tests/test_e2e.py index 3b9db801..e7efaf79 100644 --- a/clients/python/tests/test_e2e.py +++ b/clients/python/tests/test_e2e.py @@ -1,4 +1,5 @@ import gzip +import logging import os import shutil import signal @@ -10,7 +11,7 @@ import urllib.parse import urllib.request from collections.abc import Generator -from datetime import timedelta +from datetime import UTC, datetime, timedelta from io import BytesIO from pathlib import Path from typing import cast @@ -35,6 +36,18 @@ ) from objectstore_client.auth import Permission, SecretKey, TokenGenerator from objectstore_client.errors import RequestError +from objectstore_client.many import ( + DEFAULT_CONCURRENCY, + Delete, + DeleteResult, + Get, + GetResult, + Head, + HeadResult, + Operation, + Put, + PutResult, +) from objectstore_client.metadata import Compression, TimeToLive from objectstore_client.scope import Scope @@ -60,6 +73,13 @@ def tell(self) -> int: raise OSError("stream does not expose a stable position") +class UnseekableStream(BytesIO): + """Read-only stream whose size cannot be determined without reading it.""" + + def seekable(self) -> bool: + return False + + class TestTokenGenerator: _instance: TokenGenerator | None = None @@ -682,6 +702,187 @@ def test_connect_timeout() -> None: session.put(b"test data", compress="zstd") +def _many_session(server_url: str, **client_kwargs: object) -> Session: + client = Client(server_url, token=TestTokenGenerator.get(), **client_kwargs) # type: ignore[arg-type] + usecase = Usecase("test-usecase", expiration_policy=TimeToLive(timedelta(days=1))) + return client.session(usecase, org=12345, project=1337) + + +@pytest.mark.parametrize("concurrency", [1, 3]) +def test_many_operations(server_url: str, concurrency: int) -> None: + session = _many_session(server_url) + + session.many( + [ + Put(b"first object", key="many-1", compress="none", filename="report.pdf"), + Put(b"second object", key="many-2", compress="zstd"), + Put(b"third object", key="many-3", metadata={"foo": "bar"}), + Put(b"fourth object", key="many-4", compress="none"), + ], + concurrency=concurrency, + ).raise_for_failures() + + results = list( + session.many( + [Get("many-1"), Get("many-2"), Head("many-3"), Delete("many-4")], + concurrency=concurrency, + ) + ) + assert len(results) == 4 + by_key = {result.key: result for result in results} # type: ignore[union-attr] + + first = by_key["many-1"] + assert isinstance(first, GetResult) and first.response is not None + assert first.response.metadata.compression is None + assert first.response.metadata.filename == "report.pdf" + assert first.response.payload.read() == b"first object" + + # The usecase's one-day TTL reached the server and came back applied. Check + # that it matches the expected TTL within some epsilon. + expires = first.response.metadata.time_expires + assert expires is not None + assert abs(expires - (datetime.now(UTC) + timedelta(days=1))) < timedelta(minutes=5) + + second = by_key["many-2"] + assert isinstance(second, GetResult) and second.response is not None + # Transparently decompressed on the way out. + assert second.response.metadata.compression is None + assert second.response.payload.read() == b"second object" + + third = by_key["many-3"] + assert isinstance(third, HeadResult) and third.metadata is not None + assert third.metadata.custom == {"foo": "bar"} + + fourth = by_key["many-4"] + assert isinstance(fourth, DeleteResult) and fourth.error is None + assert session.head("many-4") is None + + +def test_many_reports_operation_index(server_url: str) -> None: + session = _many_session(server_url) + keys = [f"many-index-{index}" for index in range(6)] + + session.many( + [Put(f"object {key}".encode(), key=key) for key in keys] + ).raise_for_failures() + + results = list(session.many([Get(key) for key in keys])) + assert {result.index: result.key for result in results} == dict( # type: ignore[union-attr] + enumerate(keys) + ) + + +def test_many_insert_without_key(server_url: str) -> None: + session = _many_session(server_url) + + (result,) = list(session.many([Put(b"keyless object", compress="none")])) + assert isinstance(result, PutResult) + assert result.error is None + assert result.key + + retrieved = session.get(result.key) + assert retrieved is not None + assert retrieved.payload.read() == b"keyless object" + + +def test_many_partial_failures(server_url: str) -> None: + _many_session(server_url).many( + [Put(b"readable", key="many-readable", compress="none")] + ).raise_for_failures() + + # A read-only token: writes and deletes fail with 403, reads 404 on misses. + client = Client( + server_url, + token=TestTokenGenerator.create(permissions=[Permission.OBJECT_READ]), + ) + usecase = Usecase("test-usecase", expiration_policy=TimeToLive(timedelta(days=1))) + session = client.session(usecase, org=12345, project=1337) + + results = list( + session.many( + [ + Get("many-readable"), + Put(b"should fail", key="many-write-key", compress="none"), + Delete("many-delete-key"), + Get("many-nonexistent"), + ] + ) + ) + assert len(results) == 4 + by_key = {result.key: result for result in results} # type: ignore[union-attr] + + # A failed operation keeps the result type of the operation it came from, + # and carries the server's error body. + put = by_key["many-write-key"] + assert isinstance(put, PutResult) + assert isinstance(put.error, RequestError) and put.error.status == 403 + assert put.error.response is not None and "not allowed" in put.error.response + + delete = by_key["many-delete-key"] + assert isinstance(delete, DeleteResult) + assert isinstance(delete.error, RequestError) and delete.error.status == 403 + + # The failures leave the reads alone, and a missing object is a successful + # "not found" rather than an error. + readable = by_key["many-readable"] + assert isinstance(readable, GetResult) and readable.error is None + assert readable.response is not None + assert readable.response.payload.read() == b"readable" + + missing = by_key["many-nonexistent"] + assert isinstance(missing, GetResult) + assert missing.error is None and missing.response is None + + +def test_many_round_trips_every_insert_body_shape(server_url: str) -> None: + session = _many_session(server_url) + payload = b"streamed object " * 1000 + big = b"x" * (2 * 1024 * 1024) + + session.many( + [ + Put(big, key="many-big", compress="none"), + Put(b"small", key="many-small"), + Put(BytesIO(payload), key="many-stream", compress="zstd"), + Put(UnseekableStream(payload), key="many-unseekable"), + ] + ).raise_for_failures() + + for key, expected in ( + ("many-big", big), + ("many-small", b"small"), + ("many-stream", payload), + ("many-unseekable", payload), + ): + retrieved = session.get(key) + assert retrieved is not None, key + assert retrieved.payload.read() == expected, key + + +@pytest.mark.parametrize("concurrency", [None, 6]) +def test_many_pools_connections_for_concurrent_requests( + server_url: str, caplog: pytest.LogCaptureFixture, concurrency: int | None +) -> None: + session = _many_session(server_url) + keys = [f"many-pooled-{concurrency}-{index}" for index in range(8)] + + # Unsized inserts are sent individually, so this is one request per key. + inserts: list[Operation] = [ + Put(UnseekableStream(b"payload"), key=key) for key in keys + ] + with caplog.at_level(logging.WARNING, logger="urllib3"): + session.many(inserts, concurrency=concurrency).raise_for_failures() + + # A pool that cannot hold the connections logs one warning per discard + noise = [r for r in caplog.records if "pool is full" in r.getMessage()] + assert noise == [] + + pool = session._pool + assert pool.num_requests == len(keys) + in_flight = DEFAULT_CONCURRENCY if concurrency is None else concurrency + assert pool.num_connections <= in_flight, "connections were reopened, not reused" + + def _fetch(url: str, method: str = "GET") -> tuple[int, bytes]: """Fetches ``url`` with plain urllib (no auth header), returning (status, body). diff --git a/clients/python/tests/test_formdata.py b/clients/python/tests/test_formdata.py new file mode 100644 index 00000000..84cf5af5 --- /dev/null +++ b/clients/python/tests/test_formdata.py @@ -0,0 +1,191 @@ +"""Tests for the streaming ``multipart/form-data`` codec used by batch requests.""" + +from __future__ import annotations + +import io +from collections.abc import Iterable + +import pytest +from objectstore_client.formdata import ( + MalformedMultipart, + MultipartBody, + RequestPart, + ResponsePart, + _extract_boundary, + iter_multipart, +) + + +def _serialize(parts: Iterable[RequestPart]) -> tuple[str, bytes]: + body = MultipartBody(list(parts)) + return body.content_type, b"".join(body) + + +def _parse(content_type: str, body: bytes) -> list[ResponsePart]: + return list(iter_multipart(content_type, [body])) + + +def test_encodes_part_headers_and_body() -> None: + content_type, body = _serialize([RequestPart({"x-custom": "value"}, b"hello")]) + + boundary = _extract_boundary(content_type) + assert body.startswith(f"--{boundary}\r\n".encode()) + assert b"content-disposition: form-data; name=part\r\n" in body + assert b"x-custom: value\r\n" in body + assert b"\r\n\r\nhello\r\n" in body + assert body.endswith(f"--{boundary}--\r\n".encode()) + + +def test_reads_lazily() -> None: + class Recording(io.BytesIO): + reads = 0 + + def read(self, size: int | None = -1, /) -> bytes: + Recording.reads += 1 + return super().read(size) + + body = MultipartBody([RequestPart({}, Recording(b"lazy body"))]) + assert Recording.reads == 0 + + assert b"lazy body" in b"".join(body) + assert Recording.reads > 0 + + +def test_bytes_body_can_be_resent() -> None: + """urllib3 re-iterates the body when it retries a connection.""" + body = MultipartBody([RequestPart({"x-kind": "insert"}, b"payload")]) + + assert b"".join(body) == b"".join(body) + + +def test_leaves_a_stream_body_open() -> None: + stream = io.BytesIO(b"payload") + + assert b"payload" in b"".join(MultipartBody([RequestPart({}, stream)])) + + # Releasing the stream is the caller's business, not the encoder's. + assert not stream.closed + + +def test_reports_whether_a_body_can_be_resent() -> None: + # Resendable as it only contains a single non-streamed payload + assert MultipartBody([RequestPart({}, b"payload")]).resendable + # Not resendable; contains a stream that can't be rewound + assert not MultipartBody([RequestPart({}, io.BytesIO(b"payload"))]).resendable + # Not resendable; contains a stream that can't be rewound + assert not MultipartBody( + [RequestPart({}, b"payload"), RequestPart({}, io.BytesIO(b"payload"))] + ).resendable + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("x-header", "value\r\nx-sn-batch-operation-kind: delete"), + ("x-header", "value\nmore"), + ("x-header", "control\x00char"), + ("x-header", "nön-ascii"), + ("x-head\r\ner", "value"), + ("x-header:", "value"), + ("x-head er", "value"), + ("x-header", " leading space"), + ("x-header", "trailing space "), + ("x-header", "\t"), + ("x-hea(d)er", "value"), + ("nön-ascii", "value"), + ("", "value"), + ], +) +def test_rejects_unrepresentable_headers(name: str, value: str) -> None: + with pytest.raises(ValueError): + MultipartBody([RequestPart({name: value})]) + + with pytest.raises(ValueError): + RequestPart({name: value}).validate() + + +def test_extract_boundary() -> None: + assert _extract_boundary('multipart/form-data; boundary="abc"') == "abc" + assert _extract_boundary("multipart/form-data; boundary=abc") == "abc" + with pytest.raises(MalformedMultipart): + _extract_boundary("application/json") + + +def test_parses_encoded_body() -> None: + content_type, body = _serialize( + [ + RequestPart({"x-index": "0"}, b"first"), + RequestPart({"x-index": "1"}), + RequestPart({"x-index": "2"}, b"\x00\xff binary \r\n body"), + ] + ) + + parts = _parse(content_type, body) + + assert [part.headers["x-index"] for part in parts] == ["0", "1", "2"] + assert [part.body for part in parts] == [ + b"first", + b"", + b"\x00\xff binary \r\n body", + ] + assert all(part.headers["content-disposition"] for part in parts) + + +def test_parses_headers_without_space_after_colon() -> None: + body = b"--b\r\nx-index:0\r\n\r\nbody\r\n--b--\r\n" + (part,) = _parse("multipart/form-data; boundary=b", body) + assert part.headers == {"x-index": "0"} + assert part.body == b"body" + + +def test_skips_preamble() -> None: + body = b"ignored preamble\r\n--b\r\nx-index: 0\r\n\r\nbody\r\n--b--\r\n" + (part,) = _parse("multipart/form-data; boundary=b", body) + assert part.body == b"body" + + +def test_yields_parts_before_the_body_ends() -> None: + content_type, body = _serialize( + [RequestPart({"x-index": "0"}, b"first"), RequestPart({"x-index": "1"})] + ) + # Feeding the body one byte at a time exercises boundary detection across + # chunks, and shows that a part is yielded as soon as it is complete. + parts = iter_multipart(content_type, (body[i : i + 1] for i in range(len(body)))) + + first = next(parts) + assert first.body == b"first" + assert next(parts).headers["x-index"] == "1" + with pytest.raises(StopIteration): + next(parts) + + +def test_rejects_truncated_body() -> None: + content_type, body = _serialize([RequestPart({"x-index": "0"}, b"first")]) + with pytest.raises(MalformedMultipart): + _parse(content_type, body[: len(body) // 2]) + + +def test_rejects_unexpected_bytes_after_delimiter() -> None: + body = b"--b\r\nx-index: 0\r\n\r\nbody\r\n--bxx\r\n" + with pytest.raises(MalformedMultipart): + _parse("multipart/form-data; boundary=b", body) + + +def test_skips_a_header_line_it_cannot_split() -> None: + body = b"--b\r\nx-index: 0\r\ngarbage\r\nx-status: 200 OK\r\n\r\nbody\r\n--b--\r\n" + + (part,) = _parse("multipart/form-data; boundary=b", body) + + # The line is dropped, the part it belongs to is not. + assert part.headers == {"x-index": "0", "x-status": "200 OK"} + assert part.body == b"body" + + +def test_rejects_part_without_headers() -> None: + body = b"--b\r\nbroken\r\n--b--\r\n" + with pytest.raises(MalformedMultipart): + _parse("multipart/form-data; boundary=b", body) + + +def test_parses_empty_body() -> None: + assert _parse("multipart/form-data; boundary=b", b"") == [] diff --git a/clients/python/tests/test_many.py b/clients/python/tests/test_many.py new file mode 100644 index 00000000..c28ef7ee --- /dev/null +++ b/clients/python/tests/test_many.py @@ -0,0 +1,1185 @@ +from __future__ import annotations + +import io +import json +import pathlib +import subprocess +import sys +import threading +import time +from collections.abc import Iterator +from contextlib import contextmanager +from datetime import timedelta +from typing import Any + +import pytest +import sentry_sdk +import urllib3 +import zstandard +from objectstore_client import Client, Session, Usecase +from objectstore_client import many as many_module +from objectstore_client.formdata import ( + ResponsePart, + iter_multipart, +) +from objectstore_client.many import ( + MAX_BATCH_BODY_SIZE, + MAX_BATCH_OPS, + MAX_BATCH_PART_SIZE, + Delete, + DeleteResult, + ErrorResult, + Get, + GetResult, + Head, + HeadResult, + Operation, + OperationResult, + Put, + PutResult, + _classify, + _fit_to_connection_pool, + _iter_work, + _parse_status, + _zstd_compress_bound, + _ZstdBody, +) +from objectstore_client.metadata import TimeToLive +from objectstore_client.metrics import MetricsBackend, Tags, batch_size_bucket +from objectstore_client.utils import decode_header_value, encode_header_value + +_STATUS_REASONS = { + 200: "OK", + 201: "Created", + 204: "No Content", + 403: "Forbidden", + 404: "Not Found", +} + +# Headers of the batch protocol itself, which are not object metadata. +_PROTOCOL_HEADERS = { + "x-sn-batch-operation-kind", + "x-sn-batch-operation-key", + "x-sn-batch-operation-index", + "x-sn-batch-operation-status", + "content-disposition", +} + + +def _status_line(code: int) -> str: + return f"{code} {_STATUS_REASONS[code]}" + + +class Unseekable(io.BytesIO): + """A stream whose size cannot be known without reading it.""" + + def seekable(self) -> bool: + return False + + +class FakeResponse: + """A stand-in for a urllib3 response, streaming its body in small chunks.""" + + def __init__(self, status: int, headers: dict[str, str], data: bytes = b""): + self.status = status + self.headers = headers + self.data = data + self.drained = False + self.released = False + self.closed = False + self._buffer = io.BytesIO(data) + + def read(self, amt: int | None = None) -> bytes: + return self._buffer.read(amt) + + def stream(self, _amt: int | None = None) -> Iterator[bytes]: + # A deliberately tiny chunk size, so that parts and boundaries land + # across chunks the way they do on a real socket. + while chunk := self._buffer.read(13): + yield chunk + + def json(self) -> Any: + return json.loads(self.data) + + def drain_conn(self) -> None: + self.drained = True + + def release_conn(self) -> None: + self.released = True + + def close(self) -> None: + self.closed = True + + +class FakeObjectstore: + """ + An in-memory stand-in for the endpoints the many API talks to. + + Requests are parsed with the formdata codec's ``iter_multipart`` response + parser as the grammar is the same for requests and responses. + + This class shouldn't be used to write integration tests that rely on real + server behavior. Those tests should be written in ``test_e2e.py``. What it + offers instead is more granular visibilty into request/connection details + and, via subclasses, behavior that should be impossible from a real server. + + ``FakeObjectstore`` imitates the parts of a ``urllib3`` connection pool that + the client uses, as that's how it's shimmed into our test Objectstore + clients. Tests subclass it and override methods to do whatever is needed for + the test. + """ + + retries = urllib3.Retry(connect=3, read=0) + + def __init__(self, *, deny_writes: bool = False, in_order: bool = False): + self.headers: dict[str, str] = {} + self.store: dict[str, tuple[dict[str, str], bytes]] = {} + self.deny_writes = deny_writes + self.in_order = in_order + """Whether to respond in request order, rather than the reverse.""" + + self.batches: list[list[ResponsePart]] = [] + """The parsed request parts of every batch request received.""" + + self.bodies: list[Any] = [] + """The raw body object of every batch request received.""" + + self.individual: list[str] = [] + """The keys of every insert received on the individual endpoint.""" + + self.responses: list[FakeResponse] = [] + """Every response handed out, to inspect how connections were released.""" + + self.retries_used: list[Any] = [] + """The retry policy of every request, ``None`` when the pool's applies.""" + + self._generated = 0 + + @property + def requests(self) -> int: + return len(self.batches) + len(self.individual) + + def request( + self, + method: str, + url: str, + *, + body: Any = None, + headers: dict[str, str] | None = None, + **_kwargs: Any, + ) -> FakeResponse: + self.retries_used.append(_kwargs.get("retries")) + if "objects:batch" in url: + assert method == "POST" + response = self._batch(body, headers or {}) + else: + response = self._individual_insert(method, url, body, headers or {}) + self.responses.append(response) + return response + + def _individual_insert( + self, method: str, url: str, body: Any, headers: dict[str, str] + ) -> FakeResponse: + assert method in ("POST", "PUT") + key = url.rsplit("/", 1)[-1] or self._next_key() + contents = body.read() if hasattr(body, "read") else bytes(body or b"") + self.store[key] = (_metadata_headers(headers), contents) + self.individual.append(key) + return FakeResponse(200, {}, json.dumps({"key": key}).encode()) + + def _batch(self, body: Any, headers: dict[str, str]) -> FakeResponse: + self.bodies.append(body) + raw = body if isinstance(body, bytes) else b"".join(body) + parts = list(iter_multipart(headers["Content-Type"], [raw])) + self.batches.append(parts) + + responses = [self._execute(index, part) for index, part in enumerate(parts)] + if not self.in_order: + # The server executes the parts of a batch concurrently and responds + # in completion order, which is not the order they were sent in. + responses.reverse() + + boundary = "os-boundary-" + f"{len(self.batches):032x}" + return FakeResponse( + 200, + {"content-type": f'multipart/form-data; boundary="{boundary}"'}, + _serialize_response(responses, boundary), + ) + + def _execute(self, index: int, part: ResponsePart) -> ResponsePart: + kind = part.headers["x-sn-batch-operation-kind"] + raw_key = part.headers.get("x-sn-batch-operation-key") + key = decode_header_value(raw_key) if raw_key else None + + headers = { + "content-disposition": "form-data; name=part", + "x-sn-batch-operation-index": str(index), + "x-sn-batch-operation-kind": kind, + } + + def respond(status: int, key: str | None, body: bytes = b"") -> ResponsePart: + if key is not None: + headers["x-sn-batch-operation-key"] = encode_header_value(key) + headers["x-sn-batch-operation-status"] = _status_line(status) + return ResponsePart(headers, body) + + if kind == "insert": + assert len(part.body) <= MAX_BATCH_PART_SIZE, "part exceeds server limit" + if self.deny_writes: + return respond(403, key, b'{"detail":"forbidden"}') + key = key or self._next_key() + self.store[key] = (_metadata_headers(part.headers), part.body) + return respond(201, key) + + assert key is not None, f"{kind} operation without a key" + + if kind == "delete": + if self.deny_writes: + return respond(403, key, b'{"detail":"forbidden"}') + self.store.pop(key, None) + return respond(204, key) + + if key not in self.store: + return respond(404, key) + metadata, contents = self.store[key] + + if kind == "get": + headers.update(metadata) + headers.setdefault("content-type", "application/octet-stream") + headers["x-sn-time-created"] = "2024-01-01T00:00:00+00:00" + return respond(200, key, contents) + + assert kind == "head", f"unknown operation kind {kind}" + headers.update(metadata) + headers["x-sn-time-created"] = "2024-01-01T00:00:00+00:00" + headers["x-sn-size"] = str(len(contents)) + return respond(200, key) + + def _next_key(self) -> str: + self._generated += 1 + return f"generated-{self._generated}" + + +def _metadata_headers(headers: dict[str, str]) -> dict[str, str]: + """Extracts the headers the server would persist as object metadata.""" + return { + name.lower(): value + for name, value in headers.items() + if name.lower() not in _PROTOCOL_HEADERS + and (name.lower().startswith(("content-", "x-sn-", "x-snme-"))) + } + + +def _serialize_response(parts: list[ResponsePart], boundary: str) -> bytes: + """Serializes parts exactly the way the server's multipart writer does.""" + out = bytearray() + for part in parts: + out += f"--{boundary}\r\n".encode() + for name, value in part.headers.items(): + out += f"{name}: {value}\r\n".encode() + out += b"\r\n" + part.body + b"\r\n" + out += f"--{boundary}--".encode() + return bytes(out) + + +def _respond_with(parts: list[ResponsePart]) -> FakeResponse: + """A batch response carrying exactly ``parts``.""" + boundary = "os-boundary-" + "0" * 32 + return FakeResponse( + 200, + {"content-type": f'multipart/form-data; boundary="{boundary}"'}, + _serialize_response(parts, boundary), + ) + + +class RecordingMetrics(MetricsBackend): + """Records every metric the client emits.""" + + def __init__(self) -> None: + self.counters: list[tuple[str, int | float, Tags | None]] = [] + self.distributions: list[tuple[str, int | float, Tags | None]] = [] + + def increment( + self, name: str, value: int | float = 1, tags: Tags | None = None + ) -> None: + self.counters.append((name, value, tags)) + + def gauge(self, name: str, value: int | float, tags: Tags | None = None) -> None: + pass + + def distribution( + self, + name: str, + value: int | float, + tags: Tags | None = None, + unit: str | None = None, + ) -> None: + self.distributions.append((name, value, tags)) + + +def _session( + pool: Any, compression: str = "none", metrics: MetricsBackend | None = None +) -> Session: + client = Client("http://localhost:8888", metrics_backend=metrics) + client._pool = pool + usecase = Usecase( + "testing", + compression=compression, # type: ignore[arg-type] + expiration_policy=TimeToLive(timedelta(days=1)), + ) + return client.session(usecase, org=42, project=1337) + + +def _by_key(results: list[OperationResult]) -> dict[str, OperationResult]: + return { + result.key: result for result in results if not isinstance(result, ErrorResult) + } + + +@pytest.mark.parametrize("concurrency", [1, 3]) +def test_many_round_trip(concurrency: int) -> None: + pool = FakeObjectstore() + session = _session(pool) + + results = list( + session.many( + [ + Put(b"first", key="key-1", filename="report.pdf"), + Put(b"second", key="key-2", compress="zstd"), + Put(b"third", key="key-3", metadata={"foo": "bar"}), + Put(b"fourth", key="key-4"), + ], + concurrency=concurrency, + ) + ) + puts = [result for result in results if isinstance(result, PutResult)] + assert len(puts) == len(results) + assert sorted(put.key for put in puts) == ["key-1", "key-2", "key-3", "key-4"] + assert pool.requests == 1 + + results = list( + session.many( + [Get("key-1"), Get("key-2"), Head("key-3"), Delete("key-4"), Get("gone")], + concurrency=concurrency, + ) + ) + assert len(results) == 5 + by_key = _by_key(results) + + first = by_key["key-1"] + assert isinstance(first, GetResult) and first.response is not None + assert first.response.payload.read() == b"first" + assert first.response.metadata.filename == "report.pdf" + + second = by_key["key-2"] + assert isinstance(second, GetResult) and second.response is not None + # Transparently decompressed, so the metadata no longer claims compression. + assert second.response.metadata.compression is None + assert second.response.payload.read() == b"second" + + third = by_key["key-3"] + assert isinstance(third, HeadResult) and third.metadata is not None + assert third.metadata.custom == {"foo": "bar"} + + assert isinstance(by_key["key-4"], DeleteResult) + assert by_key["key-4"].error is None + assert "key-4" not in pool.store + + gone = by_key["gone"] + assert isinstance(gone, GetResult) + assert gone.response is None and gone.error is None + + # A missing object is a successful "not found" for reads, not a failure. + (missing,) = list(session.many([Head("gone")], concurrency=concurrency)) + assert isinstance(missing, HeadResult) + assert missing.metadata is None and missing.error is None + + +@pytest.mark.parametrize( + "get", [Get("k", decompress=False), Get("k", accept_encoding=["zstd"])] +) +def test_many_optionally_skips_decompression(get: Get) -> None: + """Both ways of opting out leave the payload and its encoding as stored.""" + pool = FakeObjectstore() + session = _session(pool) + session.many([Put(b"payload", key="k", compress="zstd")]).raise_for_failures() + + (result,) = list(session.many([get])) + + assert isinstance(result, GetResult) and result.response is not None + assert result.response.metadata.compression == "zstd" + raw = result.response.payload.read() + assert ( + zstandard.ZstdDecompressor().stream_reader(io.BytesIO(raw)).read() == b"payload" + ) + + +def test_many_sends_metadata_headers() -> None: + pool = FakeObjectstore() + session = _session(pool) + + session.many( + [ + Put( + b"payload", + key="k", + content_type="text/plain", + metadata={"foo": "bär", "with-newline": "a\r\nb"}, + origin="127.0.0.1", + filename="rapport.pdf", + expiration_policy=TimeToLive(timedelta(hours=2)), + ) + ] + ).raise_for_failures() + + (part,) = pool.batches[0] + assert part.headers["content-type"] == "text/plain" + assert part.headers["x-sn-origin"] == "127.0.0.1" + assert part.headers["x-sn-expiration"] == "ttl:2h" + # Values that a header cannot carry verbatim are percent-encoded, which also + # keeps CR/LF from forging additional headers. + assert part.headers["x-snme-with-newline"] == "a%0D%0Ab" + assert part.headers["x-snme-foo"] == "b%C3%A4r" + + (result,) = list(session.many([Get("k")])) + assert isinstance(result, GetResult) and result.response is not None + assert result.response.metadata.custom == {"foo": "bär", "with-newline": "a\r\nb"} + assert result.response.metadata.filename == "rapport.pdf" + + +def test_many_reports_failures() -> None: + pool = FakeObjectstore(deny_writes=True) + pool.store["exists"] = ({}, b"hi") + session = _session(pool) + ops: list[Operation] = [Put(b"x", key="k1"), Get("exists"), Delete("k2")] + + failures = session.many(ops).failures() + assert sorted(failure.key for failure in failures) == ["k1", "k2"] # type: ignore[union-attr] + + with pytest.raises(ExceptionGroup) as exc_info: + session.many(ops).raise_for_failures() + assert len(exc_info.value.exceptions) == 2 + + +class _Rejects(FakeObjectstore): + """Turns the whole batch away.""" + + def request(self, *_args: Any, **_kwargs: Any) -> FakeResponse: + return FakeResponse(500, {}, b"boom") + + +class _Unreachable(FakeObjectstore): + """Never gets the request out at all.""" + + def request(self, *_args: Any, **_kwargs: Any) -> FakeResponse: + raise urllib3.exceptions.NewConnectionError( + None, # type: ignore[arg-type] + "connection refused", + ) + + +class _SkipsLastPart(FakeObjectstore): + """Answers every operation but the last.""" + + def _batch(self, body: Any, headers: dict[str, str]) -> FakeResponse: + response = super()._batch(body, headers) + parts = list(iter_multipart(response.headers["content-type"], [response.data])) + return _respond_with(parts[:-1]) + + +class _Truncated(FakeObjectstore): + """Cuts the response short, as a dropped connection would.""" + + def _batch(self, body: Any, headers: dict[str, str]) -> FakeResponse: + response = super()._batch(body, headers) + return FakeResponse(200, response.headers, response.data[:-20]) + + +class _BadDate(FakeObjectstore): + """Gives the first part a header the client cannot parse.""" + + def _execute(self, index: int, part: ResponsePart) -> ResponsePart: + response = super()._execute(index, part) + if index == 0: + response.headers["x-sn-time-created"] = "not-a-date" + return response + + +class _CannotSign(FakeObjectstore): + """Fails while the request is prepared, before any of it is sent.""" + + @property + def headers(self) -> dict[str, str]: + # `Session._make_headers` reads these to sign the request. + raise RuntimeError("signing failed") + + @headers.setter + def headers(self, value: dict[str, str]) -> None: + pass # `FakeObjectstore.__init__` assigns them; only reads should fail. + + +@pytest.mark.parametrize("concurrency", [1, 3]) +@pytest.mark.parametrize( + ("server", "failed"), + [ + (_CannotSign, {"a", "b"}), + (_Rejects, {"a", "b"}), + (_Unreachable, {"a", "b"}), + (_SkipsLastPart, {"b"}), + (_Truncated, {"b"}), + (_BadDate, {"a"}), + ], +) +def test_many_broken_batch_has_result_for_every_op( + server: type[FakeObjectstore], failed: set[str], concurrency: int +) -> None: + """ + However a batch breaks, each operation gets exactly one result. + + Both concurrencies, because they deliver results by different routes: + straight from the generator, or through a pool worker. + """ + pool = server(in_order=True) + pool.store["a"] = ({}, b"payload") + pool.store["b"] = ({}, b"payload") + session = _session(pool) + + results = list(session.many([Get("a"), Get("b")], concurrency=concurrency)) + + assert len(results) == 2 + assert {result.key for result in results if result.error is not None} == failed # type: ignore[union-attr] + + +@pytest.mark.parametrize( + ("headers", "expected"), + [ + # No index, so nothing ties the part to an operation: it is reported on + # its own, and the operation is left to the unanswered sweep. + pytest.param({"x-sn-batch-operation-status": "200 OK"}, 2, id="unattributable"), + # An index but no status. The part names its operation, so the failure + # belongs to that operation alone and nothing sweeps it up again. + pytest.param( + {"x-sn-batch-operation-index": "0", "x-sn-batch-operation-key": "a"}, + 1, + id="attributable", + ), + ], +) +def test_many_reports_part_it_cannot_use( + headers: dict[str, str], expected: int +) -> None: + class BadPartPool(FakeObjectstore): + def _batch(self, body: Any, _headers: dict[str, str]) -> FakeResponse: + return _respond_with([ResponsePart(headers, b"")]) + + results = list(_session(BadPartPool()).many([Get("a")])) + + assert len(results) == expected + assert all(result.error is not None for result in results) + # Either way the operation is accounted for exactly once, and by a result of + # its own kind rather than a bare `ErrorResult`. + (answered,) = [result for result in results if isinstance(result, GetResult)] + assert answered.index == 0 and answered.key == "a" + + +@pytest.mark.parametrize( + ("bad", "error"), + [ + # A metadata name cannot be percent-escaped the way a value can. + (Put(b"x", key="bad", metadata={"a\r\nb": "c"}), ValueError), + (Put(b"x", key="bad", metadata={"not a header": "c"}), ValueError), + # A value that is not a string cannot be escaped at all. + (Put(b"x", key="bad", metadata={"m": 123}), TypeError), # type: ignore[dict-item] + ], +) +def test_many_fails_only_operations_it_cannot_encode( + bad: Put, error: type[Exception] +) -> None: + pool = FakeObjectstore() + session = _session(pool) + + results = list(session.many([bad, Put(b"y", key="good-1"), Delete("good-2")])) + + by_index = {result.index: result for result in results} + assert isinstance(by_index[0].error, error) + assert by_index[1].error is None and by_index[2].error is None + # Dropping the first operation shifts the rest by one on the wire, so their + # results have to report the index they had in the input. + assert by_index[1].key == "good-1" # type: ignore[union-attr] + assert "good-1" in pool.store + + +def test_put_rejects_conflicting_compression() -> None: + with pytest.raises(ValueError): + Put(b"x", compress="zstd", precompressed="zstd") + with pytest.raises(ValueError): + Put(b"x", compress="brotli") # type: ignore[arg-type] + with pytest.raises(ValueError): + Put(b"x", precompressed="none") + + +def test_many_reports_index_of_each_operation() -> None: + pool = FakeObjectstore() # responds in reverse order + session = _session(pool) + keys = [f"key-{index}" for index in range(5)] + + results = list(session.many([Get(key) for key in keys], concurrency=1)) + + assert [result.key for result in results] != keys # type: ignore[union-attr] + assert {result.index: result.key for result in results} == dict( # type: ignore[union-attr] + enumerate(keys) + ) + + +def test_many_correlates_keyless_puts_by_index() -> None: + pool = FakeObjectstore() + session = _session(pool) + payloads = [f"payload-{index}".encode() for index in range(4)] + + results = list(session.many([Put(payload) for payload in payloads], concurrency=2)) + + assert len(results) == len(payloads) + for result in results: + assert isinstance(result, PutResult) and result.error is None + assert pool.store[result.key][1] == payloads[result.index] + + +def test_many_operation_count_limit() -> None: + pool = FakeObjectstore() + session = _session(pool) + + ops: list[Operation] = [ + Delete(f"key-{index}") for index in range(MAX_BATCH_OPS + 1) + ] + results = list(session.many(ops, concurrency=1)) + + assert len(results) == MAX_BATCH_OPS + 1 + assert [len(batch) for batch in pool.batches] == [MAX_BATCH_OPS, 1] + + +def test_iter_work_splits_on_body_size_limit() -> None: + one_mb = 1024 * 1024 + classified = [(index, Delete(f"key-{index}"), one_mb) for index in range(150)] + + batches = list(_iter_work(classified)) + + per_batch = MAX_BATCH_BODY_SIZE // one_mb + assert [len(item.ops) for item in batches] == [per_batch, 150 - per_batch] # type: ignore[union-attr] + + +def test_iter_work_oversized_operation() -> None: + classified = [(0, Delete("k"), MAX_BATCH_BODY_SIZE + 1)] + batches = list(_iter_work(classified)) + assert [len(item.ops) for item in batches] == [1] # type: ignore[union-attr] + + +def test_iter_work_flushes_pending_batch_before_individual_request() -> None: + classified: list[tuple[int, Operation, int | None]] = [ + (0, Get("a"), 0), + (1, Put(b"big", key="b"), None), + (2, Get("c"), 0), + ] + + items = list(_iter_work(classified)) + + assert [type(item).__name__ for item in items] == [ + "_Batch", + "_Individual", + "_Batch", + ] + + +def test_classify_size_cutoff() -> None: + session = _session(FakeObjectstore(), compression="zstd") + + # The largest payload whose worst-case compressed size still fits a part, + # matching the boundary the Rust client tests. + size = 1_044_496 + assert _zstd_compress_bound(size) == MAX_BATCH_PART_SIZE + batchable = _classify(session, Put(b"a" * size, key="k")) + assert batchable == MAX_BATCH_PART_SIZE + + unbatchable = _classify(session, Put(b"a" * (size + 1), key="k")) + assert unbatchable is None + + +def test_classify_uses_exact_size_precompressed_body() -> None: + session = _session(FakeObjectstore(), compression="zstd") + + # A size that the worst-case bound would push over the limit, but which is + # sent verbatim because the payload is already compressed. + payload = b"a" * MAX_BATCH_PART_SIZE + size = _classify(session, Put(payload, key="k", precompressed="zstd")) + assert size == MAX_BATCH_PART_SIZE + + +def test_classify_sends_gzip_individually() -> None: + session = _session(FakeObjectstore(), compression="gzip") + + assert _classify(session, Put(b"data", key="k")) is None + assert _classify(session, Put(b"data", key="k", precompressed="gzip")) == 4 + + +def test_classify_cannot_size_unseekable_stream() -> None: + session = _session(FakeObjectstore()) + size = _classify(session, Put(Unseekable(b"data"), key="k")) + assert size is None + + +def test_classify_rejects_foreign_operations() -> None: + session = _session(FakeObjectstore()) + with pytest.raises(TypeError): + _classify(session, "not an operation") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + "contents", + [ + pytest.param(b"x" * (2 * MAX_BATCH_PART_SIZE), id="oversized"), + pytest.param(Unseekable(b"x" * 16), id="unsized"), + ], +) +def test_many_sends_unbatchable_insert_individually( + contents: bytes | io.BytesIO, +) -> None: + """An insert too big to be a part, or of unknown size, gets its own request.""" + expected = contents if isinstance(contents, bytes) else contents.getvalue() + pool = FakeObjectstore() + session = _session(pool) + + results = list( + session.many([Put(contents, key="alone"), Put(b"small", key="small")]) + ) + + assert all(result.error is None for result in results) + assert pool.individual == ["alone"] + assert [len(batch) for batch in pool.batches] == [1] + assert pool.store["alone"][1] == expected + + +def test_many_keeps_individual_insert_compressed() -> None: + pool = FakeObjectstore() + session = _session(pool, compression="zstd") + big = b"x" * (4 * MAX_BATCH_PART_SIZE) + + (result,) = list(session.many([Put(big, key="big")])) + + assert result.error is None + metadata, stored = pool.store["big"] + assert metadata["content-encoding"] == "zstd" + assert zstandard.ZstdDecompressor().stream_reader(io.BytesIO(stored)).read() == big + + +def test_many_streams_batched_stream_body() -> None: + pool = FakeObjectstore() + session = _session(pool, compression="zstd") + payload = b"streamed payload " * 100 + + (result,) = list(session.many([Put(io.BytesIO(payload), key="k")])) + + assert result.error is None + metadata, stored = pool.store["k"] + assert metadata["content-encoding"] == "zstd" + assert ( + zstandard.ZstdDecompressor().stream_reader(io.BytesIO(stored)).read() == payload + ) + + +def test_many_leaves_caller_stream_open() -> None: + pool = FakeObjectstore() + session = _session(pool, compression="zstd") + source = io.BytesIO(b"payload") + + session.many([Put(source, key="k")]).raise_for_failures() + + # The compressing reader wrapped around it is ours to close, the stream isn't. + assert not source.closed + + +def test_zstd_body_holds_compressor_only_while_reading() -> None: + """ + Zstd compression context costs ~1MB RAM. We want to defer creating them + until reading and then drop it once the body is exhausted. + """ + source = io.BytesIO(b"payload" * 100) + body = _ZstdBody(source) + assert body._compressed_stream is None + + compressed = body.read(8) + started = body._compressed_stream is not None + assert compressed and started + while chunk := body.read(8): + compressed += chunk + + assert ( + zstandard.ZstdDecompressor().stream_reader(io.BytesIO(compressed)).read() + == b"payload" * 100 + ) + # The compressor is gone once the source is exhausted, but the source, which + # belongs to the caller, is left open. + assert body._compressed_stream is None + assert not source.closed + # A stream at the end of its input is at EOF, not closed, so reading on + # returns empty bytes the way any file object does. + assert not body.closed + assert body.read(8) == b"" + + +def test_many_does_not_resend_stream_bodies() -> None: + pool = FakeObjectstore() + session = _session(pool, compression="zstd") + + session.many([Put(io.BytesIO(b"streamed"), key="k")]).raise_for_failures() + + # Connect retries stay on: they happen before the body is read. + (retries,) = pool.retries_used + assert (retries.read, retries.other, retries.redirect) == (0, 0, 0) + assert retries.connect == 3 + + +def test_many_does_not_override_retry_policy_for_non_stream_bodies() -> None: + pool = FakeObjectstore() + session = _session(pool, compression="zstd") + + # Only inserts carry a body, and only a compressed one is a stream. Both of + # these are re-sendable even though the usecase compresses by default. + session.many([Get("a"), Delete("b"), Head("c")]).raise_for_failures() + session.many([Put(b"payload", key="k", compress="none")]).raise_for_failures() + + assert pool.retries_used == [None, None] + + +def test_many_reads_operations_lazily() -> None: + consumed = [] + + def operations() -> Iterator[Operation]: + for index in range(MAX_BATCH_OPS + 100): + consumed.append(index) + yield Delete(f"key-{index}") + + pool = FakeObjectstore() + session = _session(pool) + + results = session.many(operations(), concurrency=1) + assert consumed == [] + assert pool.requests == 0 + + next(iter(results)) + assert len(consumed) <= MAX_BATCH_OPS + 1 + assert pool.requests == 1 + + results.close() + + +def test_many_drained_batch_closes_connection() -> None: + pool = FakeObjectstore() + session = _session(pool) + + list(session.many([Get("a"), Get("b")], concurrency=1)) + + (response,) = pool.responses + assert response.released and not response.closed + + +@pytest.mark.parametrize("concurrency", [1, 2]) +def test_many_abandoned_batch_closes_connection(concurrency: int) -> None: + pool = FakeObjectstore() + for index in range(50): + pool.store[f"key-{index}"] = ({}, b"payload") + session = _session(pool) + + # One batch of fifty, so results are still coming when the caller walks away + # after the first: on the calling thread at a concurrency of one, and from a + # worker parked on the result queue above it. + with session.many( + [Get(f"key-{index}") for index in range(50)], concurrency=concurrency + ) as results: + next(iter(results)) + + # The response was abandoned before its body was read to the end, so the + # connection cannot be returned to the pool. + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + if all(response.closed for response in pool.responses): + break + time.sleep(0.05) + assert all(response.closed and not response.released for response in pool.responses) + + +def test_many_runs_on_calling_thread_at_concurrency_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An effective concurrency of one means no thread pool, however it arose.""" + threads: list[int] = [] + + class RecordingPool(FakeObjectstore): + def _batch(self, body: Any, headers: dict[str, str]) -> FakeResponse: + threads.append(threading.get_ident()) + return super()._batch(body, headers) + + # How the pool arrives at a concurrency of one is `_fit_to_connection_pool`'s + # business, and is covered by its own tests. + monkeypatch.setattr(many_module, "_fit_to_connection_pool", lambda *_args: 1) + + session = _session(RecordingPool()) + session.many([Get("a"), Delete("b")], concurrency=8).raise_for_failures() + + assert threads == [threading.get_ident()] + + +def test_many_runs_batches_concurrently() -> None: + barrier = threading.Barrier(3, timeout=10) + + class BarrierPool(FakeObjectstore): + def _batch(self, body: Any, headers: dict[str, str]) -> FakeResponse: + # Only completes once three requests are in flight at the same time. + barrier.wait() + return super()._batch(body, headers) + + pool = BarrierPool() + session = _session(pool) + # Enough operations to fill two batches and start a third. + ops: list[Operation] = [ + Delete(f"key-{index}") for index in range(2 * MAX_BATCH_OPS + 1) + ] + + results = list(session.many(ops, concurrency=3)) + + assert len(pool.batches) == 3 + assert len(results) == len(ops) + assert all(result.error is None for result in results) + + +def test_many_drops_work_when_abandoned() -> None: + pool = FakeObjectstore() + for index in range(40): + pool.store[f"key-{index}"] = ({}, b"payload") + session = _session(pool) + + # An unsized insert cuts the batch it lands in, so this is four work items: + # a batch, an insert, another batch, and the insert that must not be sent. + ops: list[Operation] = [Get(f"key-{index}") for index in range(20)] + ops.append(Put(Unseekable(b"payload"), key="cuts-the-batch")) + ops += [Get(f"key-{index}") for index in range(20, 40)] + ops.append(Put(Unseekable(b"payload"), key="never-sent")) + + results = session.many(ops, concurrency=2) + next(iter(results)) + results.close() + + # The pool is internal, so there is nothing to join on: give a worker that + # wrongly picked the insert up long enough to prove it. + deadline = time.monotonic() + 0.5 + while time.monotonic() < deadline and "never-sent" not in pool.store: + time.sleep(0.01) + + # The first two items were dispatched, so this is work stopping rather than + # work never having started. + assert "cuts-the-batch" in pool.store + assert "never-sent" not in pool.store + + +def _pool_queue(client: Client) -> Any: + connections = client._pool.pool + assert connections is not None + return connections + + +def test_fit_to_connection_pool_block_is_false() -> None: + """`concurrency` is the only knob; the pool is sized from it.""" + client = Client("http://localhost:8888") + session = client.session(Usecase("testing", compression="none"), org=42) + connections = _pool_queue(client) + assert connections.maxsize == 1 + + assert _fit_to_connection_pool(session, 8) == 8 + assert connections.maxsize == 8 + + # Only ever grows, so a pooled connection is never left without a slot. + assert _fit_to_connection_pool(session, 2) == 2 + assert connections.maxsize == 8 + + +def test_fit_to_connection_pool_block_is_true() -> None: + """`block=True` makes the pool size a deliberate cap on connections.""" + client = Client( + "http://localhost:8888", connection_kwargs={"maxsize": 2, "block": True} + ) + session = client.session(Usecase("testing", compression="none"), org=42) + + assert _fit_to_connection_pool(session, 8) == 2 + # Asking for less than the cap is still honoured, and the cap stays put. + assert _fit_to_connection_pool(session, 1) == 1 + assert _pool_queue(client).maxsize == 2 + + +@contextmanager +def _initialized_sentry() -> Iterator[None]: + """Sentry's integrations are only installed for an initialized SDK.""" + previous = sentry_sdk.get_global_scope().client + sentry_sdk.init(dsn=None, traces_sample_rate=1.0) + try: + yield + finally: + sentry_sdk.get_global_scope().set_client(previous) + + +def test_many_pool_thread_sentry_context() -> None: + """ + Workers must run with the scope of the thread that submitted them. + + Nothing here does that by hand: `ThreadingIntegration` patches + `ThreadPoolExecutor.submit` to carry the submitting scope into the task. We + just want to make sure it continues working. + """ + seen: list[tuple[int, str | None]] = [] + + class RecordingPool(FakeObjectstore): + def request(self, method: str, url: str, **kwargs: Any) -> FakeResponse: + seen.append( + ( + threading.get_ident(), + sentry_sdk.get_current_scope().get_traceparent(), + ) + ) + return super().request(method, url, **kwargs) + + with _initialized_sentry(): + session = _session(RecordingPool()) + traceparent = sentry_sdk.get_current_scope().get_traceparent() + + # A batch request and an individual one, so both paths are covered. + ops: list[Operation] = [Get("a"), Put(Unseekable(b"data"), key="b")] + results = list(session.many(ops, concurrency=2)) + + # No errors + assert all(result.error is None for result in results) + # Every thread had the same traceparent + assert [scope for _, scope in seen] == [traceparent, traceparent] + # Those threads actually weren't just the caller thread, proving the scopes + # were inherited + assert all(ident != threading.get_ident() for ident, _ in seen) + + +def test_many_does_not_hang_interpreter_exit_when_iterating() -> None: + """A consumer waiting on results must not outlive the workers feeding it.""" + script = f""" +import sys, threading, time +sys.path.insert(0, {str(pathlib.Path(__file__).parent)!r}) +from test_many import FakeObjectstore, _session + +from objectstore_client.many import Get + +class Slow(FakeObjectstore): + def _batch(self, body, headers): + time.sleep(0.5) + return super()._batch(body, headers) + +pool = Slow() +session = _session(pool) +started = threading.Event() + +def consume(): + started.set() + # Never abandoned, so the only thing that can release this is the shutdown. + for _ in session.many([Get(f"key-{{i}}") for i in range(50)], concurrency=2): + pass + +# Not a daemon, so the interpreter joins it on the way out. +threading.Thread(target=consume, daemon=False).start() +started.wait() +time.sleep(0.1) # let the request get in flight, then fall off the main thread +""" + subprocess.run([sys.executable, "-c", script], timeout=60, check=True) + + +def test_many_does_not_hang_interpreter_exit_when_abandoned() -> None: + """Workers park on a full queue; nothing may be left holding up the exit.""" + script = f""" +import sys +sys.path.insert(0, {str(pathlib.Path(__file__).parent)!r}) +from test_many import FakeObjectstore, _session + +from objectstore_client.many import Get + +pool = FakeObjectstore() +session = _session(pool) +# One batch of many operations, so the worker has far more results to hand over +# than the queue holds and is still parked on it when the interpreter exits. +ops = [Get(f"key-{{i}}") for i in range(50)] + +results = iter(session.many(ops, concurrency=2)) +next(results) +# `results` stays referenced, so the generator is never closed and nothing +# releases the worker. +""" + subprocess.run([sys.executable, "-c", script], timeout=60, check=True) + + +def test_many_batch_operations_metric() -> None: + metrics = RecordingMetrics() + pool = FakeObjectstore() + session = _session(pool, metrics=metrics) + + session.many( + [Get("a"), Get("b"), Put(b"x", key="c"), Delete("d"), Head("e")], + concurrency=1, + ).raise_for_failures() + + counts = { + (tags or {}).get("operation"): value + for name, value, tags in metrics.counters + if name == "storage.batch.operations" + } + assert counts == {"get": 2, "insert": 1, "delete": 1, "head": 1} + assert all( + (tags or {}).get("usecase") == "testing" for _, _, tags in metrics.counters + ) + + +def test_many_batch_latency_metric() -> None: + metrics = RecordingMetrics() + session = _session(FakeObjectstore(), metrics=metrics) + + session.many([Get(f"key-{index}") for index in range(5)]).raise_for_failures() + + (latency,) = [m for m in metrics.distributions if m[0] == "storage.batch.latency"] + assert (latency[2] or {})["operations"] == "4-7" + + +def test_batch_size_bucket() -> None: + assert [batch_size_bucket(count) for count in (0, 1, 2, 3, 4, 7, 8, 1000)] == [ + "0", + "1", + "2-3", + "2-3", + "4-7", + "4-7", + "8-15", + "512-1023", + ] + + +def test_many_individual_insert_metrics() -> None: + metrics = RecordingMetrics() + pool = FakeObjectstore() + session = _session(pool, metrics=metrics) + + session.many( + [Put(b"x" * (2 * MAX_BATCH_PART_SIZE), key="big")] + ).raise_for_failures() + + names = {name for name, _, _ in metrics.distributions} + assert "storage.put.latency" in names + # Nothing went through the batch endpoint, so it is not counted as a batch. + assert "storage.batch.latency" not in names + assert metrics.counters == [] + + +def test_parse_status() -> None: + assert _parse_status("200 OK") == 200 + assert _parse_status("404 Not Found") == 404 + assert _parse_status("204") == 204 + assert _parse_status(None) is None + assert _parse_status("") is None + assert _parse_status("garbage") is None