leblob 0.4.0.dev1__py3-none-any.whl

leblob/blobstore.py

@@ -0,0 +1,594 @@
+"""Blob files on one or more volumes: durable writes, range reads and deletion.
+
+A blob's ID is random, so user-supplied keys never reach the filesystem.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import logging
+import os
+import secrets
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import aclosing, suppress
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Self
+
+from leblob import metrics
+from leblob.model import Blob, BlobSlice
+from leblob.s3 import errors
+from leblob.s3.errors import S3Error
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        AsyncGenerator,
+        AsyncIterable,
+        Iterable,
+        Mapping,
+        Sequence,
+    )
+    from types import TracebackType
+
+    from leblob.model import VolumeRecord
+    from leblob.s3.payload import Hasher
+
+BUFFER_BYTES = 1024 * 1024
+TMP_MAX_AGE_SECONDS = 24 * 3600
+VOLUME_FILE = "volume.json"
+
+_SHARD_LENGTH = 3
+_SHARD_COUNT = 16**_SHARD_LENGTH
+_ALLOWED_IN_A_NEW_VOLUME = {"lost+found"}  # A fresh ext4 filesystem has it.
+_ID_NOT_ASSIGNED_YET = 0  # The metadata assigns it, once the files are in place.
+_AVOID_SECONDS = 30
+_RANDOM = secrets.SystemRandom()
+
+logger = logging.getLogger(__name__)
+
+
+class VolumeError(Exception):
+    """A volume that can't be used. The message says why."""
+
+
+def existing_files(directory: Path) -> Iterable[tuple[str, os.stat_result]]:
+    """Yield the files in a directory that other workers are changing concurrently.
+
+    A file listed a moment ago may have been committed or collected since.
+    """
+    with os.scandir(directory) as entries:
+        for entry in entries:
+            try:
+                yield entry.path, entry.stat()
+            except FileNotFoundError:
+                continue
+
+
+def fsync_directory(path: Path) -> None:
+    descriptor = os.open(path, os.O_RDONLY)
+    try:
+        os.fsync(descriptor)
+    finally:
+        os.close(descriptor)
+
+
+def place(free_bytes: Mapping[int, int]) -> list[int]:
+    """Order volumes at random, weighted by free space, so that they fill evenly.
+
+    The first in the list are the ones to use, and the others take over if one fails.
+    A full volume keeps a tiny chance: if all are full, the write gets the real error.
+    """
+    remaining = dict(free_bytes)
+    order = []
+    while remaining:
+        weights = [free + 1 for free in remaining.values()]
+        chosen = _RANDOM.choices(list(remaining), weights=weights)[0]
+        order.append(chosen)
+        del remaining[chosen]
+    return order
+
+
+@dataclass(frozen=True, slots=True)
+class Volume:
+    """A directory of immutable blob files, usually a whole disk."""
+
+    root: Path
+    volume_id: int
+
+    @property
+    def tmp_dir(self) -> Path:
+        return self.root / "tmp"
+
+    @property
+    def blobs_dir(self) -> Path:
+        return self.root / "blobs"
+
+    def initialize(self) -> None:
+        """Create every shard up front, so writes never create a directory."""
+        self.tmp_dir.mkdir(parents=True, exist_ok=True)
+        for number in range(_SHARD_COUNT):
+            shard = f"{number:0{_SHARD_LENGTH}x}"
+            (self.blobs_dir / shard).mkdir(parents=True, exist_ok=True)
+
+    def path(self, blob_id: str) -> Path:
+        return self.blobs_dir / blob_id[:_SHARD_LENGTH] / blob_id
+
+    def free_bytes(self) -> int:
+        stat = os.statvfs(self.root)
+        return stat.f_bavail * stat.f_frsize
+
+    def claim(self, identity: str) -> None:
+        """Durably write the file naming this volume, the last step of creating it."""
+        partial_file = self.root / f"{VOLUME_FILE}.partial"
+        with partial_file.open("w", encoding="utf-8") as file:
+            json.dump({"identity": identity}, file)
+            file.flush()
+            os.fsync(file.fileno())
+        partial_file.rename(self.root / VOLUME_FILE)
+        fsync_directory(self.root)
+
+    def delete(self, blob_id: str) -> None:
+        self.path(blob_id).unlink(missing_ok=True)
+
+    def identifies_itself(self) -> bool:
+        """False once the disk is unmounted, which empties the directory underneath."""
+        return (self.root / VOLUME_FILE).exists()
+
+    def copy_is_intact(self, blob: Blob) -> bool:
+        """Check a copy against its MD5, reading the disk and not the page cache."""
+        digest = hashlib.md5(usedforsecurity=False)
+        with self.path(blob.blob_id).open("rb") as file:
+            if sys.platform == "linux":  # The deployment target.
+                os.posix_fadvise(file.fileno(), 0, 0, os.POSIX_FADV_DONTNEED)
+            while data := file.read(BUFFER_BYTES):
+                digest.update(data)
+        return digest.hexdigest() == blob.md5
+
+    def sweep_tmp(self, now: float) -> int:
+        """Remove temp files abandoned by a crash. Returns how many were removed."""
+        removed = 0
+        for path, stat in existing_files(self.tmp_dir):
+            if now - stat.st_mtime > TMP_MAX_AGE_SECONDS:
+                Path(path).unlink(missing_ok=True)
+                removed += 1
+        return removed
+
+    def scan(self) -> Iterable[tuple[str, int, float]]:
+        """Yield (blob ID, size, modification time) for every blob on disk."""
+        for shard in self.blobs_dir.iterdir():
+            if not shard.is_dir():  # Such as a `.DS_Store`: not ours, and no harm.
+                continue
+            for path, stat in existing_files(shard):
+                yield Path(path).name, stat.st_size, stat.st_mtime
+
+
+@dataclass(slots=True)
+class BlobStore:
+    """Every volume of a store that can be used now.
+
+    A volume that fails is tried last for a while, in this worker: otherwise a dead
+    disk would keep costing requests, until a person or repair deals with it.
+    """
+
+    volumes: Mapping[int, Volume]
+    copies: int = 1
+    draining: frozenset[int] = frozenset()
+    _avoid_until: dict[int, float] = field(
+        default_factory=dict, init=False, compare=False
+    )
+
+    def path(self, blob: Blob) -> Path:
+        return self.volumes[blob.volume_ids[0]].path(blob.blob_id)
+
+    def report_failure(self, volume_id: int, error: OSError) -> None:
+        logger.warning("Volume %s failed: %s", volume_id, error)
+        metrics.counters.add(metrics.VOLUME_FAILURES, volume=str(volume_id))
+        self._avoid_until[volume_id] = time.monotonic() + _AVOID_SECONDS
+
+    def _is_avoided(self, volume_id: int) -> bool:
+        return self._avoid_until.get(volume_id, 0) > time.monotonic()
+
+    def _free_bytes(self, volume_id: int) -> int:
+        try:
+            return self.volumes[volume_id].free_bytes()
+        except OSError as error:
+            self.report_failure(volume_id, error)
+            return 0
+
+    def follow(self, current: BlobStore) -> None:
+        """Take on the volumes and settings of a store opened more recently."""
+        self.volumes = current.volumes
+        self.copies = current.copies
+        self.draining = current.draining
+
+    def placement_order(self, exclude: Iterable[int] = ()) -> list[Volume]:
+        """The volumes new data could go on, the ones to use first."""
+        unusable = self.draining.union(exclude)
+        free_bytes = {
+            volume_id: self._free_bytes(volume_id)
+            for volume_id in self.volumes
+            if volume_id not in unusable
+        }
+        order = place({
+            key: free for key, free in free_bytes.items() if not self._is_avoided(key)
+        }) + place({
+            key: free for key, free in free_bytes.items() if self._is_avoided(key)
+        })
+        return [self.volumes[key] for key in order]
+
+    def writer(self, hashers: Mapping[str, Hasher]) -> BlobWriter:
+        return BlobWriter(self, self.placement_order(), hashers)
+
+    def can_read(self, blob: Blob) -> bool:
+        """Whether any copy is on a volume this worker has open."""
+        return any(key in self.volumes for key in blob.volume_ids)
+
+    def read_order(self, blob: Blob) -> list[Volume]:
+        """The volumes to read a blob from, spreading reads over its copies."""
+        available = [key for key in blob.volume_ids if key in self.volumes]
+        order = _RANDOM.sample(available, len(available))
+        order.sort(key=self._is_avoided)
+        return [self.volumes[key] for key in order]
+
+    def delete(self, blob_id: str, volume_ids: Iterable[int]) -> None:
+        """Delete what can be reached. `fsck` finds any copy left on a dead disk."""
+        for volume_id in volume_ids:
+            volume = self.volumes.get(volume_id)
+            if volume is None:
+                continue
+            try:
+                volume.delete(blob_id)
+            except OSError as error:
+                self.report_failure(volume_id, error)
+
+    def sweep_tmp(self, now: float) -> int:
+        removed = 0
+        for volume_id, volume in self.volumes.items():
+            try:
+                removed += volume.sweep_tmp(now)
+            except OSError as error:
+                self.report_failure(volume_id, error)
+        return removed
+
+
+def create_volume(root: Path, identity: str) -> None:
+    """Lay out a new volume in an existing directory that holds nothing else.
+
+    The directory must exist already: creating it would hide an unmounted disk.
+    """
+    if not root.is_dir():
+        message = f"{root} isn't a directory"
+        raise VolumeError(message)
+    if {entry.name for entry in root.iterdir()} - _ALLOWED_IN_A_NEW_VOLUME:
+        message = f"{root} isn't empty"
+        raise VolumeError(message)
+    root.chmod(0o700)
+    volume = Volume(root, _ID_NOT_ASSIGNED_YET)
+    volume.initialize()
+    volume.claim(identity)
+
+
+def open_volume(data_dir: Path, record: VolumeRecord) -> Volume:
+    """Check a volume is the one the metadata expects.
+
+    An unmounted disk leaves an empty directory, and writing there would fill the disk
+    underneath and hide those blobs later. So a volume must identify itself.
+    """
+    root = data_dir / record.path
+    try:
+        identity = json.loads((root / VOLUME_FILE).read_text("utf-8"))["identity"]
+    except (OSError, ValueError, KeyError) as error:
+        message = (
+            f"volume {record.volume_id} at {root} can't be used, "
+            f"is its disk mounted? {error}"
+        )
+        raise VolumeError(message) from error
+    if identity != record.identity:
+        message = (
+            f"{root} holds volume {identity}, but the metadata expects "
+            f"{record.identity} there: these belong to different stores"
+        )
+        raise VolumeError(message)
+    return Volume(root, record.volume_id)
+
+
+def open_store(
+    data_dir: Path,
+    records: Sequence[VolumeRecord],
+    copies: int = 1,
+    *,
+    allow_degraded: bool = False,
+) -> BlobStore:
+    """Open every volume, or, if allowed, those that are available."""
+    volumes = {}
+    for record in records:
+        try:
+            volumes[record.volume_id] = open_volume(data_dir, record)
+        except VolumeError as error:
+            if not allow_degraded:
+                raise
+            logger.warning("Starting without a volume: %s", error)
+    if not volumes:
+        message = f"none of the volumes of {data_dir} can be used"
+        raise VolumeError(message)
+    draining = frozenset(record.volume_id for record in records if record.draining)
+    return BlobStore(volumes, copies, draining)
+
+
+class _Copy:
+    """One volume's file of a blob being written."""
+
+    def __init__(self, volume: Volume, blob_id: str) -> None:
+        self.volume = volume
+        self.final_path = volume.path(blob_id)
+        # Its own name: two workers may repair a blob at once, and both are good.
+        self._tmp_path = volume.tmp_dir / secrets.token_hex(16)
+        self._file = self._tmp_path.open("xb")
+        self._committed = False
+
+    def write(self, data: bytes) -> None:
+        self._file.write(data)
+
+    def commit(self) -> None:
+        self._file.flush()
+        # ponytail: on macOS fsync doesn't flush the drive cache. Linux is the
+        # target; add fcntl F_FULLFSYNC if it is ever deployed on macOS.
+        # `make power` removes this line for its control: keep it as it is.
+        os.fsync(self._file.fileno())
+        self._file.close()
+        self._tmp_path.rename(self.final_path)
+        fsync_directory(self.final_path.parent)
+        self._committed = True
+
+    def discard(self) -> None:
+        """Remove what can be removed. On a dead disk, the temp sweep finishes it."""
+        with suppress(OSError):
+            self._file.close()
+        if not self._committed:
+            with suppress(OSError):
+                self._tmp_path.unlink(missing_ok=True)
+
+
+class CorruptSourceError(Exception):
+    """A copy read back with the wrong bytes, so it can't be the source of another."""
+
+
+class TargetVolumeError(Exception):
+    """The volume receiving a new copy failed; the source may be fine."""
+
+
+async def copy_blob(store: BlobStore, blob: Blob, target: Volume) -> None:
+    """Make another copy of a blob durable on a volume, checking every byte first.
+
+    The caller records it afterwards, so a crash leaves at worst an orphan.
+    """
+    try:
+        copy = _Copy(target, blob.blob_id)
+    except OSError as error:
+        raise TargetVolumeError(str(error)) from error
+    digest = hashlib.md5(usedforsecurity=False)
+
+    def write(data: bytes) -> None:
+        digest.update(data)
+        try:
+            copy.write(data)
+        except OSError as error:
+            raise TargetVolumeError(str(error)) from error
+
+    try:
+        whole = BlobSlice(blob, 0, blob.size)
+        async with aclosing(read_slices(store, [whole])) as chunks:
+            async for data in chunks:
+                await asyncio.to_thread(write, data)
+        if digest.hexdigest() != blob.md5:
+            message = f"Blob {blob.blob_id} doesn't match its MD5"
+            raise CorruptSourceError(message)
+        try:
+            await asyncio.to_thread(copy.commit)
+        except OSError as error:
+            raise TargetVolumeError(str(error)) from error
+    finally:
+        copy.discard()
+
+
+class BlobWriter:
+    """One blob being written, to as many volumes as the store keeps copies on.
+
+    It only becomes durable and named on `commit`. A copy that fails is dropped, and
+    the blob survives as long as one remains: refusing the write wouldn't make the
+    remaining copies any safer, and repair restores the number later.
+    """
+
+    def __init__(
+        self, store: BlobStore, volumes: Iterable[Volume], hashers: Mapping[str, Hasher]
+    ) -> None:
+        self._store = store
+        self._hashers = hashers
+        self.blob_id = secrets.token_hex(16)
+        self._copies: list[_Copy] = []
+        self._dropped: list[_Copy] = []
+        self._last_error: OSError | None = None
+        self.size = 0
+        for volume in volumes:
+            if len(self._copies) == store.copies:
+                break
+            try:
+                self._copies.append(_Copy(volume, self.blob_id))
+            except OSError as error:
+                self._record_failure(volume, error)
+        self._require_a_copy()
+
+    def _record_failure(self, volume: Volume, error: OSError) -> None:
+        self._last_error = error
+        self._store.report_failure(volume.volume_id, error)
+
+    def _require_a_copy(self) -> None:
+        if not self._copies:
+            message = f"no volume could store blob {self.blob_id}"
+            raise OSError(message) from self._last_error
+
+    def _drop(self, copy: _Copy, error: OSError) -> None:
+        self._record_failure(copy.volume, error)
+        self._copies.remove(copy)
+        self._dropped.append(copy)
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        for copy in (*self._copies, *self._dropped):
+            copy.discard()
+
+    def write(self, data: bytes) -> None:
+        """Hash once and write every copy. Runs on a worker thread."""
+        for hasher in self._hashers.values():
+            hasher.update(data)
+        for copy in list(self._copies):
+            try:
+                copy.write(data)
+            except OSError as error:
+                self._drop(copy, error)
+        self._require_a_copy()
+        self.size += len(data)
+
+    def digests(self) -> dict[str, bytes]:
+        return {name: hasher.digest() for name, hasher in self._hashers.items()}
+
+    def commit(self) -> Blob:
+        """Make the blob durable under its final name. Runs on a worker thread.
+
+        Copies commit in parallel, so that a PUT waits for the slowest disk and not
+        for the sum of them.
+        """
+        copies = list(self._copies)
+        if len(copies) == 1:
+            errors = [_commit_copy(copies[0])]
+        else:
+            with ThreadPoolExecutor(len(copies)) as pool:
+                errors = list(pool.map(_commit_copy, copies))
+        for copy, error in zip(copies, errors, strict=True):
+            if error is not None:
+                self._drop(copy, error)
+        self._require_a_copy()
+        return Blob(
+            self.blob_id,
+            self.size,
+            self._hashers["md5"].digest().hex(),
+            tuple(copy.volume.volume_id for copy in self._copies),
+        )
+
+
+def _commit_copy(copy: _Copy) -> OSError | None:
+    try:
+        copy.commit()
+    except OSError as error:
+        return error
+    return None
+
+
+async def write_stream(
+    writer: BlobWriter, chunks: AsyncIterable[bytes], max_bytes: int
+) -> None:
+    """Feed a stream to a writer, one thread-pool trip per buffer of data."""
+    buffer = bytearray()
+    received = 0
+    async for chunk in chunks:
+        received += len(chunk)
+        if received > max_bytes:
+            raise S3Error(errors.ENTITY_TOO_LARGE)
+        buffer += chunk
+        if len(buffer) >= BUFFER_BYTES:
+            await asyncio.to_thread(writer.write, bytes(buffer))
+            buffer.clear()
+    if buffer:
+        await asyncio.to_thread(writer.write, bytes(buffer))
+
+
+class _SliceReader:
+    """Reads a byte range of a blob, moving to the next copy when one fails.
+
+    Moving is always safe at the same offset: blobs are immutable, so their copies are
+    identical.
+    """
+
+    def __init__(self, store: BlobStore, blob_slice: BlobSlice) -> None:
+        self._store = store
+        self._blob = blob_slice.blob
+        self._position = blob_slice.offset
+        self._remaining = blob_slice.length
+        self._volumes = store.read_order(self._blob)
+        self._volume: Volume | None = None
+        self._descriptor: int | None = None
+        self._failure: OSError | None = None
+
+    async def _open_next_copy(self) -> int:
+        while self._volumes:
+            volume = self._volumes.pop(0)
+            path = volume.path(self._blob.blob_id)
+            try:
+                descriptor = await asyncio.to_thread(os.open, path, os.O_RDONLY)
+            except FileNotFoundError as error:
+                self._failure = error  # The blob is gone. The disk may be fine.
+            except OSError as error:
+                self._failure = error
+                self._store.report_failure(volume.volume_id, error)
+            else:
+                self._volume = volume
+                self._descriptor = descriptor
+                return descriptor
+        message = f"No copy of blob {self._blob.blob_id} could be read: {self._failure}"
+        raise OSError(message) from self._failure
+
+    async def _read_buffer(self, descriptor: int) -> bytes:
+        size = min(BUFFER_BYTES, self._remaining)
+        data = await asyncio.to_thread(os.pread, descriptor, size, self._position)
+        if not data:
+            message = f"Blob {self._blob.blob_id} is shorter than its metadata."
+            raise OSError(message)
+        return data
+
+    async def next_buffer(self) -> bytes:
+        """Return the next buffer of data, empty at the end."""
+        while self._remaining > 0:
+            descriptor = self._descriptor
+            if descriptor is None:
+                descriptor = await self._open_next_copy()
+            try:
+                data = await self._read_buffer(descriptor)
+            except OSError as error:
+                self._failure = error
+                if self._volume is not None:
+                    self._store.report_failure(self._volume.volume_id, error)
+                self.close()
+                continue
+            self._position += len(data)
+            self._remaining -= len(data)
+            return data
+        return b""
+
+    def close(self) -> None:
+        if self._descriptor is not None:
+            os.close(self._descriptor)
+            self._descriptor = None
+
+
+async def read_slices(
+    store: BlobStore, slices: Iterable[BlobSlice]
+) -> AsyncGenerator[bytes]:
+    """Stream byte ranges of blobs, opening each blob only when it is reached."""
+    for blob_slice in slices:
+        reader = _SliceReader(store, blob_slice)
+        try:
+            while data := await reader.next_buffer():
+                yield data
+        finally:
+            reader.close()