"""Amazon S3 backend built on conditional writes.
S3 has no object generation. It identifies a version with an ETag, and offers
two conditional headers that give exactly the guarantees this package needs:
``If-None-Match: *`` makes a write create-only, and ``If-Match: <etag>`` makes
it a compare-and-swap.
Persisted references written against an S3 root therefore carry no version, the
same as a local root. The integer this store returns from ``read_version`` and
``generation`` is a runtime compare-and-swap handle only, derived from the ETag
(see :mod:`.versioning`), and never reaches a control document.
"""
from __future__ import annotations
import logging
import threading
import time
from collections import OrderedDict
from pathlib import Path
from typing import IO, Any, cast
from signal_dataset._internal.retry import RetryPolicy
from signal_dataset.errors import PublicationCollisionError, StorageError
from signal_dataset.storage.backends.s3 import errors
from signal_dataset.storage.backends.s3.client import (
build_client,
load_boto3,
loaded_transport_errors,
)
from signal_dataset.storage.backends.s3.options import S3Options
from signal_dataset.storage.backends.s3.versioning import normalize_etag, version_token
from signal_dataset.storage.contracts import ObjectInfo, ObjectStore, ObjectVersion
from signal_dataset.storage.paths import s3_parts
logger = logging.getLogger(__name__)
ANY_OBJECT = "*"
[docs]
class S3ObjectStore(ObjectStore):
"""Create-only object transport over S3 conditional writes."""
def __init__(
self,
client: Any | None = None,
*,
options: S3Options | None = None,
client_error: type[Exception] | tuple[type[Exception], ...] | None = None,
) -> None:
self._options = S3Options() if options is None else options
if client is None:
_, transport_errors = load_boto3()
client = build_client(self._options)
client_error = transport_errors
self._client = client
if client_error is None:
client_error = loaded_transport_errors()
# An empty tuple is a valid `except` target that catches nothing, which
# is the right answer only when botocore is absent entirely.
self._client_error: type[Exception] | tuple[type[Exception], ...] = client_error or ()
# Maps a token back to the ETag it was minted from, so a swap can send
# the caller's own version rather than whatever is current. A store is
# shared across threads by Dataset, so the memo is guarded.
self._etags: OrderedDict[int, str] = OrderedDict()
self._etags_lock = threading.Lock()
@property
def options(self) -> S3Options:
return self._options
[docs]
def read(self, uri: str, *, generation: int | None = None) -> bytes:
self.reject_generation(generation)
bucket, key = s3_parts(uri)
try:
response = self._client.get_object(Bucket=bucket, Key=key)
# The body streams after the headers arrive, so a mid-transfer
# stall raises here rather than from `get_object`. Reading inside
# the `try` is what keeps a botocore timeout from escaping an API
# whose contract is StorageError.
payload = cast(bytes, response["Body"].read())
except self._client_error as exc:
raise errors.translate(exc, uri) from exc
return payload
[docs]
def read_version(self, uri: str) -> ObjectVersion:
bucket, key = s3_parts(uri)
try:
response = self._client.get_object(Bucket=bucket, Key=key)
payload = cast(bytes, response["Body"].read())
except self._client_error as exc:
raise errors.translate(exc, uri) from exc
return ObjectVersion(payload, self.remember(self.require_etag(response)))
[docs]
def info(self, uri: str, *, generation: int | None = None) -> ObjectInfo:
self.reject_generation(generation)
return ObjectInfo(self.head(uri)[0])
[docs]
def generation(self, uri: str) -> int:
return self.remember(self.head(uri)[1])
[docs]
def create(self, uri: str, data: bytes, *, content_type: str) -> None:
self.put(
uri,
data,
content_type=content_type,
condition={"IfNoneMatch": ANY_OBJECT},
content_length=len(data),
)
[docs]
def create_file(self, uri: str, path: Path, *, content_type: str) -> None:
size = path.stat().st_size
limit = self._options.max_single_part_bytes
if size > limit:
raise StorageError(
f"object of {size} bytes exceeds the {limit} byte single-request limit for {uri}; "
"cap shard size with StorageOptions.max_encoded_shard_bytes"
)
with path.open("rb") as stream:
self.put(
uri,
stream,
content_type=content_type,
condition={"IfNoneMatch": ANY_OBJECT},
# Declared explicitly so the bytes on the wire are the bytes the
# limit above was checked against, even if the file grows.
content_length=size,
)
[docs]
def compare_and_swap(
self,
uri: str,
data: bytes,
*,
expected_generation: int | None,
content_type: str,
) -> None:
if expected_generation is None:
self.create(uri, data, content_type=content_type)
return
with self._etags_lock:
etag = self._etags.get(expected_generation)
if etag is None:
# No record of minting this token, so fall back to comparing the
# current version. A token collision here would be undetectable,
# which is why the token is 128 bits wide.
try:
etag = self.head(uri)[1]
except FileNotFoundError as exc:
raise PublicationCollisionError(uri) from exc
if version_token(etag) != expected_generation:
raise PublicationCollisionError(uri)
try:
self.put(
uri,
data,
content_type=content_type,
condition={"IfMatch": etag},
content_length=len(data),
)
except FileNotFoundError as exc:
# A racing delete answers a conditional PUT with 404, not 412. The
# local backend reports the same situation as a collision, and the
# annotation catalog's retry loop only catches that.
raise PublicationCollisionError(uri) from exc
[docs]
def download_to_fileobj(
self, uri: str, stream: IO[bytes], *, generation: int | None = None
) -> None:
"""Stream an object into a file object without buffering it whole."""
self.reject_generation(generation)
bucket, key = s3_parts(uri)
logger.debug("Downloading S3 object", extra={"uri": uri})
try:
self._client.download_fileobj(bucket, key, stream)
except self._client_error as exc:
raise errors.translate(exc, uri) from exc
[docs]
def list(self, prefix_uri: str) -> list[str]:
bucket, prefix = s3_parts(prefix_uri, allow_empty_name=True)
try:
pages = self._client.get_paginator("list_objects_v2").paginate(
Bucket=bucket, Prefix=prefix
)
return sorted(
f"s3://{bucket}/{item['Key']}"
for page in pages
for item in page.get("Contents", ())
)
except self._client_error as exc:
raise errors.translate(exc, prefix_uri) from exc
[docs]
def head(self, uri: str) -> tuple[int, str]:
"""Return an object's size and ETag in one request."""
bucket, key = s3_parts(uri)
try:
response = self._client.head_object(Bucket=bucket, Key=key)
except self._client_error as exc:
raise errors.translate(exc, uri) from exc
size = response.get("ContentLength")
if not isinstance(size, int):
raise StorageError(f"S3 response for {uri} omitted ContentLength")
return size, self.require_etag(response)
[docs]
def put(
self,
uri: str,
body: bytes | IO[bytes],
*,
content_type: str,
condition: dict[str, str],
content_length: int,
) -> None:
"""Write conditionally, retrying only the racing-writer conflict."""
bucket, key = s3_parts(uri)
limit = self._options.conflict_attempt_limit
for attempt in range(limit):
try:
self._client.put_object(
Bucket=bucket,
Key=key,
Body=body,
ContentType=content_type,
ContentLength=content_length,
**condition,
)
return
except self._client_error as exc:
if not errors.is_conflict(exc):
raise errors.translate(exc, uri) from exc
if attempt == limit - 1:
raise StorageError(
f"S3 reported concurrent conditional writes for {uri} "
f"after {limit} attempts"
) from exc
logger.debug(
"Retrying S3 conditional write after a conflict",
extra={"uri": uri, "attempt": attempt + 1},
)
time.sleep(self.backoff(attempt))
self.rewind(body)
[docs]
def backoff(self, attempt: int) -> float:
"""Capped exponential delay with jitter.
A 409 means writers collided, so retrying them in lockstep reproduces
the collision. The cap keeps a generous attempt limit from turning into
an unbounded sleep.
The schedule itself is :class:`RetryPolicy`'s, so this backend and the
operation-level retry share one implementation. What differs is the
tuning: this retries a conflict on a single key, which clears in
milliseconds, so its defaults are far shorter.
"""
# `attempts` is deliberately not passed: `delay` does not read it, and
# `conflict_attempt_limit` has its own bounds. Forwarding it would
# subject a legal S3Options to RetryPolicy's separate attempt ceiling.
return RetryPolicy(
base_seconds=self._options.conflict_retry_base_seconds,
max_seconds=self._options.conflict_retry_max_seconds,
).delay(attempt)
[docs]
def remember(self, etag: str) -> int:
"""Mint a token and record the ETag it came from."""
token = version_token(etag)
with self._etags_lock:
self._etags[token] = etag
self._etags.move_to_end(token)
while len(self._etags) > self._options.etag_memo_size:
self._etags.popitem(last=False)
return token
[docs]
@staticmethod
def rewind(body: bytes | IO[bytes]) -> None:
if not isinstance(body, bytes) and getattr(body, "seekable", None) and body.seekable():
body.seek(0)
[docs]
def require_etag(self, response: dict[str, Any]) -> str:
etag = response.get("ETag")
if not isinstance(etag, str) or not normalize_etag(etag):
raise StorageError("S3 response omitted an ETag")
return normalize_etag(etag)
[docs]
def reject_generation(self, generation: int | None) -> None:
if generation is not None:
raise ValueError("S3 objects do not have generations")