import math
import time
from dataclasses import dataclass
from logging import getLogger
from multiprocessing import Pipe, Queue
from multiprocessing.connection import Connection
from os import PathLike
from typing import Any, Callable
from edri.dataclass.event import Event
from edri.dataclass.persistence_failure_policy import PersistenceFailurePolicy
from edri.utility.function import install_graceful_sigterm
from ._store_persistence import _SQLitePersistence
from edri.abstract import ManagerBase
from edri.events.edri.store import (
Set, Get, Delete, GetCallback,
Increment, CompareAndSet, SetIfAbsent, StorePurge,
)
__all__ = ["Callback", "PersistenceFailurePolicy"]
[docs]
@dataclass
class Callback:
condition: Callable[[Any], bool] | None
class Store(ManagerBase):
"""
A key-value store that processes set/get/delete requests plus change callbacks,
per-key TTL expiry, and lock-free atomic operations.
Because the Store is a single Manager that processes events sequentially, its
read-modify-write operations (:class:`Increment`, :class:`CompareAndSet`,
:class:`SetIfAbsent`) are atomic by construction — no locks are required. This
makes it a natural coordination primitive for an appliance: shared counters
(rate limiting), version tokens (optimistic concurrency), and TTL'd locks
(leader election).
Args:
router_queue: Queue for Router-bound events.
path: Optional trusted SQLite file to restore and write asynchronously.
With None, no persistence resources are created.
failure_policy: BUFFER keeps serving from RAM while failed writes retry
in an unbounded queue.
startup_timeout: Seconds to wait for persistent restoration, or None
to wait without a deadline. Ignored when path is None.
Attributes:
store (dict[str, Any]): The key-value data.
expiry (dict[str, float]): Per-key expiry deadlines (``time.monotonic()``).
A key absent from this map never expires.
callbacks (dict[str, list[Callback]]): Registered change callbacks per key.
"""
def __init__(
self, router_queue: "Queue[Event]", path: str | PathLike[str] | None = None,
*, failure_policy: PersistenceFailurePolicy = PersistenceFailurePolicy.BUFFER,
startup_timeout: float | None = 30.0,
) -> None:
if failure_policy is not PersistenceFailurePolicy.BUFFER:
raise ValueError("Only BUFFER persistence failure policy is supported")
if startup_timeout is not None and (
not math.isfinite(startup_timeout) or startup_timeout <= 0
):
raise ValueError("startup_timeout must be positive and finite, or None")
super().__init__(router_queue, getLogger(__name__))
self._path = path
self._persistence: _SQLitePersistence | None = None
self._startup_pipe: Connection | None = None
self._startup_timeout = startup_timeout
self.store: dict[str, Any] = {}
self.expiry: dict[str, float] = {}
self.callbacks: dict[str, list[Callback]] = {}
def _initialize_persistence(self) -> None:
if self._path is None or self._persistence is not None:
return
persistence = _SQLitePersistence(self._path, self.logger)
try:
for name, value, expires_at in persistence.load():
remaining = None if expires_at is None else expires_at - time.time()
if remaining is not None and remaining <= 0:
continue
self.store[name] = value
if remaining is not None:
self.expiry[name] = time.monotonic() + remaining
persistence.start()
except BaseException:
persistence.close()
raise
self._persistence = persistence
def run(self) -> None:
try:
if self._path is not None:
install_graceful_sigterm()
# Restore before subscribing: EDRI.start_store() normally runs
# before the Router, so startup confirmation cannot depend on it.
self._initialize_persistence()
self._report_startup(None)
super().run()
except BaseException as error:
self._report_startup(f"{type(error).__qualname__}: {error}")
raise
finally:
if self._persistence is not None:
self._persistence.close()
def _report_startup(self, error: str | None) -> None:
pipe = self._startup_pipe
if pipe is None:
return
self._startup_pipe = None
try:
pipe.send(error)
except OSError:
self.logger.warning("Store startup receiver unavailable", exc_info=True)
finally:
pipe.close()
def start(self, router_queue: Any | None = None) -> None:
"""Start the child. Persistent mode waits for restoration and propagates
bootstrap failures as RuntimeError before the application starts routing."""
if self._path is None:
super().start(router_queue)
return
receiver, sender = Pipe(duplex=False)
self._startup_pipe = sender
started = False
try:
super().start(router_queue)
started = True
sender.close()
if not receiver.poll(self._startup_timeout):
raise RuntimeError("Store persistence startup timed out")
try:
error = receiver.recv()
except EOFError as error:
raise RuntimeError("Store exited before persistence startup completed") from error
if error is not None:
raise RuntimeError(f"Store persistence startup failed ({error})")
except BaseException:
if started:
if self.is_alive():
self.terminate()
self.join(1.0)
if self.is_alive():
self.kill()
self.join(1.0)
raise
finally:
receiver.close()
sender.close()
self._startup_pipe = None
# --- internal helpers -------------------------------------------------
def _expired(self, name: str) -> bool:
deadline = self.expiry.get(name)
return deadline is not None and time.monotonic() >= deadline
def _drop_if_expired(self, name: str) -> None:
if self._expired(name):
self.store.pop(name, None)
self.expiry.pop(name, None)
if self._persistence is not None:
self._persistence.delete(name)
self.logger.debug("Key '%s' expired", name)
def _present(self, name: str) -> bool:
"""Return whether ``name`` currently holds a live (non-expired) value,
dropping it lazily if it has expired."""
self._drop_if_expired(name)
return name in self.store
def _set_value(self, name: str, value: Any, ttl: float | None) -> None:
self.store[name] = value
if ttl:
self.expiry[name] = time.monotonic() + ttl
else:
self.expiry.pop(name, None)
if self._persistence is not None:
expires_at = time.time() + ttl if ttl else None
self._persistence.upsert(name, value, expires_at)
def _fire_callbacks(self, name: str, value: Any) -> None:
for callback in self.callbacks.get(name, []):
if callback.condition is None or callback.condition(value):
# `response` is a non-init field (init=False, default_factory) added by the
# @event dataclass_transform, which mypy cannot see through and therefore
# treats as a required positional argument. The runtime call is valid.
event_callback = GetCallback(name=name) # type: ignore[call-arg]
event_callback.response.value = value
self.send_to_router(event_callback)
self.logger.debug("Sending callback item for '%s'", name)
# --- basic operations -------------------------------------------------
def solve_req_set(self, event: Set) -> None:
"""Store ``value`` under ``name`` with an optional ``ttl`` (seconds).
Triggers any matching change callbacks."""
self._set_value(event.name, event.value, event.ttl)
self.logger.debug("Value set for key '%s': %s", event.name, event.value)
self._fire_callbacks(event.name, event.value)
def solve_req_get(self, event: Get) -> None:
"""Retrieve the live value for ``name`` (``None`` if missing or expired)."""
if self._present(event.name):
event.response.value = self.store[event.name]
self.logger.debug("Value retrieved for key '%s': %s", event.name, event.response.value)
else:
event.response.value = None
self.logger.debug("No value found for key '%s'", event.name)
def solve_req_delete(self, event: Delete) -> None:
"""Remove ``name`` from the store; no-op if absent."""
self.expiry.pop(event.name, None)
if event.name in self.store:
del self.store[event.name]
if self._persistence is not None:
self._persistence.delete(event.name)
self.logger.debug("Deleted item with key '%s'", event.name)
else:
self.logger.debug("Attempted to delete non-existent key '%s'", event.name)
def solve_req_get_callback(self, event: GetCallback) -> None:
callbacks = self.callbacks.get(event.name, [])
callbacks.append(Callback(event.condition))
self.callbacks[event.name] = callbacks
self.logger.debug("Callback added for key '%s'", event.name)
# --- atomic operations (single-threaded → lock-free) ------------------
def solve_req_increment(self, event: Increment) -> None:
"""Atomically add ``amount`` (default 1) to a numeric counter, starting
from 0 when the key is missing or expired. Returns the new value."""
base = self.store[event.name] if self._present(event.name) else 0
new_value = base + event.amount
self._set_value(event.name, new_value, event.ttl)
event.response.value = new_value
self.logger.debug("Incremented '%s' by %s -> %s", event.name, event.amount, new_value)
self._fire_callbacks(event.name, new_value)
def solve_req_compare_and_set(self, event: CompareAndSet) -> None:
"""Set ``value`` only if the current live value equals ``expected``
(a missing/expired key compares as ``None``). Returns whether it applied."""
current = self.store[event.name] if self._present(event.name) else None
if current == event.expected:
self._set_value(event.name, event.value, event.ttl)
event.response.success = True
self.logger.debug("CAS on '%s' succeeded (%s -> %s)", event.name, event.expected, event.value)
self._fire_callbacks(event.name, event.value)
else:
event.response.success = False
self.logger.debug("CAS on '%s' failed (expected %s, was %s)", event.name, event.expected, current)
def solve_req_set_if_absent(self, event: SetIfAbsent) -> None:
"""Set ``value`` only if the key is currently absent or expired. Returns
whether it was acquired — the building block for a TTL'd lock (SETNX)."""
if self._present(event.name):
event.response.acquired = False
self.logger.debug("SetIfAbsent on '%s' not acquired (already present)", event.name)
else:
self._set_value(event.name, event.value, event.ttl)
event.response.acquired = True
self.logger.debug("SetIfAbsent on '%s' acquired", event.name)
self._fire_callbacks(event.name, event.value)
def solve_req_store_purge(self, event: StorePurge) -> None:
"""Sweep and drop every expired key (reclaims memory held by keys that
expired but were never read again)."""
expired = [name for name in list(self.expiry) if self._expired(name)]
for name in expired:
self._drop_if_expired(name)
if expired:
self.logger.debug("Purged %d expired key(s)", len(expired))