Repository navigation
fix(store): the local object store writes atomically and keeps content types #9
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
97a7770
e91eea9
1271813
5896423
a2d2793
d46296d
1b8b548
35ea9bc
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3,24 +3,48 @@ | |
| from __future__ import annotations | ||
|
|
||
| import contextlib | ||
| import errno | ||
| import hashlib | ||
| import json | ||
| import logging | ||
| import os | ||
| import secrets | ||
| import shutil | ||
| import tempfile | ||
| import time | ||
| from collections.abc import Iterator | ||
| from datetime import datetime, timezone | ||
| from pathlib import Path | ||
| from typing import BinaryIO | ||
|
|
||
| from agent_env.store.base import ObjectAlreadyExistsError, ObjectNotFoundError | ||
|
|
||
| if os.name != "nt": | ||
| import fcntl | ||
| from agent_env.store.local_state import ensure_state_dir | ||
| from agent_env.store.object_store.object_store import DEFAULT_CONTENT_TYPE, ObjectMetadata, ObjectStore | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| _META_DIR = ".agentenv-meta" # each object's content type, at the object's own key | ||
| _STAGING_DIR = ".agentenv-tmp" # files being written, inside the root so a rename into place is atomic | ||
| _RESERVED = (".gitignore", _META_DIR, _STAGING_DIR) | ||
| # A filesystem without hard links: a no-overwrite write relies on the key's lock instead. | ||
| _NO_HARD_LINKS = frozenset({errno.EPERM, errno.ENOTSUP, errno.EOPNOTSUPP, errno.EXDEV, errno.EMLINK}) | ||
| _LOCK_STRIPES = 256 # keys share this many lock files, so locking leaves no file behind per key | ||
| _STAGED_PREFIX = "staged-" | ||
| _STAGED_GRACE_SECONDS = 60 # a staged file younger than this may not be locked by its writer yet | ||
|
|
||
|
|
||
| class LocalFilesystemObjectStore(ObjectStore): | ||
| """ObjectStore backed by a local directory tree; content types are not persisted. | ||
| """ObjectStore backed by a local directory tree. | ||
|
|
||
| The root-level ``.gitignore`` key is reserved for the ignore rule the store writes.""" | ||
| A write lands whole or not at all, and keeps any content type but the default. The root-level ``.gitignore``, | ||
| ``.agentenv-meta`` and ``.agentenv-tmp`` keys are reserved for the store's own files.""" | ||
|
|
||
| def __init__(self, root: str) -> None: | ||
| self._root = Path(root) | ||
| self._swept = False | ||
|
|
||
| @property | ||
| def root(self) -> Path: | ||
|
|
@@ -36,18 +60,18 @@ def put( | |
| path = self._resolve(key) | ||
| if path.exists() and not allow_overwrite: | ||
| raise ObjectAlreadyExistsError(f"Object already exists at {path}.") | ||
| ensure_state_dir(self._root) | ||
| path.parent.mkdir(parents=True, exist_ok=True) | ||
| path.write_bytes(data) | ||
| with self._staged() as staged: | ||
| staged.write_bytes(data) | ||
| self._commit(staged, path, content_type, allow_overwrite=allow_overwrite) | ||
| return self._to_url(path) | ||
|
|
||
| def put_file(self, key: str, file_path: str, content_type: str = DEFAULT_CONTENT_TYPE) -> str: | ||
| path = self._resolve(key) | ||
| if path.exists(): | ||
| raise ObjectAlreadyExistsError(f"Object already exists at {path}.") | ||
| ensure_state_dir(self._root) | ||
| path.parent.mkdir(parents=True, exist_ok=True) | ||
| shutil.copyfile(file_path, path) | ||
| with self._staged() as staged: | ||
| shutil.copyfile(file_path, staged) | ||
| self._commit(staged, path, content_type, allow_overwrite=False) | ||
| return self._to_url(path) | ||
|
|
||
| def put_file_at(self, object_url: str, file_path: str, content_type: str = DEFAULT_CONTENT_TYPE) -> str: | ||
|
|
@@ -74,34 +98,30 @@ def open(self, object_url: str) -> BinaryIO: | |
|
|
||
| def get_object_metadata(self, key: str) -> ObjectMetadata | None: | ||
| path = self._resolve(key) | ||
| if not path.is_file(): | ||
| return None | ||
| st = path.stat() | ||
| return ObjectMetadata(content_type=None, size=st.st_size, last_modified=datetime.fromtimestamp(st.st_mtime, tz=timezone.utc)) | ||
| with self._locked(path, shared=True): # never between another write's bytes and its type | ||
| if not path.is_file(): | ||
| return None | ||
| st = path.stat() | ||
| content_type = self._read_content_type(path, st) | ||
| return ObjectMetadata( | ||
| content_type=content_type, | ||
| size=st.st_size, | ||
| last_modified=datetime.fromtimestamp(st.st_mtime, tz=timezone.utc), | ||
| ) | ||
|
|
||
| def get_object_metadata_at(self, object_url: str) -> ObjectMetadata | None: | ||
| return self.get_object_metadata(self.get_object_key(object_url)) | ||
|
|
||
| def list(self, prefix: str) -> list[str]: | ||
| keys = [] | ||
| for path in self._root.rglob("*"): | ||
| if not path.is_file() or path == self._root / ".gitignore": | ||
| continue | ||
| key = path.relative_to(self._root).as_posix() | ||
| if key.startswith(prefix): | ||
| keys.append(key) | ||
| return keys | ||
| root = self._root.resolve() | ||
| return [key for key in (p.relative_to(root).as_posix() for p in self._objects_below(root)) if key.startswith(prefix)] | ||
|
|
||
| def list_at(self, url_prefix: str) -> list[str]: | ||
| base = url_prefix if url_prefix.endswith("/") else url_prefix + "/" | ||
| dir_path = self._resolve(self.get_object_key(url_prefix)) | ||
| if not dir_path.is_dir(): | ||
| return [] | ||
| return [ | ||
| base + p.relative_to(dir_path).as_posix() | ||
| for p in sorted(dir_path.rglob("*")) | ||
| if p.is_file() and p != self._root.resolve() / ".gitignore" | ||
| ] | ||
| return [base + p.relative_to(dir_path).as_posix() for p in sorted(self._objects_below(dir_path))] | ||
|
|
||
| def object_url(self, key: str) -> str: | ||
| return self._to_url(self._resolve(key)) | ||
|
|
@@ -119,13 +139,156 @@ def get_object_key(self, object_url: str) -> str: | |
| return path.relative_to(root).as_posix() | ||
|
|
||
| def _resolve(self, key: str) -> Path: | ||
| root = self._root.resolve() | ||
| resolved = (self._root / key).resolve() | ||
| if not resolved.is_relative_to(self._root.resolve()): | ||
| if not resolved.is_relative_to(root): | ||
| raise ValueError(f"Key {key!r} escapes the store root.") | ||
| if resolved == self._root.resolve() / ".gitignore": | ||
| raise ValueError(f"Key {key!r} is reserved: the store keeps its .gitignore there.") | ||
| parts = resolved.relative_to(root).parts | ||
| if parts and parts[0] in _RESERVED: | ||
| raise ValueError(f"Key {key!r} is reserved: the store keeps its own {parts[0]} there.") | ||
| return resolved | ||
|
|
||
| def _objects_below(self, dir_path: Path) -> Iterator[Path]: | ||
| """Every object file below ``dir_path``, leaving out the store's own files at the root.""" | ||
| root = self._root.resolve() | ||
| if not dir_path.is_dir(): | ||
| return | ||
| for entry in dir_path.iterdir(): | ||
| if dir_path == root and entry.name in _RESERVED: | ||
| continue | ||
| if entry.is_file(): | ||
| yield entry | ||
| elif entry.is_dir(): | ||
| yield from (p for p in entry.rglob("*") if p.is_file()) | ||
|
|
||
| @contextlib.contextmanager | ||
| def _staged(self) -> Iterator[Path]: | ||
| """An empty file to write an object into before it is committed; removed if it never is.""" | ||
| ensure_state_dir(self._root) | ||
| staging = self._root / _STAGING_DIR | ||
| staging.mkdir(exist_ok=True) | ||
| if not self._swept: | ||
| self._swept = True | ||
| _sweep(staging) | ||
| fd, name = tempfile.mkstemp(dir=staging, prefix=_STAGED_PREFIX) | ||
| staged = Path(name) | ||
| try: | ||
| if os.name == "nt": # an open file cannot be renamed there, and there is no flock to hold | ||
| os.close(fd) | ||
| fd = None | ||
| else: | ||
| fcntl.flock(fd, fcntl.LOCK_EX) # held while the write lives, so a sweep never takes its file | ||
| staged.chmod(0o644) | ||
| yield staged | ||
| finally: | ||
| if fd is not None: | ||
| os.close(fd) | ||
| staged.unlink(missing_ok=True) | ||
|
|
||
| def _commit(self, staged: Path, path: Path, content_type: str, *, allow_overwrite: bool) -> None: | ||
| """Put a staged file in place at ``path`` and record its content type, both under the key's lock, so no | ||
| other write to the key lands in between and no reader sees one write's type with another's bytes. An | ||
| overwrite drops the old type first, so one cut short leaves the type unknown, never wrong. | ||
|
|
||
| The type is also recorded with the identity of the file it describes, which is what keeps them paired | ||
| where there is no lock (Windows, whose NTFS keeps times fine enough to tell writes apart).""" | ||
| identity = _stamp(staged) | ||
| meta = self._meta_path(path) | ||
| path.parent.mkdir(parents=True, exist_ok=True) | ||
| with self._locked(path): | ||
| if allow_overwrite: | ||
| # The old type is set aside while the bytes land: a write cut short reads as unknown, never as the | ||
| # old type, and one that fails before its bytes land puts the old type back. | ||
| aside = self._set_aside(meta) | ||
| try: | ||
| os.replace(staged, path) | ||
| except BaseException: | ||
| if aside is not None: | ||
| with contextlib.suppress(FileNotFoundError): # swept meanwhile: the type reads as unknown | ||
| os.replace(aside, meta) | ||
| raise | ||
| if aside is not None: | ||
| with contextlib.suppress(OSError): # the bytes have landed; a sweep takes what is left | ||
| aside.unlink(missing_ok=True) | ||
|
Comment on lines
+211
to
+212
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. If removing a saved type record fails after an overwrite, the write succeeds but leaves that record in Prompt To Fix With AIThis is a comment left during a code review.
Path: src/agent_env/store/object_store/local_object_store.py
Line: 211-212
Comment:
**Old type records linger**
If removing a saved type record fails after an overwrite, the write succeeds but leaves that record in `.agentenv-tmp`. This store sweeps only on its first write, so a long-lived store will not remove it. Windows does not sweep at all. Repeated cleanup failures can leave old records behind and use disk space.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Leaving this as is. A record is left behind only when deleting a file the store itself just renamed into its own staging directory fails. That means the filesystem is refusing deletes in a directory the store otherwise writes freely. The cost is one small hidden file per failure, and it never makes a type wrong. The next store instance on a POSIX filesystem sweeps it. On Windows it stays until removed by hand. Sweeping on every write, or retrying in the background, would add cost to the common path for a fault that already breaks writes elsewhere. On the summary line: an overwrite whose rename fails puts the old type back. The type is lost (read as unknown, never wrong) only if the set-aside record also vanished in the meantime, which requires another process's sweep while this commit was paused for over a minute. |
||
| else: | ||
| try: | ||
| os.link(staged, path) # exclusive even against a writer that takes no lock | ||
| except FileExistsError: | ||
| raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None | ||
| except OSError as e: | ||
| if e.errno not in _NO_HARD_LINKS: | ||
| raise | ||
| # The lock makes the check and the rename one step; on Windows, rename refuses an existing target. | ||
| if path.exists(): | ||
| raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None | ||
| try: | ||
| (os.rename if os.name == "nt" else os.replace)(staged, path) | ||
| except FileExistsError: | ||
| raise ObjectAlreadyExistsError(f"Object already exists at {path}.") from None | ||
| # The default type is what a writer that named none gets; it reads back as unknown, so readers | ||
| # still guess from the name as they did before types were kept. | ||
| try: | ||
| if content_type == DEFAULT_CONTENT_TYPE: | ||
| meta.unlink(missing_ok=True) | ||
| return | ||
| meta.parent.mkdir(parents=True, exist_ok=True) | ||
| with self._staged() as staged_meta: | ||
| staged_meta.write_text(json.dumps({"content_type": content_type, "object": identity})) | ||
| os.replace(staged_meta, meta) | ||
| except OSError as e: # the object is in place and its type reads back as unknown | ||
| logger.warning("Could not record the content type of %s: %s", path, e) | ||
|
|
||
| @contextlib.contextmanager | ||
| def _locked(self, path: Path, *, shared: bool = False) -> Iterator[None]: | ||
| """The key's lock, exclusive to write it and shared to read its type: a flock on one of a fixed set of | ||
| files in the staging directory, which the kernel drops if its holder dies. None on Windows.""" | ||
| if os.name == "nt": | ||
| yield | ||
| return | ||
| try: | ||
| f = self._lock_path(path).open("r" if shared else "a") | ||
| except OSError: | ||
| if not shared: | ||
| raise | ||
| f = None # no write has taken this lock yet, or the store is read-only: nothing to wait for | ||
| if f is None: | ||
| yield | ||
| return | ||
| with f: | ||
| fcntl.flock(f, fcntl.LOCK_SH if shared else fcntl.LOCK_EX) | ||
| yield | ||
|
|
||
| def _set_aside(self, meta: Path) -> Path | None: | ||
| """Move a type record into the staging directory, fresh enough that no sweep takes it; None when there is | ||
| none.""" | ||
| aside = self._root / _STAGING_DIR / f"{_STAGED_PREFIX}aside-{secrets.token_hex(8)}" | ||
| try: | ||
| os.replace(meta, aside) | ||
| except FileNotFoundError: | ||
| return None | ||
| with contextlib.suppress(OSError): | ||
| os.utime(aside) | ||
| return aside | ||
|
|
||
| def _lock_path(self, path: Path) -> Path: | ||
| stripe = int(hashlib.sha256(str(path).encode()).hexdigest()[:8], 16) % _LOCK_STRIPES | ||
| return self._root / _STAGING_DIR / f"lock-{stripe:02x}" | ||
|
greptile-apps[bot] marked this conversation as resolved.
|
||
|
|
||
| def _read_content_type(self, path: Path, st: os.stat_result | None = None) -> str | None: | ||
| """The content type recorded for the object at ``path``, if it was recorded for this very file; None | ||
| otherwise, as for one written before types were kept.""" | ||
| try: | ||
| recorded = json.loads(self._meta_path(path).read_text()) | ||
| if recorded.get("object") != _identity(st or path.stat()): | ||
| return None | ||
| return recorded.get("content_type") | ||
| except (OSError, ValueError, AttributeError): | ||
| return None | ||
|
|
||
| def _meta_path(self, path: Path) -> Path: | ||
| root = self._root.resolve() | ||
| return root / _META_DIR / path.relative_to(root) | ||
|
|
||
| @staticmethod | ||
| def _to_url(path: Path) -> str: | ||
| return f"file://{path.resolve()}" | ||
|
|
@@ -135,6 +298,35 @@ def _from_url(object_url: str) -> str: | |
| return object_url[len("file://"):] if object_url.startswith("file://") else object_url | ||
|
|
||
|
|
||
| def _stamp(staged: Path) -> list[int]: | ||
| """Give a staged file a modification time no other write shares (now, to the millisecond, plus a random | ||
| remainder) and return its identity. A rename keeps it, and so does a copy that keeps modification times.""" | ||
| now_ns = time.time_ns() | ||
| stamp = now_ns - now_ns % 1_000_000 + secrets.randbelow(1_000_000) | ||
| os.utime(staged, ns=(stamp, stamp)) | ||
| return _identity(staged.stat()) | ||
|
|
||
|
|
||
| def _sweep(staging: Path) -> None: | ||
| """Remove files staged by writes that died: each live write holds a lock on its file, which the kernel drops | ||
| when the writer exits, so a file no one holds will never be committed. Not on Windows, which has no flock.""" | ||
| if os.name == "nt": | ||
| return | ||
| cutoff = time.time() - _STAGED_GRACE_SECONDS | ||
| for staged in staging.glob(f"{_STAGED_PREFIX}*"): | ||
| with contextlib.suppress(OSError): # BlockingIOError, the one that matters: a live write holds it | ||
| if staged.stat().st_mtime > cutoff: | ||
| continue | ||
| with staged.open("rb") as f: | ||
| fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB) | ||
| staged.unlink() | ||
|
|
||
|
|
||
| def _identity(st: os.stat_result) -> list[int]: | ||
| return [st.st_size, st.st_mtime_ns] | ||
|
|
||
|
|
||
|
|
||
| @contextlib.contextmanager | ||
| def _missing_as_not_found(object_url: str) -> Iterator[None]: | ||
| try: | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.