968 lines
33 KiB
Python
968 lines
33 KiB
Python
"""
|
|
State management module with pluggable backends.
|
|
|
|
This module provides state storage backends for pipeline execution,
|
|
supporting both in-memory storage and Redis for distributed execution.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from collections import deque
|
|
from collections.abc import Callable, Iterable, Iterator
|
|
from contextlib import contextmanager
|
|
from threading import RLock
|
|
from typing import Any, cast
|
|
|
|
import orjson
|
|
import redis
|
|
|
|
from fluksio.flow.namespace import HISTORY, NAMES, Rename
|
|
|
|
#: A node may return a dict keyed by something other than a string, and the
|
|
#: stdlib encoder this replaced turned those keys into strings rather than
|
|
#: refusing them. Two differences remain and are both improvements: a
|
|
#: non-finite number is stored as ``null`` instead of the bare ``NaN`` that is
|
|
#: not JSON at all, and a datetime serialises rather than raising.
|
|
_JSON_OPTS = orjson.OPT_NON_STR_KEYS
|
|
|
|
# A sparkline only means something for a value that can be placed on an axis,
|
|
# so the history keeps those and nothing else. 120 points fill a panel-wide
|
|
# chart while leaving Redis a cache rather than a time-series database.
|
|
#: How long a Redis call may take before it is treated as a failure, and
|
|
#: how long connecting may take. A blocking read must ask for less than the
|
|
#: first of these.
|
|
SOCKET_TIMEOUT_S = 5.0
|
|
CONNECT_TIMEOUT_S = 2.0
|
|
|
|
HISTORY_LIMIT = 120
|
|
#: How many keys a rename moves per round trip, and between two calls of its
|
|
#: ``progress``, which is where the rename journal can stop.
|
|
RENAME_BATCH = 500
|
|
|
|
|
|
def as_number(value: Any) -> float | None:
|
|
"""The plottable form of a value, or None if it cannot be placed on an axis."""
|
|
# A flag is not a measurement, but when it was on is worth seeing, and 0/1
|
|
# is what a step curve is drawn from. bool is an int subclass, so it would
|
|
# fall through the isinstance below either way.
|
|
if isinstance(value, bool):
|
|
return 1.0 if value else 0.0
|
|
if not isinstance(value, (int, float)):
|
|
return None
|
|
return float(value)
|
|
|
|
|
|
def move_keys(client: redis.Redis, moves: list[tuple[str, str]]) -> int:
|
|
"""``RENAMENX`` each ``(key, target)`` in one round trip, which keeps the
|
|
Redis type and the remaining TTL. Returns how many moved.
|
|
|
|
A key gone since it was listed has expired and is skipped. A target that
|
|
exists is refused rather than overwritten: a rename checks that nothing
|
|
lives under the new name before it starts, so one there now is a fault.
|
|
"""
|
|
pipe = client.pipeline(transaction=False)
|
|
for key, target in moves:
|
|
pipe.renamenx(key, target)
|
|
replies = pipe.execute(raise_on_error=False)
|
|
moved = 0
|
|
for (key, target), reply in zip(moves, replies, strict=True):
|
|
if isinstance(reply, redis.ResponseError) and "no such key" in str(reply):
|
|
continue
|
|
if isinstance(reply, Exception):
|
|
raise reply
|
|
if not reply:
|
|
raise RuntimeError(f"Will not move '{key}' onto '{target}', which exists")
|
|
moved += 1
|
|
return moved
|
|
|
|
|
|
class StateBackend(ABC):
|
|
"""
|
|
Abstract base class for state storage backends.
|
|
|
|
Provides a dict-like interface for storing and retrieving state values
|
|
with support for distributed locking.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get(self, key: str, default: Any = None) -> Any:
|
|
"""
|
|
Get a value by key.
|
|
|
|
:param key: The key to retrieve.
|
|
:type key: str
|
|
:param default: Default value if key not found.
|
|
:type default: Any
|
|
:returns: The stored value or default.
|
|
:rtype: Any
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def set(self, key: str, value: Any) -> None:
|
|
"""
|
|
Set a value by key.
|
|
|
|
:param key: The key to set.
|
|
:type key: str
|
|
:param value: The value to store.
|
|
:type value: Any
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def exists(self, key: str) -> bool:
|
|
"""
|
|
Check if a key exists.
|
|
|
|
:param key: The key to check.
|
|
:type key: str
|
|
:returns: True if key exists.
|
|
:rtype: bool
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def delete(self, key: str) -> None:
|
|
"""Forget one key and any history kept for it."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def clear(self) -> None:
|
|
"""Clear all keys in the state."""
|
|
...
|
|
|
|
def message_names(self) -> list[str]:
|
|
"""The message names in state, without the bookkeeping keys.
|
|
|
|
Its own method because the Redis backend can answer it from a set it
|
|
maintains, rather than by scanning a namespace that holds five
|
|
bookkeeping keys for every message.
|
|
"""
|
|
return [k for k in self.keys() if not k.startswith("__")]
|
|
|
|
def close(self) -> None: # noqa: B027
|
|
"""Release whatever the backend holds outside this process.
|
|
|
|
A run builds a state backend of its own and drops it when it
|
|
finishes; on Redis that is a client and its connection pool, which
|
|
nothing was giving back. Concrete rather than abstract: a backend that
|
|
holds nothing has nothing to answer here.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def keys(self) -> list[str]:
|
|
"""
|
|
Get all keys.
|
|
|
|
:returns: List of all keys.
|
|
:rtype: list[str]
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def update(self, mapping: dict[str, Any]) -> None:
|
|
"""
|
|
Update multiple keys atomically.
|
|
|
|
:param mapping: Dict of key-value pairs to set.
|
|
:type mapping: dict[str, Any]
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
@contextmanager
|
|
def lock(self, timeout: float = 10.0) -> Iterator[None]:
|
|
"""
|
|
Acquire a lock for thread-safe operations.
|
|
|
|
:param timeout: Lock timeout in seconds.
|
|
:type timeout: float
|
|
:yields: Lock context.
|
|
"""
|
|
...
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Atomic operations for synchronous node support
|
|
# -------------------------------------------------------------------------
|
|
|
|
@abstractmethod
|
|
def increment(self, key: str) -> int:
|
|
"""
|
|
Atomically increment a counter and return the new value.
|
|
|
|
:param key: The key to increment.
|
|
:type key: str
|
|
:returns: The new value after incrementing.
|
|
:rtype: int
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def get_multi(self, keys: list[str]) -> dict[str, Any]:
|
|
"""
|
|
Get multiple values atomically.
|
|
|
|
:param keys: List of keys to retrieve.
|
|
:type keys: list[str]
|
|
:returns: Dict mapping keys to values (missing keys have None).
|
|
:rtype: dict[str, Any]
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def get_present(self, keys: list[str]) -> dict[str, Any]:
|
|
"""The values of those keys that exist, in one round trip.
|
|
|
|
Unlike :meth:`get_multi`, a key that is absent is left out rather than
|
|
mapped to None — which is what reading a node's inputs needs, since a
|
|
message holding ``null`` is a message the node has, and one that was
|
|
never published is not.
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def increment_multi(self, keys: list[str]) -> None:
|
|
"""Bump several counters at once.
|
|
|
|
The new values are not returned: the only reader compares them against
|
|
what it saw last time and reads them back itself.
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def compare_and_swap_multi(
|
|
self,
|
|
expected: dict[str, Any],
|
|
updates: dict[str, Any],
|
|
) -> bool:
|
|
"""
|
|
Atomically check expected values and apply updates if all match.
|
|
|
|
This is the key primitive for race-condition-free synchronous node execution.
|
|
If any expected value doesn't match the current value, no updates are applied.
|
|
|
|
:param expected: Dict of keys to their expected values.
|
|
:type expected: dict[str, Any]
|
|
:param updates: Dict of keys to their new values (applied only if expected matches).
|
|
:type updates: dict[str, Any]
|
|
:returns: True if the swap succeeded, False if expected values didn't match.
|
|
:rtype: bool
|
|
"""
|
|
...
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Message history
|
|
# -------------------------------------------------------------------------
|
|
|
|
@abstractmethod
|
|
def append_history(
|
|
self, values: dict[str, Any], ts: float, limits: dict[str, int] | None = None
|
|
) -> None:
|
|
"""
|
|
Record freshly published values in each message's capped history.
|
|
|
|
Non-numeric payloads are skipped — the history exists to be plotted.
|
|
|
|
:param values: Message names mapped to the value just published.
|
|
:type values: dict[str, Any]
|
|
:param ts: When they were published.
|
|
:param limits: How many points to keep per message, where a chart asks
|
|
for more than the default.
|
|
:type ts: float
|
|
"""
|
|
...
|
|
|
|
def record(
|
|
self,
|
|
values: dict[str, Any],
|
|
stamps: dict[str, float],
|
|
counters: list[str],
|
|
ts: float,
|
|
limits: dict[str, int] | None = None,
|
|
drop: list[str] | None = None,
|
|
) -> None:
|
|
"""Everything a published value owes state, in one round trip.
|
|
|
|
The value itself, the timestamp beside it, its series and its version
|
|
counter used to be four calls and four round trips; a backend that can
|
|
batch them does so here. This default is those four calls, so a
|
|
backend that cannot batch gains nothing and breaks nothing.
|
|
|
|
:param values: Message names mapped to the value just published.
|
|
:param stamps: The timestamp key of each of those messages.
|
|
:param counters: The version keys to bump.
|
|
:param ts: When they were published.
|
|
:param limits: How many points to keep per message.
|
|
:param drop: Keys to forget in the same trip — a released rate-limit
|
|
hold, which is written and cleared on this same path.
|
|
"""
|
|
for key in drop or ():
|
|
self.delete(key)
|
|
self.update({**values, **stamps})
|
|
self.append_history(values, ts, limits)
|
|
self.increment_multi(counters)
|
|
|
|
@abstractmethod
|
|
def history(self, key: str) -> list[tuple[float, float]]:
|
|
"""
|
|
The recorded ``(timestamp, value)`` pairs of one message, oldest first.
|
|
|
|
:param key: The message name.
|
|
:type key: str
|
|
:returns: At most ``HISTORY_LIMIT`` points; empty if nothing was recorded.
|
|
:rtype: list[tuple[float, float]]
|
|
"""
|
|
...
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Renaming a flow key
|
|
# -------------------------------------------------------------------------
|
|
|
|
@abstractmethod
|
|
def rename_keys(
|
|
self, rename: Rename, progress: Callable[[str], None] | None = None
|
|
) -> int:
|
|
"""Move every key that names something of ``rename.old`` to the new
|
|
name (`Rename.state_key`), with its history.
|
|
|
|
Keys move, values do not: a message's payload and a node's private
|
|
state are its own, whatever names they mention. A key already under
|
|
the new name stays where it is, so a rerun after a crash finishes the
|
|
job. A target that exists is refused rather than overwritten.
|
|
|
|
:param progress: Called after each batch of keys.
|
|
:returns: How many keys moved.
|
|
:raises ValueError: For a bookkeeping key of an unregistered family.
|
|
"""
|
|
...
|
|
|
|
def count(self, rename: Rename) -> int:
|
|
"""How many keys `rename_keys` would move."""
|
|
return sum(rename.state_key(key) is not None for key in self.keys())
|
|
|
|
def __contains__(self, key: str) -> bool:
|
|
return self.exists(key)
|
|
|
|
def __getitem__(self, key: str) -> Any:
|
|
if not self.exists(key):
|
|
raise KeyError(key)
|
|
return self.get(key)
|
|
|
|
def __setitem__(self, key: str, value: Any) -> None:
|
|
self.set(key, value)
|
|
|
|
def __iter__(self) -> Iterator[str]:
|
|
return iter(self.keys())
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.keys())
|
|
|
|
|
|
class MemoryState(StateBackend):
|
|
"""
|
|
In-memory state backend using a thread-safe dict.
|
|
|
|
Suitable for single-process execution and testing.
|
|
|
|
:example:
|
|
>>> state = MemoryState()
|
|
>>> state["key"] = "value"
|
|
>>> state["key"]
|
|
'value'
|
|
"""
|
|
|
|
__slots__ = ("_data", "_history", "_lock")
|
|
|
|
def __init__(self) -> None:
|
|
self._data: dict[str, Any] = {}
|
|
self._history: dict[str, deque[tuple[float, float]]] = {}
|
|
self._lock = RLock() # Reentrant lock for nested access
|
|
|
|
def get(self, key: str, default: Any = None) -> Any:
|
|
with self._lock:
|
|
return self._data.get(key, default)
|
|
|
|
def set(self, key: str, value: Any) -> None:
|
|
with self._lock:
|
|
self._data[key] = value
|
|
|
|
def exists(self, key: str) -> bool:
|
|
with self._lock:
|
|
return key in self._data
|
|
|
|
def delete(self, key: str) -> None:
|
|
with self._lock:
|
|
self._data.pop(key, None)
|
|
self._history.pop(key, None)
|
|
|
|
def clear(self) -> None:
|
|
with self._lock:
|
|
self._data.clear()
|
|
self._history.clear()
|
|
|
|
def keys(self) -> list[str]:
|
|
with self._lock:
|
|
return list(self._data.keys())
|
|
|
|
def update(self, mapping: dict[str, Any]) -> None:
|
|
with self._lock:
|
|
self._data.update(mapping)
|
|
|
|
@contextmanager
|
|
def lock(self, timeout: float = 10.0) -> Iterator[None]:
|
|
"""Acquire the internal lock."""
|
|
acquired = self._lock.acquire(timeout=timeout)
|
|
if not acquired:
|
|
raise TimeoutError("Failed to acquire lock")
|
|
try:
|
|
yield
|
|
finally:
|
|
self._lock.release()
|
|
|
|
def increment(self, key: str) -> int:
|
|
"""Atomically increment a counter and return the new value."""
|
|
with self._lock:
|
|
value: int = self._data.get(key, 0)
|
|
new_value = value + 1
|
|
self._data[key] = new_value
|
|
return new_value
|
|
|
|
def get_multi(self, keys: list[str]) -> dict[str, Any]:
|
|
"""Get multiple values atomically."""
|
|
with self._lock:
|
|
return {k: self._data.get(k) for k in keys}
|
|
|
|
def get_present(self, keys: list[str]) -> dict[str, Any]:
|
|
"""The values of those keys that are there."""
|
|
with self._lock:
|
|
return {k: self._data[k] for k in keys if k in self._data}
|
|
|
|
def increment_multi(self, keys: list[str]) -> None:
|
|
with self._lock:
|
|
for key in keys:
|
|
self._data[key] = self._data.get(key, 0) + 1
|
|
|
|
def compare_and_swap_multi(
|
|
self,
|
|
expected: dict[str, Any],
|
|
updates: dict[str, Any],
|
|
) -> bool:
|
|
"""Atomically check expected values and apply updates if all match."""
|
|
with self._lock:
|
|
# Check all expected values
|
|
for key, exp_value in expected.items():
|
|
actual = self._data.get(key)
|
|
if actual != exp_value:
|
|
return False
|
|
|
|
# All matched, apply updates
|
|
self._data.update(updates)
|
|
return True
|
|
|
|
def append_history(
|
|
self, values: dict[str, Any], ts: float, limits: dict[str, int] | None = None
|
|
) -> None:
|
|
"""Append the numeric values to their message's bounded series."""
|
|
with self._lock:
|
|
for key, value in values.items():
|
|
number = as_number(value)
|
|
if number is None:
|
|
continue
|
|
cap = (limits or {}).get(key, HISTORY_LIMIT)
|
|
series = self._history.get(key)
|
|
if series is None or series.maxlen != cap:
|
|
# A widget asking for a deeper series re-caps it in place.
|
|
series = deque(series or (), maxlen=cap)
|
|
self._history[key] = series
|
|
series.append((ts, number))
|
|
|
|
def record(
|
|
self,
|
|
values: dict[str, Any],
|
|
stamps: dict[str, float],
|
|
counters: list[str],
|
|
ts: float,
|
|
limits: dict[str, int] | None = None,
|
|
drop: list[str] | None = None,
|
|
) -> None:
|
|
"""The four writes under one acquisition of the one lock."""
|
|
with self._lock:
|
|
for key in drop or ():
|
|
self._data.pop(key, None)
|
|
self._history.pop(key, None)
|
|
self._data.update(values)
|
|
self._data.update(stamps)
|
|
for key, value in values.items():
|
|
number = as_number(value)
|
|
if number is None:
|
|
continue
|
|
cap = (limits or {}).get(key, HISTORY_LIMIT)
|
|
series = self._history.get(key)
|
|
if series is None or series.maxlen != cap:
|
|
series = deque(series or (), maxlen=cap)
|
|
self._history[key] = series
|
|
series.append((ts, number))
|
|
for key in counters:
|
|
self._data[key] = self._data.get(key, 0) + 1
|
|
|
|
def history(self, key: str) -> list[tuple[float, float]]:
|
|
"""The recorded points of one message, oldest first."""
|
|
with self._lock:
|
|
return list(self._history.get(key, ()))
|
|
|
|
def _moves(self, rename: Rename) -> tuple[dict[str, str], dict[str, str]]:
|
|
"""Where the keys of values and of histories go. Histories are kept
|
|
by bare message name here."""
|
|
data = {k: t for k in self._data if (t := rename.state_key(k)) is not None}
|
|
history = {k: rename.name(k) for k in self._history if rename.name(k) != k}
|
|
return data, history
|
|
|
|
def rename_keys(
|
|
self, rename: Rename, progress: Callable[[str], None] | None = None
|
|
) -> int:
|
|
"""One pass under the lock. A history moves as the deque it is, so it
|
|
keeps its capacity."""
|
|
with self._lock:
|
|
data, history = self._moves(rename)
|
|
for store, moves in ((self._data, data), (self._history, history)):
|
|
clash = next((t for t in moves.values() if t in store), None)
|
|
if clash is not None:
|
|
raise RuntimeError(
|
|
f"Will not move a key onto '{clash}', which exists"
|
|
)
|
|
for key, target in data.items():
|
|
self._data[target] = self._data.pop(key)
|
|
for key, target in history.items():
|
|
self._history[target] = self._history.pop(key)
|
|
if progress is not None:
|
|
progress("keys")
|
|
return len(data) + len(history)
|
|
|
|
def count(self, rename: Rename) -> int:
|
|
with self._lock:
|
|
data, history = self._moves(rename)
|
|
return len(data) + len(history)
|
|
|
|
|
|
class RedisState(StateBackend):
|
|
"""
|
|
Redis-based state backend for distributed execution.
|
|
|
|
Values are stored as JSON — anything a node passes through must be
|
|
serializable. Uses a namespace prefix to isolate pipeline executions.
|
|
|
|
:param host: Redis host address.
|
|
:type host: str
|
|
:param port: Redis port number.
|
|
:type port: int
|
|
:param db: Redis database number.
|
|
:type db: int
|
|
:param namespace: Key prefix for isolation.
|
|
:type namespace: str
|
|
:param password: Redis password (optional).
|
|
:type password: str | None
|
|
:param ttl: Time-to-live for keys in seconds (optional).
|
|
:type ttl: int | None
|
|
|
|
:example:
|
|
>>> state = RedisState(host="localhost", namespace="my_pipeline")
|
|
>>> state["result"] = {"value": 42}
|
|
>>> state["result"]
|
|
{'value': 42}
|
|
"""
|
|
|
|
__slots__ = ("_client", "_namespace", "_ttl", "_lock_name")
|
|
|
|
def __init__(
|
|
self,
|
|
host: str = "localhost",
|
|
port: int = 6379,
|
|
db: int = 0,
|
|
namespace: str = "pipeline",
|
|
password: str | None = None,
|
|
ttl: int | None = None,
|
|
) -> None:
|
|
self._client = redis.Redis(
|
|
host=host,
|
|
port=port,
|
|
db=db,
|
|
password=password,
|
|
decode_responses=False,
|
|
# Without these a Redis that stops answering without closing the
|
|
# connection — a partition, a firewall dropping state — hangs the
|
|
# caller until the kernel gives up, minutes later. That takes the
|
|
# health endpoint whose job is to notice with it.
|
|
socket_timeout=SOCKET_TIMEOUT_S,
|
|
socket_connect_timeout=CONNECT_TIMEOUT_S,
|
|
health_check_interval=30,
|
|
)
|
|
self._namespace = namespace
|
|
self._ttl = ttl
|
|
self._lock_name = f"{namespace}:_lock"
|
|
|
|
def _key(self, key: str) -> str:
|
|
"""Generate namespaced key."""
|
|
return f"{self._namespace}:{key}"
|
|
|
|
def _serialize(self, value: Any) -> bytes:
|
|
"""Serialize value for storage."""
|
|
return orjson.dumps(value, option=_JSON_OPTS)
|
|
|
|
def _deserialize(self, data: bytes | None) -> Any:
|
|
"""Deserialize value from storage."""
|
|
return orjson.loads(data) if data else None
|
|
|
|
# redis-py types every command as a sync/async union; this is the
|
|
# synchronous client, so the results are narrowed where they are consumed.
|
|
def get(self, key: str, default: Any = None) -> Any:
|
|
data = cast(bytes | None, self._client.get(self._key(key)))
|
|
return self._deserialize(data) if data is not None else default
|
|
|
|
def _names_key(self) -> str:
|
|
return self._key(NAMES)
|
|
|
|
def message_names(self) -> list[str]:
|
|
"""The message names, from the set kept beside them.
|
|
|
|
`keys()` is a SCAN of the whole namespace, which holds a version, a
|
|
timestamp, a history list and a last-seen marker for every message —
|
|
so reading the message names cost several times as many round trips
|
|
as there are messages. Every websocket connect asks for this.
|
|
"""
|
|
members = cast(set[bytes], self._client.smembers(self._names_key()))
|
|
return [m.decode("utf-8") for m in members]
|
|
|
|
def _note(self, pipe: Any, keys: Iterable[str]) -> None:
|
|
"""Record the message names among `keys`, in the caller's pipeline."""
|
|
names = [k for k in keys if not k.startswith("__")]
|
|
if names:
|
|
pipe.sadd(self._names_key(), *names)
|
|
if self._ttl:
|
|
pipe.expire(self._names_key(), self._ttl)
|
|
|
|
def _incr(self, pipe: Any, key: str) -> None:
|
|
"""INCR a counter in the caller's pipeline, expiring with the values."""
|
|
pipe.incr(self._key(key))
|
|
if self._ttl:
|
|
pipe.expire(self._key(key), self._ttl)
|
|
|
|
def set(self, key: str, value: Any) -> None:
|
|
data = self._serialize(value)
|
|
pipe = self._client.pipeline()
|
|
if self._ttl:
|
|
pipe.setex(self._key(key), self._ttl, data)
|
|
else:
|
|
pipe.set(self._key(key), data)
|
|
self._note(pipe, [key])
|
|
pipe.execute()
|
|
|
|
def exists(self, key: str) -> bool:
|
|
return bool(self._client.exists(self._key(key)))
|
|
|
|
def delete(self, key: str) -> None:
|
|
pipe = self._client.pipeline()
|
|
pipe.delete(self._key(key), self._history_key(key))
|
|
pipe.srem(self._names_key(), key)
|
|
pipe.execute()
|
|
|
|
def clear(self) -> None:
|
|
"""Clear all keys in the namespace."""
|
|
pattern = f"{self._namespace}:*"
|
|
cursor = 0
|
|
while True:
|
|
cursor, keys = cast(
|
|
tuple[int, list[bytes]],
|
|
self._client.scan(cursor, match=pattern, count=100),
|
|
)
|
|
if keys:
|
|
self._client.delete(*keys)
|
|
if cursor == 0:
|
|
break
|
|
|
|
def keys(self) -> list[str]:
|
|
"""Get all keys in the namespace (without prefix)."""
|
|
pattern = f"{self._namespace}:*"
|
|
prefix_len = len(self._namespace) + 1
|
|
result = []
|
|
cursor = 0
|
|
while True:
|
|
cursor, keys = cast(
|
|
tuple[int, list[bytes]],
|
|
self._client.scan(cursor, match=pattern, count=100),
|
|
)
|
|
for k in keys:
|
|
key_str = k.decode("utf-8")
|
|
# Skip the lock key
|
|
if not key_str.endswith(":_lock"):
|
|
result.append(key_str[prefix_len:])
|
|
if cursor == 0:
|
|
break
|
|
return result
|
|
|
|
def update(self, mapping: dict[str, Any]) -> None:
|
|
"""Update multiple keys atomically using a pipeline."""
|
|
if not mapping:
|
|
return
|
|
pipe = self._client.pipeline()
|
|
for key, value in mapping.items():
|
|
data = self._serialize(value)
|
|
if self._ttl:
|
|
pipe.setex(self._key(key), self._ttl, data)
|
|
else:
|
|
pipe.set(self._key(key), data)
|
|
self._note(pipe, mapping)
|
|
pipe.execute()
|
|
|
|
@contextmanager
|
|
def lock(self, timeout: float = 10.0) -> Iterator[None]:
|
|
"""
|
|
Acquire a distributed lock using Redis.
|
|
|
|
:param timeout: Lock timeout in seconds.
|
|
:type timeout: float
|
|
:yields: Lock context.
|
|
:raises TimeoutError: If lock cannot be acquired.
|
|
"""
|
|
redis_lock = self._client.lock(
|
|
self._lock_name,
|
|
timeout=timeout,
|
|
blocking_timeout=timeout,
|
|
)
|
|
acquired = redis_lock.acquire()
|
|
if not acquired:
|
|
raise TimeoutError("Failed to acquire Redis lock")
|
|
try:
|
|
yield
|
|
finally:
|
|
try:
|
|
redis_lock.release()
|
|
except redis.exceptions.LockError:
|
|
pass # Lock may have expired
|
|
|
|
def ping(self) -> bool:
|
|
"""
|
|
Check if Redis connection is available.
|
|
|
|
:returns: True if connected.
|
|
:rtype: bool
|
|
"""
|
|
try:
|
|
return bool(self._client.ping())
|
|
except redis.ConnectionError:
|
|
return False
|
|
|
|
def increment(self, key: str) -> int:
|
|
"""Atomically increment a counter using Redis INCR."""
|
|
pipe = self._client.pipeline()
|
|
self._incr(pipe, key)
|
|
return cast(int, pipe.execute()[0])
|
|
|
|
def increment_multi(self, keys: list[str]) -> None:
|
|
"""Bump every counter in one round trip rather than one INCR each."""
|
|
if not keys:
|
|
return
|
|
pipe = self._client.pipeline()
|
|
for key in keys:
|
|
self._incr(pipe, key)
|
|
pipe.execute()
|
|
|
|
def get_multi(self, keys: list[str]) -> dict[str, Any]:
|
|
"""Get multiple values atomically using Redis MGET."""
|
|
if not keys:
|
|
return {}
|
|
|
|
full_keys = [self._key(k) for k in keys]
|
|
values = cast(list[bytes | None], self._client.mget(full_keys))
|
|
|
|
result = {}
|
|
for key, value in zip(keys, values, strict=True):
|
|
result[key] = self._deserialize(value) if value is not None else None
|
|
|
|
return result
|
|
|
|
def get_present(self, keys: list[str]) -> dict[str, Any]:
|
|
"""The keys that are there, from one MGET.
|
|
|
|
A missing key comes back as ``None`` from Redis and a stored ``null``
|
|
comes back as the four bytes; only the raw reply tells the two apart,
|
|
which is why this is not ``get_multi`` with the Nones filtered out.
|
|
"""
|
|
if not keys:
|
|
return {}
|
|
values = cast(
|
|
list[bytes | None], self._client.mget([self._key(k) for k in keys])
|
|
)
|
|
return {
|
|
key: self._deserialize(raw)
|
|
for key, raw in zip(keys, values, strict=True)
|
|
if raw is not None
|
|
}
|
|
|
|
def __getitem__(self, key: str) -> Any:
|
|
# One GET rather than the base class's EXISTS-then-GET: a missing key
|
|
# and a stored `null` differ in the raw reply, so nothing has to ask
|
|
# twice.
|
|
raw = cast(bytes | None, self._client.get(self._key(key)))
|
|
if raw is None:
|
|
raise KeyError(key)
|
|
return self._deserialize(raw)
|
|
|
|
def compare_and_swap_multi(
|
|
self,
|
|
expected: dict[str, Any],
|
|
updates: dict[str, Any],
|
|
) -> bool:
|
|
"""
|
|
Atomically check expected values and apply updates using Redis WATCH/MULTI.
|
|
|
|
Uses optimistic locking with WATCH to detect concurrent modifications.
|
|
If any watched key changes before EXEC, the transaction is aborted.
|
|
|
|
:param expected: Dict of keys to their expected values.
|
|
:type expected: dict[str, Any]
|
|
:param updates: Dict of keys to their new values.
|
|
:type updates: dict[str, Any]
|
|
:returns: True if the swap succeeded, False otherwise.
|
|
:rtype: bool
|
|
"""
|
|
if not expected:
|
|
# No expectations, just set the updates
|
|
self.update(updates)
|
|
return True
|
|
|
|
watch_keys = [self._key(k) for k in expected.keys()]
|
|
|
|
try:
|
|
# Start watching the keys for changes
|
|
self._client.watch(*watch_keys)
|
|
|
|
# Check current values
|
|
current = self.get_multi(list(expected.keys()))
|
|
|
|
for key, exp_value in expected.items():
|
|
if current.get(key) != exp_value:
|
|
self._client.unwatch()
|
|
return False
|
|
|
|
# All values match - execute the transaction
|
|
pipe = self._client.pipeline(True) # True = use MULTI/EXEC
|
|
|
|
for key, value in updates.items():
|
|
data = self._serialize(value)
|
|
if self._ttl:
|
|
pipe.setex(self._key(key), self._ttl, data)
|
|
else:
|
|
pipe.set(self._key(key), data)
|
|
self._note(pipe, updates)
|
|
|
|
pipe.execute()
|
|
return True
|
|
|
|
except redis.WatchError:
|
|
# Another client modified one of the watched keys
|
|
return False
|
|
|
|
def close(self) -> None:
|
|
"""Give the connection pool back."""
|
|
self._client.close()
|
|
|
|
def _history_key(self, key: str) -> str:
|
|
return self._key(f"{HISTORY}:{key}")
|
|
|
|
def append_history(
|
|
self, values: dict[str, Any], ts: float, limits: dict[str, int] | None = None
|
|
) -> None:
|
|
"""Push the numeric values onto their capped list, in one round-trip."""
|
|
pipe = self._client.pipeline()
|
|
queued = False
|
|
for key, value in values.items():
|
|
number = as_number(value)
|
|
if number is None:
|
|
continue
|
|
history_key = self._history_key(key)
|
|
cap = (limits or {}).get(key, HISTORY_LIMIT)
|
|
pipe.lpush(history_key, orjson.dumps([ts, number]))
|
|
pipe.ltrim(history_key, 0, cap - 1)
|
|
if self._ttl:
|
|
pipe.expire(history_key, self._ttl)
|
|
queued = True
|
|
if queued:
|
|
pipe.execute()
|
|
|
|
def record(
|
|
self,
|
|
values: dict[str, Any],
|
|
stamps: dict[str, float],
|
|
counters: list[str],
|
|
ts: float,
|
|
limits: dict[str, int] | None = None,
|
|
drop: list[str] | None = None,
|
|
) -> None:
|
|
"""Value, timestamp, series and version counter in one round trip.
|
|
|
|
These were four pipelines — four round trips — for one published
|
|
value, and a value crossing an edge pays them twice. Redis executes a
|
|
pipeline as one transaction, so batching them changes nothing about
|
|
what lands together and removes three of the four trips.
|
|
"""
|
|
pipe = self._client.pipeline()
|
|
for key in drop or ():
|
|
pipe.delete(self._key(key), self._history_key(key))
|
|
for key, value in {**values, **stamps}.items():
|
|
data = self._serialize(value)
|
|
if self._ttl:
|
|
pipe.setex(self._key(key), self._ttl, data)
|
|
else:
|
|
pipe.set(self._key(key), data)
|
|
for key, value in values.items():
|
|
number = as_number(value)
|
|
if number is None:
|
|
continue
|
|
history_key = self._history_key(key)
|
|
cap = (limits or {}).get(key, HISTORY_LIMIT)
|
|
pipe.lpush(history_key, orjson.dumps([ts, number]))
|
|
pipe.ltrim(history_key, 0, cap - 1)
|
|
if self._ttl:
|
|
pipe.expire(history_key, self._ttl)
|
|
for key in counters:
|
|
self._incr(pipe, key)
|
|
self._note(pipe, values)
|
|
pipe.execute()
|
|
|
|
def history(self, key: str) -> list[tuple[float, float]]:
|
|
"""The recorded points of one message, oldest first."""
|
|
entries = cast(list[bytes], self._client.lrange(self._history_key(key), 0, -1))
|
|
# LPUSH puts the newest first, a chart reads the other way round.
|
|
points: list[tuple[float, float]] = []
|
|
for entry in reversed(entries):
|
|
ts, value = orjson.loads(entry)
|
|
points.append((float(ts), float(value)))
|
|
return points
|
|
|
|
def rename_keys(
|
|
self, rename: Rename, progress: Callable[[str], None] | None = None
|
|
) -> int:
|
|
"""``RENAMENX`` in batches, then the names set in one transaction.
|
|
|
|
The keys are listed before any moves, since a SCAN running over keys
|
|
that are being renamed can return one twice.
|
|
"""
|
|
moves = [(k, t) for k in self.keys() if (t := rename.state_key(k)) is not None]
|
|
moved = 0
|
|
for start in range(0, len(moves), RENAME_BATCH):
|
|
batch = moves[start : start + RENAME_BATCH]
|
|
moved += move_keys(
|
|
self._client, [(self._key(k), self._key(t)) for k, t in batch]
|
|
)
|
|
if progress is not None:
|
|
progress("keys")
|
|
names = [name for name in self.message_names() if rename.name(name) != name]
|
|
if names:
|
|
pipe = self._client.pipeline()
|
|
pipe.srem(self._names_key(), *names)
|
|
pipe.sadd(self._names_key(), *map(rename.name, names))
|
|
pipe.execute()
|
|
return moved
|