"""Staging that materializes a remote shard as an unnamed local file.
The object is written to a temporary file whose name is removed immediately, and
the reader is handed ``/dev/fd/N`` for the descriptor still holding it. On a
local filesystem the kernel reclaims the space when the last descriptor closes
or the process dies, including under SIGKILL.
That guarantee is filesystem-dependent, not universal. NFS and SMB implement
unlink-while-open by renaming to a hidden ``.nfsXXXX`` file and removing it when
the descriptor closes, so a process killed mid-read leaves a full shard behind.
``stage`` checks for this after unlinking and warns, because both the default
(``TMPDIR``) and the ``directory`` argument can land on such a mount.
This is already a house idiom: ``LocalObjectStore.indexed_uri`` hands a confined
file to a native reader the same way.
**Cost.** A staged shard is fetched whole, because the reader needs a complete
file to index into. Reading one record therefore costs one shard. Descriptors
are retained and reused across calls — see ``max_open_shards`` — so consecutive
reads from the same shard pay once rather than once per record; a pattern that
alternates between more shards than are retained pays per read. Staged bytes are
unnamed, so they are invisible to ``ls`` and ``du`` while still consuming the
filesystem: ``max_open_shards`` times the largest shard is the working set.
``/dev/fd`` is POSIX-only. The package is already Linux and macOS only — the
local backend's atomic publication needs ``os.link`` and ``fcntl.flock`` — so
this adds no new platform constraint. The reader must use positional reads:
on macOS ``/dev/fd/N`` duplicates the description and therefore shares its file
offset, where on Linux it reopens the inode at zero.
"""
from __future__ import annotations
import logging
import os
import tempfile
import threading
from collections import OrderedDict
from collections.abc import Iterator
from contextlib import contextmanager
from typing import IO
from signal_dataset._internal.constants import DEFAULT_MAX_OPEN_SHARDS
from signal_dataset.errors import StorageError
from signal_dataset.storage.contracts import ObjectDownloader, ObjectStore
from signal_dataset.storage.registry import object_store_for
from signal_dataset.storage.staging.contracts import ShardStagingArea
logger = logging.getLogger(__name__)
TEMPORARY_PREFIX = "signal-dataset-shard-"
[docs]
class EphemeralFileStaging(ShardStagingArea):
"""Materializes a shard as an unnamed local file, retaining it for reuse."""
def __init__(
self,
*,
objects: ObjectStore | None = None,
directory: str | os.PathLike[str] | None = None,
max_object_bytes: int | None = None,
max_open_shards: int = DEFAULT_MAX_OPEN_SHARDS,
) -> None:
if max_object_bytes is not None and not positive(max_object_bytes):
raise ValueError("max_object_bytes must be a positive integer or None")
if not positive(max_open_shards):
raise ValueError("max_open_shards must be a positive integer")
self._objects = objects
# Defaults to the platform temporary directory, which honours TMPDIR.
# Callers whose /tmp is small or slow pass their own; nothing here
# hard-codes a location.
self._directory = None if directory is None else os.fspath(directory)
self._max_object_bytes = max_object_bytes
self._max_open_shards = max_open_shards
self._open: OrderedDict[tuple[str, int | None], int] = OrderedDict()
#: How many `stage()` blocks are currently reading each descriptor.
self._holders: dict[int, int] = {}
#: Evicted while still being read. Closed by the last reader to leave.
self._retired: set[int] = set()
self._lock = threading.Lock()
[docs]
@contextmanager
def stage(self, uri: str, *, generation: int | None = None) -> Iterator[str]:
# The lock is never held across the yield: a reader may stage again
# while inside the block, and holding it would serialize every read.
descriptor = self.acquire(uri, generation)
try:
yield f"/dev/fd/{descriptor}"
finally:
self.release(descriptor)
[docs]
def acquire(self, uri: str, generation: int | None) -> int:
"""Return a descriptor for `uri`, counting this caller as a reader.
Every acquire must be matched by a `release`. Eviction alone must never
close a descriptor: another thread may be reading through it right now,
and closing it underneath would hand that reader a descriptor number the
kernel is free to reuse for something else.
"""
key = (uri, generation)
with self._lock:
retained = self._open.get(key)
if retained is not None:
self._open.move_to_end(key)
self._holders[retained] = self._holders.get(retained, 0) + 1
logger.debug("Reusing a staged shard", extra={"uri": uri})
return retained
descriptor = self.materialize(uri, generation)
closable: list[int] = []
with self._lock:
if key in self._open:
# Another thread materialized the same shard first; keep theirs
# so both callers see one descriptor, and drop this one. It was
# never handed out, so nobody can be reading it.
closable.append(descriptor)
descriptor = self._open[key]
self._open.move_to_end(key)
else:
self._open[key] = descriptor
self._holders[descriptor] = self._holders.get(descriptor, 0) + 1
while len(self._open) > self._max_open_shards:
stale = self._open.popitem(last=False)[1]
if self._holders.get(stale, 0):
self._retired.add(stale)
else:
closable.append(stale)
for stale in closable:
os.close(stale)
return descriptor
[docs]
def release(self, descriptor: int) -> None:
"""Mark one reader as finished, closing the descriptor if it was evicted."""
with self._lock:
remaining = self._holders.get(descriptor, 0) - 1
if remaining > 0:
self._holders[descriptor] = remaining
return
self._holders.pop(descriptor, None)
if descriptor not in self._retired:
return # still retained for reuse
self._retired.discard(descriptor)
os.close(descriptor)
[docs]
def materialize(self, uri: str, generation: int | None) -> int:
"""Fetch an object into a fresh unnamed file and return its descriptor."""
store = self._objects if self._objects is not None else object_store_for(uri)
self.check_size(store, uri, generation)
descriptor, name = tempfile.mkstemp(prefix=TEMPORARY_PREFIX, dir=self._directory)
try:
# Unlink inside the guard: if it fails, the descriptor still has to
# be closed and the name is still on disk to report.
os.unlink(name)
self.warn_if_still_named(descriptor, name)
with os.fdopen(descriptor, "wb", closefd=False) as stream:
self.fetch(store, uri, stream, generation)
stream.flush()
# boto3 writes ranges out of order, so the offset is wherever the
# last chunk landed. On macOS a reader inherits it; rewind.
os.lseek(descriptor, 0, os.SEEK_SET)
except BaseException:
os.close(descriptor)
raise
logger.debug("Staged shard for a path-only reader", extra={"uri": uri})
return descriptor
[docs]
def close(self) -> None:
"""Release every retained descriptor.
Descriptors still being read are retired rather than closed, so a
concurrent `stage()` block finishes against a descriptor that is still
open and the last reader out closes it.
"""
with self._lock:
descriptors = [
descriptor
for descriptor in self._open.values()
if not self._holders.get(descriptor)
]
self._retired.update(
descriptor
for descriptor in self._open.values()
if self._holders.get(descriptor)
)
self._open.clear()
for descriptor in descriptors:
try:
os.close(descriptor)
except OSError: # pragma: no cover - already closed
logger.warning("Staged descriptor was already closed")
def __del__(self) -> None: # pragma: no cover - interpreter teardown
try:
self.close()
except Exception:
pass
[docs]
def check_size(self, store: ObjectStore, uri: str, generation: int | None) -> None:
if self._max_object_bytes is None:
return
size = store.info(uri, generation=generation).size
if size > self._max_object_bytes:
raise StorageError(
f"shard of {size} bytes exceeds the {self._max_object_bytes} byte staging "
f"limit for {uri}"
)
[docs]
@staticmethod
def warn_if_still_named(descriptor: int, name: str) -> None:
"""Detect a filesystem that renames rather than unlinks an open file."""
if os.fstat(descriptor).st_nlink != 0:
logger.warning(
"Staging directory does not reclaim unlinked open files, so a shard "
"may be left behind if this process is killed",
extra={"path": name},
)
[docs]
@staticmethod
def fetch(
store: ObjectStore, uri: str, stream: IO[bytes], generation: int | None
) -> None:
if isinstance(store, ObjectDownloader):
store.download_to_fileobj(uri, stream, generation=generation)
return
# No streaming capability, so the whole object lands in memory once.
stream.write(store.read(uri, generation=generation))
def positive(value: object) -> bool:
return not isinstance(value, bool) and isinstance(value, int) and value > 0