92 lines
3.0 KiB
Python
92 lines
3.0 KiB
Python
# -*- coding: UTF-8 -*-
|
|
"""Helpers for cross-process safe pickle persistence."""
|
|
|
|
import os
|
|
import pickle
|
|
import tempfile
|
|
import time
|
|
from contextlib import contextmanager
|
|
from typing import Callable, TypeVar
|
|
|
|
T = TypeVar("T")
|
|
|
|
try: # pragma: no cover - platform specific
|
|
import fcntl
|
|
except ImportError: # pragma: no cover - platform specific
|
|
fcntl = None
|
|
|
|
try: # pragma: no cover - platform specific
|
|
import msvcrt
|
|
except ImportError: # pragma: no cover - platform specific
|
|
msvcrt = None
|
|
|
|
|
|
def _acquire_platform_lock(lock_file, *, timeout_seconds: float) -> None:
|
|
deadline = time.monotonic() + max(0.1, float(timeout_seconds or 0.0))
|
|
while True:
|
|
try:
|
|
if fcntl is not None:
|
|
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
elif msvcrt is not None:
|
|
msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1)
|
|
return
|
|
except (BlockingIOError, OSError):
|
|
if time.monotonic() >= deadline:
|
|
raise TimeoutError(f"lock acquire timed out for {lock_file.name}")
|
|
time.sleep(0.05)
|
|
|
|
|
|
def _release_platform_lock(lock_file) -> None:
|
|
try:
|
|
if fcntl is not None:
|
|
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
|
|
elif msvcrt is not None:
|
|
lock_file.seek(0)
|
|
msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
@contextmanager
|
|
def _locked_path(path: str, *, timeout_seconds: float = 10.0):
|
|
normalized_path = os.path.abspath(path)
|
|
directory = os.path.dirname(normalized_path) or "."
|
|
os.makedirs(directory, exist_ok=True)
|
|
lock_path = f"{normalized_path}.lock"
|
|
with open(lock_path, "a+b") as lock_file:
|
|
_acquire_platform_lock(lock_file, timeout_seconds=timeout_seconds)
|
|
try:
|
|
yield normalized_path
|
|
finally:
|
|
_release_platform_lock(lock_file)
|
|
|
|
|
|
def save_pickle_atomic(path: str, value, *, timeout_seconds: float = 10.0) -> None:
|
|
with _locked_path(path, timeout_seconds=timeout_seconds) as normalized_path:
|
|
directory = os.path.dirname(normalized_path) or "."
|
|
fd, temp_path = tempfile.mkstemp(
|
|
prefix=f".{os.path.basename(normalized_path)}.",
|
|
suffix=".tmp",
|
|
dir=directory,
|
|
)
|
|
try:
|
|
with os.fdopen(fd, "wb") as temp_file:
|
|
pickle.dump(value, temp_file)
|
|
temp_file.flush()
|
|
os.fsync(temp_file.fileno())
|
|
os.replace(temp_path, normalized_path)
|
|
finally:
|
|
if os.path.exists(temp_path):
|
|
try:
|
|
os.remove(temp_path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def load_pickle_locked(path: str, *, default_factory: Callable[[], T], timeout_seconds: float = 10.0) -> T:
|
|
with _locked_path(path, timeout_seconds=timeout_seconds) as normalized_path:
|
|
if not os.path.exists(normalized_path):
|
|
return default_factory()
|
|
with open(normalized_path, "rb") as source:
|
|
return pickle.load(source)
|