file_lock_utils.py

# file_lock_utils.py
# Cross-process file locking utilities for parallel session generation.
#
# Wraps filelock to provide simple helpers for the common pattern of
# load-JSON -> modify -> save-JSON that appears throughout the codebase.

import os
import json
from filelock import FileLock

# Cache of FileLock instances keyed by file path.
# Each JSON file gets a companion .lock file in the same directory.
_locks: dict[str, FileLock] = {}

# Default timeout (seconds) for acquiring a lock.
LOCK_TIMEOUT = 30


def get_lock(path: str, timeout: float = LOCK_TIMEOUT) -> FileLock:
    """
    Get (or create) a FileLock for the given file path.

    The lock file is placed alongside the target file with a '.lock' suffix.
    Lock instances are cached so repeated calls for the same path reuse the
    same FileLock object (important for thread safety within a single process).

    Args:
        path: Absolute path to the file being protected.
        timeout: Seconds to wait for the lock before raising Timeout.

    Returns:
        A FileLock instance (use as a context manager).
    """
    if path not in _locks:
        lock_path = path + ".lock"
        _locks[path] = FileLock(lock_path, timeout=timeout)
    return _locks[path]


def locked_json_read(path: str) -> dict | list:
    """
    Read a JSON file under a file lock.

    Args:
        path: Path to the JSON file.

    Returns:
        Parsed JSON content (dict or list).
    """
    with get_lock(path):
        if not os.path.exists(path):
            return {}
        with open(path, "r") as f:
            return json.load(f)


def locked_json_update(path: str, update_fn, default=None):
    """
    Atomically read-modify-write a JSON file under a file lock.

    Usage:
        def add_item(data):
            data["items"].append("new_item")
            return data

        locked_json_update("state.json", add_item, default={"items": []})

    Args:
        path: Path to the JSON file.
        update_fn: Callable(data) -> data.  Receives the current file content
                   (or *default* if the file doesn't exist) and must return
                   the updated content to be written back.
        default: Value to use if the file does not yet exist.
                 Defaults to an empty dict.

    Returns:
        The updated data (as returned by update_fn).
    """
    if default is None:
        default = {}

    with get_lock(path):
        # Load
        if os.path.exists(path):
            try:
                with open(path, "r") as f:
                    data = json.load(f)
            except (json.JSONDecodeError, IOError):
                data = default if callable(default) else (default.copy() if hasattr(default, 'copy') else default)
        else:
            data = default if callable(default) else (default.copy() if hasattr(default, 'copy') else default)

        # Modify
        data = update_fn(data)

        # Save
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        with open(path, "w") as f:
            json.dump(data, f, indent=2)

    return data