From 536e29a265102475caf6d5d2582a1f3a4aab3b82 Mon Sep 17 00:00:00 2001 From: Mike Langmayr <1809691+mikelangmayr@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:54:58 -0700 Subject: [PATCH 1/3] Add the keygrabber collector: config, scheduling and the daemon --- README.md | 5 +- docs/source/api/index.md | 18 ++ docs/source/index.md | 8 +- docs/source/keygrabber.md | 144 +++++++++++++++ libby/keygrabber/__init__.py | 17 ++ libby/keygrabber/collection.py | 145 +++++++++++++++ libby/keygrabber/config.py | 270 ++++++++++++++++++++++++++++ libby/keygrabber/daemon.py | 264 +++++++++++++++++++++++++++ libby/keygrabber/main.py | 41 +++++ pyproject.toml | 1 + tests/test_keygrabber_collection.py | 199 ++++++++++++++++++++ tests/test_keygrabber_config.py | 223 +++++++++++++++++++++++ tests/test_keygrabber_daemon.py | 191 ++++++++++++++++++++ 13 files changed, 1522 insertions(+), 4 deletions(-) create mode 100644 docs/source/keygrabber.md create mode 100644 libby/keygrabber/collection.py create mode 100644 libby/keygrabber/config.py create mode 100644 libby/keygrabber/daemon.py create mode 100644 libby/keygrabber/main.py create mode 100644 tests/test_keygrabber_collection.py create mode 100644 tests/test_keygrabber_config.py create mode 100644 tests/test_keygrabber_daemon.py diff --git a/README.md b/README.md index a6c31ed..edb2e83 100644 --- a/README.md +++ b/README.md @@ -4,8 +4,9 @@ Libby: a tiny messaging library which uses Bamboo with pluggable transports (ZMQ ## Documentation -Full docs (installation, keywords, the `Client` library, the `libby` CLI, and -how to build a `LibbyDaemon` peer, plus the generated API reference) are +Full docs (installation, keywords, the `Client` library, the `libby` CLI, how +to build a `LibbyDaemon` peer, the keygrabber, plus the generated API +reference) are built with Sphinx + the Shibuya theme and published to GitHub Pages: **[caltechopticalobservatories.github.io/libby](https://caltechopticalobservatories.github.io/libby/)** diff --git a/docs/source/api/index.md b/docs/source/api/index.md index 6746733..d6f2389 100644 --- a/docs/source/api/index.md +++ b/docs/source/api/index.md @@ -73,6 +73,24 @@ the docs build does not install the optional `influxdb` extra. :show-inheritance: ``` +```{eval-rst} +.. automodule:: libby.keygrabber.config + :members: + :show-inheritance: +``` + +```{eval-rst} +.. automodule:: libby.keygrabber.collection + :members: + :show-inheritance: +``` + +```{eval-rst} +.. automodule:: libby.keygrabber.daemon + :members: + :show-inheritance: +``` + ## Responses and errors ```{eval-rst} diff --git a/docs/source/index.md b/docs/source/index.md index f01d7df..4bb1a2d 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -4,13 +4,15 @@ Libby is a tiny messaging library built on [Bamboo](https://github.com/CaltechOp with pluggable transports (ZMQ or RabbitMQ). It gives you: - **Keywords** — typed, named values (`show` / `modify`) served over RPC, with - a registry, auto-generated `keys.list` / `keys.describe` services, and CLI - coercion. + a registry, auto-generated `keys.list` / `keys.describe` / `keys.read` + services, and CLI coercion. - **`LibbyDaemon`** — a base class for peers: lifecycle, discovery, RPC handlers, and pub/sub, in a few overrides. - **`Client`** — a long-lived, in-process handle for reading and writing keywords from scripts, and for blocking on a keyword condition with `wait_for` (libby's `ktl.waitFor`). +- **Keygrabber**: a daemon that polls keywords from other peers and writes them + to a time-series database for dashboarding. - **`libby` CLI** — a command-line front end for keyword peers (`show` / `modify` / `list` / `describe` / `waitfor`). @@ -23,6 +25,7 @@ keywords client cli daemon +keygrabber api/index ``` @@ -32,4 +35,5 @@ api/index - Building a peer that serves keywords? Read {doc}`keywords` then {doc}`daemon`. - Writing a script or tool that talks to peers? Read {doc}`client`. - Poking at peers interactively? Read {doc}`cli`. +- Recording keywords for Grafana? Read {doc}`keygrabber`. - Looking for a specific class or function? See {doc}`api/index`. diff --git a/docs/source/keygrabber.md b/docs/source/keygrabber.md new file mode 100644 index 0000000..e75a85e --- /dev/null +++ b/docs/source/keygrabber.md @@ -0,0 +1,144 @@ +# Keygrabber + +The keygrabber is a libby daemon that polls keywords from other peers on a +configurable cadence and writes them to a time-series database, so Grafana can +dashboard an instrument without every daemon growing its own database code. + +One process serves the whole fleet. Cadence lives in this daemon's config +rather than in each hardware daemon, adding a metric never restarts a daemon +that owns moving hardware, and the database credential lives in one place. + +InfluxDB 2.x is the only backend today. It needs the optional extra: + +```bash +pip install libby[influxdb] +``` + +## Running it + +```bash +keygrabber -c /etc/hispec/keygrabber.yaml +``` + +It is an ordinary `LibbyDaemon`, so `SIGTERM` stops it cleanly and the +`shutdown` keyword will too once the control surface lands. On the way out it +gives its retry queue a bounded chance to drain, so a graceful stop does not +lose the last tick. + +## Config + +A keygrabber config is an ordinary daemon config (`peer_id`, `group_id`, +`transport`, and the transport's own settings) plus three sections of its own. + +```yaml +peer_id: keygrabber +group_id: hispec +transport: rabbitmq +rabbitmq_url: amqp://localhost + +sink: + type: influxdb + url: http://influx.hispec:8086 + org: hispec + bucket: telemetry + token_env: HISPEC_INFLUX_TOKEN + +workers: 4 + +defaults: + interval_s: 10.0 + timeout_s: 2.0 + refresh_s: 300.0 + +collections: + adc: + peer: hsfei.adc + interval_s: 5.0 + keywords: ["positionvalue%", "ismoving", "isconnected"] + pressure: + peer: hsfei.atcpress + interval_s: 60.0 + keywords: ["%"] + exclude: ["units_code"] +``` + +### sink + +`type` selects the backend; `influxdb` is the only one implemented. + +The token is **never** written in the config. `token_env` names an environment +variable to read it from, and a `token` key in the file is rejected outright. +`timeout_ms` is optional. + +### collections + +One entry per peer and cadence; several entries may target the same peer at +different cadences. A collection name must match `[a-z0-9_]+`, because it +becomes the prefix of that collection's control keywords. + +- `peer` is `.`, the address of one peer +- `keywords` is a list of names or `%` patterns to record +- `exclude` removes names the includes matched +- `interval_s`, `timeout_s` and `refresh_s` fall back to `defaults` + +`uptime` and `lasterror` are excluded by default: `uptime` changes every second +and says nothing a timestamp does not, and `lasterror` is null most of the +time. Naming either one in `keywords` explicitly opts it back in. + +Patterns are resolved against the live peer at startup and again every +`refresh_s`, so keywords added by a restarted daemon get picked up without +restarting the keygrabber. + +### Cadence and timeouts + +Config load rejects an `interval_s` at or below `1.5 x timeout_s`. bamboo waits +`timeout_s` for the acknowledgement and a further `timeout_s / 2` for the +reply, so one unanswered read can occupy a worker for one and a half timeouts, +and a tighter interval would be overrun by a single slow peer. + +A tick whose predecessor is still running is skipped rather than queued behind +it, so a wedged peer cannot accumulate overlapping reads. + +## How it reads + +A tick is one `keys.read` request per peer, not one per keyword. This matters +more than the smaller number suggests: a daemon dispatches requests inline on +its receive thread, so reading twenty keywords individually does not overlap +anything on that daemon. It serializes exactly as a batch would, while paying +twenty dispatch cycles instead of one. The saving is contention on a control +daemon's only dispatch thread. + +A peer running a libby without `keys.read` is read one keyword at a time +instead. That is detected from the `services` field of `keys.list`, not by +trying `keys.read` and seeing what happens: an unknown key is dropped without +an acknowledgement, so a probe cannot tell an old peer from a dead one. + +Every value in a tick carries one timestamp, taken by the keygrabber rather +than by each daemon, which keeps clock skew between daemon hosts out of the +data. + +## Storage + +Samples reach the backend through a `Sink`, and nothing in a `Sample` is +Influx-shaped, so a second backend is a new sink rather than a change to the +collector. + +The InfluxDB schema is one measurement per keyword name, tagged with `group`, +`peer` and `units`, with a single `value` field. A keyword name carries one +type across peers, so field types stay consistent, and a Grafana query is a +measurement plus a `peer` tag filter. + +Three details follow from what Influx can store: + +- Integers are written as floats, so one peer reporting `0` and another `0.5` + for the same keyword cannot collide as int against float and be rejected. +- A null value is skipped, since Influx has no null field. It is not an error. +- A keyword with no units gets a `units=none` tag, because Influx drops an + empty tag value and the keyword would otherwise split into two series. + +A failed write goes to a bounded retry queue with exponential backoff, which +drops its oldest batch when full, so a database that stays down cannot grow the +daemon's memory without limit. + +See {mod}`libby.keygrabber.sink` in the {doc}`API reference ` for +the sink contract. diff --git a/libby/keygrabber/__init__.py b/libby/keygrabber/__init__.py index 064b9c0..c65b2f6 100644 --- a/libby/keygrabber/__init__.py +++ b/libby/keygrabber/__init__.py @@ -5,6 +5,15 @@ (``pip install libby[influxdb]``), so a deployment that only uses another backend never has to install it. """ +from .collection import Collection, TickResult +from .config import ( + CollectionConfig, + KeygrabberConfig, + build_sink, + parse_config, + select_keywords, +) +from .daemon import KeygrabberDaemon from .sink import ( RetryingWriter, RetryPolicy, @@ -15,10 +24,18 @@ ) __all__ = [ + "Collection", + "CollectionConfig", + "KeygrabberConfig", + "KeygrabberDaemon", "RetryingWriter", "RetryPolicy", "Sample", "Sink", "SinkError", "SinkWriteError", + "TickResult", + "build_sink", + "parse_config", + "select_keywords", ] diff --git a/libby/keygrabber/collection.py b/libby/keygrabber/collection.py new file mode 100644 index 0000000..59f744d --- /dev/null +++ b/libby/keygrabber/collection.py @@ -0,0 +1,145 @@ +"""One configured peer, resolved against the live peer and read on a cadence.""" +from __future__ import annotations + +import time +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Callable, Dict, List, Optional, Tuple + +from ..client import Client +from ..errors import LibbyError, LibbyTimeout +from .config import CollectionConfig, select_keywords +from .sink import Sample + +BULK_READ_SERVICE = "keys.read" + + +@dataclass(frozen=True) +class TickResult: + """What one read of a collection produced.""" + + samples: Tuple[Sample, ...] + read_errors: int + + +class Collection: + """Tracks what one peer exposes and turns a read of it into samples. + + Resolution is refreshed periodically rather than once, so keywords added by + a restarted daemon are picked up without restarting the keygrabber. + """ + + def __init__( + self, + config: CollectionConfig, + *, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self.config = config + self._clock = clock + self._names: Tuple[str, ...] = () + self._bulk_read = False + self._resolved_at: Optional[float] = None + + @property + def name(self) -> str: + """Return the collection's configured name.""" + return self.config.name + + @property + def keyword_count(self) -> int: + """Return how many keywords the last resolve selected.""" + return len(self._names) + + @property + def bulk_read(self) -> bool: + """Return whether the peer advertised the bulk read service.""" + return self._bulk_read + + def needs_resolve(self) -> bool: + """Return whether the keyword selection is due to be refreshed.""" + if self._resolved_at is None: + return True + return self._clock() - self._resolved_at >= self.config.refresh_s + + def resolve(self, client: Client) -> Tuple[str, ...]: + """Ask the peer what it serves and select the configured keywords. + + One ``keys.list`` covers both: the peer's keyword names, and whether it + serves ``keys.read``. Selection is then local, so a collection with + several patterns still costs one request. + """ + listing = client.listing(f"{self.config.peer}.%", + timeout_s=self.config.timeout_s) + available = [name.rsplit(".", 1)[-1] for name in listing.names] + self._names = select_keywords(self.config, available) + self._bulk_read = BULK_READ_SERVICE in listing.services + self._resolved_at = self._clock() + return self._names + + def tick(self, client: Client, timestamp: datetime) -> TickResult: + """Read every selected keyword once and return the samples.""" + if not self._names: + return TickResult((), 0) + + qualified = [f"{self.config.peer}.{name}" for name in self._names] + responses = ( + client.read(qualified, timeout_s=self.config.timeout_s) + if self._bulk_read + else self._read_individually(client, qualified) + ) + + samples: List[Sample] = [] + read_errors = 0 + for qualified_name, response in responses.items(): + if not response.get("ok"): + read_errors += 1 + continue + samples.append(self._sample(qualified_name, response, timestamp)) + return TickResult(tuple(samples), read_errors) + + def _read_individually( + self, + client: Client, + qualified: List[str], + ) -> Dict[str, Dict[str, Any]]: + """Read one keyword at a time, for a peer without ``keys.read``. + + Abandons the rest of the tick after the first timeout: a peer that has + stopped answering would otherwise cost ``timeout_s`` per keyword and + overrun the interval many times over. + """ + responses: Dict[str, Dict[str, Any]] = {} + timed_out = False + for name in qualified: + if timed_out: + responses[name] = {"ok": False, "error": "skipped after timeout"} + continue + try: + responses[name] = client.show(name, timeout_s=self.config.timeout_s) + except LibbyTimeout as exc: + timed_out = True + responses[name] = {"ok": False, "error": str(exc)} + except LibbyError as exc: + responses[name] = {"ok": False, "error": str(exc)} + return responses + + def _sample( + self, + qualified_name: str, + response: Dict[str, Any], + timestamp: datetime, + ) -> Sample: + """Build a sample from one keyword response. + + A null value is kept rather than dropped here: whether it can be stored + is the sink's business, and a backend other than Influx may hold it. + """ + return Sample( + keyword=qualified_name.rsplit(".", 1)[-1], + group=self.config.group, + peer=self.config.daemon, + value=response.get("value"), + units=response.get("units"), + timestamp=timestamp, + ) diff --git a/libby/keygrabber/config.py b/libby/keygrabber/config.py new file mode 100644 index 0000000..6d93cc1 --- /dev/null +++ b/libby/keygrabber/config.py @@ -0,0 +1,270 @@ +"""Parse and validate the keygrabber's own configuration sections. + +A keygrabber config file is an ordinary ``LibbyDaemon`` config (``peer_id``, +``group_id``, ``transport``, ...) plus three sections this module owns: +``sink``, ``defaults`` and ``collections``. +""" +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from typing import Any, List, Mapping, Sequence, Tuple + +from ..config import ConfigError +from ..keyword import match_pattern +from .sink import DEFAULT_MAX_BATCHES, RetryPolicy, Sink + +# Keywords no collection should record unless it names one explicitly. +# ``uptime`` changes every second and carries no information a timestamp does +# not already give; ``lasterror`` is null most of the time and Influx cannot +# store a null field. +DEFAULT_EXCLUDE: Tuple[str, ...] = ("lasterror", "uptime") + +DEFAULT_INTERVAL_S = 10.0 +DEFAULT_TIMEOUT_S = 2.0 +DEFAULT_REFRESH_S = 300.0 +DEFAULT_WORKERS = 4 + +# bamboo waits ``timeout_s`` for the ACK and a further ``timeout_s / 2`` for +# the response, so one wedged peer can occupy a worker for 1.5x the timeout. +TIMEOUT_HEADROOM = 1.5 + +_COLLECTION_NAME = re.compile(r"^[a-z0-9_]+$") + + +# A config record, so the field count is the schema rather than complexity +@dataclass(frozen=True) +class CollectionConfig: # pylint: disable=too-many-instance-attributes + """One configured peer, keyword selection and cadence.""" + + name: str + group: str + daemon: str + keywords: Tuple[str, ...] + exclude: Tuple[str, ...] + interval_s: float + timeout_s: float + refresh_s: float + + @property + def peer(self) -> str: + """Return the qualified peer address this collection reads.""" + return f"{self.group}.{self.daemon}" + + def excluded(self) -> Tuple[str, ...]: + """Return every exclusion, including the defaults not named explicitly.""" + defaults = tuple(name for name in DEFAULT_EXCLUDE + if name not in self.keywords) + return self.exclude + defaults + + +@dataclass(frozen=True) +class KeygrabberConfig: + """The keygrabber's parsed configuration.""" + + collections: Tuple[CollectionConfig, ...] + sink: Mapping[str, Any] + workers: int + retry: RetryPolicy + + +def parse_config(config: Mapping[str, Any]) -> KeygrabberConfig: + """Validate a daemon config's keygrabber sections into typed objects.""" + defaults = _mapping(config, "defaults") + collections = _mapping(config, "collections") + + parsed = tuple( + _parse_collection(name, collections[name], defaults) + for name in sorted(collections) + ) + return KeygrabberConfig( + collections=parsed, + sink=_mapping(config, "sink"), + workers=_positive_int(config.get("workers", DEFAULT_WORKERS), "workers"), + retry=_parse_retry(_mapping(config, "retry")), + ) + + +def build_sink(sink_config: Mapping[str, Any]) -> Sink: + """Build the configured sink, importing its backend client on demand.""" + kind = str(sink_config.get("type", "")).strip().lower() + if not kind: + raise ConfigError("sink.type is required") + if kind != "influxdb": + raise ConfigError(f"unsupported sink type {kind!r}; expected 'influxdb'") + + if "token" in sink_config: + raise ConfigError( + "sink.token must not appear in the config; name an environment " + "variable with sink.token_env instead" + ) + token_env = _required(sink_config, "token_env", "sink") + token = os.environ.get(str(token_env)) + if not token: + raise ConfigError( + f"environment variable {token_env} is unset or empty, so the " + "sink has no token" + ) + + # Validate every field before importing the backend, so a malformed config + # is reported the same way whether or not the extra is installed + settings = { + "url": str(_required(sink_config, "url", "sink")), + "org": str(_required(sink_config, "org", "sink")), + "bucket": str(_required(sink_config, "bucket", "sink")), + "token": token, + } + if "timeout_ms" in sink_config: + settings["timeout_ms"] = _positive_int( + sink_config["timeout_ms"], "sink.timeout_ms") + + # Deferred: influxdb-client is an optional extra, so a deployment using a + # different backend never has to install it + from .influx import ( # pylint: disable=import-outside-toplevel + InfluxConfig, InfluxSink, + ) + return InfluxSink(InfluxConfig(**settings)) + + +def _parse_collection( + name: str, + raw: Any, + defaults: Mapping[str, Any], +) -> CollectionConfig: + """Validate one entry of the collections section.""" + if not _COLLECTION_NAME.match(name): + raise ConfigError( + f"collection name {name!r} must match [a-z0-9_]+, so it composes " + "into control keyword names" + ) + if not isinstance(raw, Mapping): + raise ConfigError(f"collection {name!r} must be a mapping") + + group, daemon = _parse_peer(name, _required(raw, "peer", f"collection {name!r}")) + interval_s = _positive_float( + raw.get("interval_s", defaults.get("interval_s", DEFAULT_INTERVAL_S)), + f"collection {name!r} interval_s") + timeout_s = _positive_float( + raw.get("timeout_s", defaults.get("timeout_s", DEFAULT_TIMEOUT_S)), + f"collection {name!r} timeout_s") + + if interval_s <= TIMEOUT_HEADROOM * timeout_s: + raise ConfigError( + f"collection {name!r} has interval_s {interval_s} at or below " + f"{TIMEOUT_HEADROOM} x timeout_s ({TIMEOUT_HEADROOM * timeout_s}); " + "one unanswered read would overrun the interval" + ) + + return CollectionConfig( + name=name, + group=group, + daemon=daemon, + keywords=_patterns(raw.get("keywords"), f"collection {name!r} keywords", + required=True), + exclude=_patterns(raw.get("exclude"), f"collection {name!r} exclude", + required=False), + interval_s=interval_s, + timeout_s=timeout_s, + refresh_s=_positive_float( + raw.get("refresh_s", defaults.get("refresh_s", DEFAULT_REFRESH_S)), + f"collection {name!r} refresh_s"), + ) + + +def _parse_retry(raw: Mapping[str, Any]) -> RetryPolicy: + """Build the retry policy, letting RetryPolicy enforce its own bounds.""" + try: + return RetryPolicy( + max_batches=_positive_int( + raw.get("max_batches", DEFAULT_MAX_BATCHES), "retry.max_batches"), + **{ + key: _positive_float(raw[key], f"retry.{key}") + for key in ("base_backoff_s", "max_backoff_s") if key in raw + }, + ) + except ValueError as exc: + raise ConfigError(f"retry: {exc}") from exc + + +def _parse_peer(name: str, raw: Any) -> Tuple[str, str]: + """Split a collection's peer into group and daemon.""" + parts = str(raw).split(".") + if len(parts) != 2 or not all(parts) or "%" in str(raw): + raise ConfigError( + f"collection {name!r} peer {raw!r} must be '.'" + ) + return parts[0], parts[1] + + +def _patterns(raw: Any, label: str, *, required: bool) -> Tuple[str, ...]: + """Validate a list of keyword patterns.""" + if raw is None: + if required: + raise ConfigError(f"{label} is required") + return () + if isinstance(raw, str) or not isinstance(raw, Sequence): + raise ConfigError(f"{label} must be a list of strings") + if not all(isinstance(item, str) and item for item in raw): + raise ConfigError(f"{label} must be a list of non-empty strings") + if required and not raw: + raise ConfigError(f"{label} must name at least one pattern") + return tuple(raw) + + +def _mapping(config: Mapping[str, Any], key: str) -> Mapping[str, Any]: + """Return a config section, defaulting to empty.""" + section = config.get(key) + if section is None: + return {} + if not isinstance(section, Mapping): + raise ConfigError(f"{key} must be a mapping") + return section + + +def _required(section: Mapping[str, Any], key: str, label: str) -> Any: + """Return a required value, or raise naming what is missing.""" + value = section.get(key) + if value is None or (isinstance(value, str) and not value.strip()): + raise ConfigError(f"{label} is missing {key}") + return value + + +def _positive_float(raw: Any, label: str) -> float: + """Coerce a config value to a positive float.""" + try: + value = float(raw) + except (TypeError, ValueError) as exc: + raise ConfigError(f"{label} must be a number, got {raw!r}") from exc + if value <= 0: + raise ConfigError(f"{label} must be positive, got {value}") + return value + + +def _positive_int(raw: Any, label: str) -> int: + """Coerce a config value to a positive int.""" + if isinstance(raw, bool) or not isinstance(raw, int): + raise ConfigError(f"{label} must be an integer, got {raw!r}") + if raw < 1: + raise ConfigError(f"{label} must be at least 1, got {raw}") + return raw + + +def select_keywords( + collection: CollectionConfig, + available: Sequence[str], +) -> Tuple[str, ...]: + """Return the bare keyword names this collection should read. + + ``available`` is the peer's full keyword list. Includes are applied first, + then exclusions, so a default exclusion is dropped even when a wildcard + would otherwise have matched it. + """ + included: List[str] = [] + for pattern in collection.keywords: + included.extend(name for name in match_pattern(pattern, available) + if name not in included) + + excluded = {name for pattern in collection.excluded() + for name in match_pattern(pattern, included)} + return tuple(name for name in included if name not in excluded) diff --git a/libby/keygrabber/daemon.py b/libby/keygrabber/daemon.py new file mode 100644 index 0000000..ea5a510 --- /dev/null +++ b/libby/keygrabber/daemon.py @@ -0,0 +1,264 @@ +"""The keygrabber daemon: schedule reads, collect samples, write them out.""" +from __future__ import annotations + +import heapq +import queue +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from typing import Callable, List, Optional, Sequence, Tuple + +from ..client import Client +from ..daemon import LibbyDaemon +from ..errors import LibbyError +from ..libby import Libby +from .collection import Collection +from .config import KeygrabberConfig, build_sink, parse_config +from .sink import RetryingWriter, Sample, Sink + +# How long the writer waits for a batch before checking the retry queue +WRITER_POLL_S = 0.25 + +# Scheduler wake interval, so a stop request is noticed promptly even when the +# next tick is far off +SCHEDULER_TICK_S = 0.25 + +# Batches held between the readers and the writer. Bounded so a stalled writer +# cannot grow memory; the retry queue inside RetryingWriter handles a stalled +# backend. +QUEUE_DEPTH_PER_WORKER = 4 + +# Time allowed for the retry queue to drain during shutdown. Must stay well +# under the systemd unit's TimeoutStopSec so a wedged sink cannot turn a stop +# into a SIGKILL. +DRAIN_DEADLINE_S = 5.0 + + +class Counters: # pylint: disable=too-few-public-methods + """Tallies the daemon reports, guarded for cross-thread increments. + + The lock is only ever held around an increment, never across a read or a + sink write, because these are served on the transport's receive thread and + blocking it would time out every read already in flight. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self.points_written = 0 + self.read_errors = 0 + self.write_errors = 0 + self.skipped_ticks = 0 + self.dropped_batches = 0 + + def add(self, **deltas: int) -> None: + """Add to one or more counters.""" + with self._lock: + for name, delta in deltas.items(): + setattr(self, name, getattr(self, name) + delta) + + +# Coordinates config, client, sink, collections and three kinds of thread; +# the attribute count is collaborators, not state +class KeygrabberDaemon(LibbyDaemon): # pylint: disable=too-many-instance-attributes + """Polls keywords from other libby peers and writes them to a sink. + + One process serves the whole fleet: cadence lives in this daemon's config + rather than in each hardware daemon, and the database credential lives in + one place. + """ + + transport = "rabbitmq" + discovery_enabled = False + + def __init__(self) -> None: + super().__init__() + self.counters = Counters() + self._settings: Optional[KeygrabberConfig] = None + self._client: Optional[Client] = None + self._writer: Optional[RetryingWriter] = None + self._collections: List[Collection] = [] + self._queue: "queue.Queue[Tuple[Sample, ...]]" = queue.Queue() + self._pool: Optional[ThreadPoolExecutor] = None + self._threads: List[threading.Thread] = [] + self._halt = threading.Event() + self._in_flight: set[str] = set() + self._in_flight_lock = threading.Lock() + self._clock: Callable[[], float] = time.monotonic + + def on_start(self, libby: Libby) -> None: + """Parse config, open the sink, and start the reader and writer threads.""" + self._settings = parse_config(self._config) + self._client = Client(libby) + self._writer = RetryingWriter(self.make_sink(), + policy=self._settings.retry) + self._writer.connect() + self._collections = [Collection(config) + for config in self._settings.collections] + self._queue = queue.Queue( + maxsize=max(1, self._settings.workers * QUEUE_DEPTH_PER_WORKER)) + self._pool = ThreadPoolExecutor(max_workers=self._settings.workers, + thread_name_prefix="keygrabber-read") + self._halt.clear() + self._spawn(self._write_loop, "keygrabber-write") + self._spawn(self._schedule_loop, "keygrabber-schedule") + self.logger.info("collecting %d collections with %d workers", + len(self._collections), self._settings.workers) + + def on_stop(self, libby: Optional[Libby] = None) -> None: + """Stop the threads, then give the retry queue a bounded chance to drain.""" + self._halt.set() + if self._pool is not None: + self._pool.shutdown(wait=True) + self._pool = None + for thread in self._threads: + thread.join(timeout=WRITER_POLL_S * 8) + self._threads = [] + if self._writer is not None: + self._drain() + self._writer.close() + self._writer = None + + def make_sink(self) -> Sink: + """Build the configured sink. + + Override to supply a sink the config cannot describe, or to inject one + in a test without installing a database client. + """ + if self._settings is None: + raise LibbyError("keygrabber config has not been parsed yet") + return build_sink(self._settings.sink) + + def _spawn(self, target: Callable[[], None], name: str) -> None: + thread = threading.Thread(target=target, name=name, daemon=True) + thread.start() + self._threads.append(thread) + + ### Scheduling + + def _schedule_loop(self) -> None: + """Submit each collection's tick when it comes due.""" + now = self._clock() + pending: List[Tuple[float, str]] = [ + (now, collection.name) for collection in self._collections + ] + heapq.heapify(pending) + by_name = {collection.name: collection for collection in self._collections} + + while not self._halt.is_set(): + if not pending: + self._halt.wait(SCHEDULER_TICK_S) + continue + + due_at, name = pending[0] + delay = due_at - self._clock() + if delay > 0: + self._halt.wait(min(delay, SCHEDULER_TICK_S)) + continue + + heapq.heappop(pending) + collection = by_name[name] + self._submit(collection) + # Never schedule into the past: a long stall would otherwise queue + # a burst of catch-up ticks that can only skip + interval = collection.config.interval_s + heapq.heappush(pending, (max(due_at + interval, self._clock()), name)) + + def _submit(self, collection: Collection) -> None: + """Run a tick unless the previous one is still going.""" + with self._in_flight_lock: + if collection.name in self._in_flight: + self.counters.add(skipped_ticks=1) + self.logger.warning( + "collection %s skipped: previous read still in flight", + collection.name) + return + self._in_flight.add(collection.name) + + pool = self._pool + if pool is None: + return + pool.submit(self._run_tick, collection) + + def _run_tick(self, collection: Collection) -> None: + """Resolve if due, read once, and hand the samples to the writer.""" + try: + if collection.needs_resolve(): + collection.resolve(self._require_client()) + result = collection.tick(self._require_client(), + datetime.now(timezone.utc)) + if result.read_errors: + self.counters.add(read_errors=result.read_errors) + if result.samples: + self._enqueue(result.samples) + except LibbyError as exc: + self.counters.add(read_errors=1) + self.logger.error("collection %s read failed: %s", collection.name, exc) + finally: + with self._in_flight_lock: + self._in_flight.discard(collection.name) + + def _enqueue(self, samples: Sequence[Sample]) -> None: + try: + self._queue.put_nowait(tuple(samples)) + except queue.Full: + self.counters.add(dropped_batches=1) + self.logger.error("sample queue full; dropped %d samples", + len(samples)) + + ### Writing + + def _write_loop(self) -> None: + """Own every sink call, so no reader thread ever touches the backend.""" + while not self._halt.is_set(): + try: + batch = self._queue.get(timeout=WRITER_POLL_S) + except queue.Empty: + self._flush() + continue + self._write(batch) + self._flush() + + def _write(self, batch: Tuple[Sample, ...]) -> None: + writer = self._writer + if writer is None: + return + try: + self.counters.add(points_written=writer.write(batch)) + except LibbyError as exc: + self.counters.add(write_errors=1) + self.logger.error("sink write failed: %s", exc) + + def _flush(self) -> None: + writer = self._writer + if writer is None: + return + try: + self.counters.add(points_written=writer.flush_due()) + except LibbyError as exc: + self.counters.add(write_errors=1) + self.logger.error("sink retry failed: %s", exc) + + def _drain(self) -> None: + """Write whatever is still queued, under a deadline.""" + deadline = self._clock() + DRAIN_DEADLINE_S + while self._clock() < deadline: + try: + self._write(self._queue.get_nowait()) + continue + except queue.Empty: + pass + writer = self._writer + if writer is None or writer.queue_depth == 0: + return + self._flush() + if writer.queue_depth: + time.sleep(WRITER_POLL_S) + if self._writer is not None and self._writer.queue_depth: + self.logger.error("gave up draining %d batches after %.1fs", + self._writer.queue_depth, DRAIN_DEADLINE_S) + + def _require_client(self) -> Client: + if self._client is None: + raise LibbyError("keygrabber has no client; it is not started") + return self._client diff --git a/libby/keygrabber/main.py b/libby/keygrabber/main.py new file mode 100644 index 0000000..200e70e --- /dev/null +++ b/libby/keygrabber/main.py @@ -0,0 +1,41 @@ +"""Console entry point for the keygrabber daemon. + +Exits live here rather than in the daemon, so the daemon stays importable and +raises catchable exceptions. +""" +from __future__ import annotations + +import argparse +import sys +from typing import List, Optional + +from ..errors import LibbyError +from .daemon import KeygrabberDaemon + + +def main(argv: Optional[List[str]] = None) -> int: + """Run the keygrabber until it is stopped. Return an exit code.""" + parser = argparse.ArgumentParser( + prog="keygrabber", + description="Poll libby keywords and write them to a time-series database", + ) + parser.add_argument("-c", "--config", required=True, + help="path to the keygrabber config (YAML or JSON)") + parser.add_argument("-d", "--daemon-id", default=None, + help="daemon id, for a subsystem config with several") + namespace = parser.parse_args(argv) + + try: + daemon = KeygrabberDaemon.from_config_file(namespace.config, + namespace.daemon_id) + daemon.serve() + except KeyboardInterrupt: + return 0 + except (LibbyError, OSError, ValueError) as exc: + print(f"keygrabber: {exc}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index 616e507..87ebed0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,3 +47,4 @@ exclude = ["peers*", "package*"] [project.scripts] libby = "libby.cli.libby_cli:main" +keygrabber = "libby.keygrabber.main:main" diff --git a/tests/test_keygrabber_collection.py b/tests/test_keygrabber_collection.py new file mode 100644 index 0000000..60363ea --- /dev/null +++ b/tests/test_keygrabber_collection.py @@ -0,0 +1,199 @@ +"""Unit tests for Collection: resolution, refresh, and one tick's samples.""" +from __future__ import annotations + +import unittest +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional, Sequence, Tuple + +from libby import KeyListing, LibbyTimeout +from libby.keygrabber import Collection, parse_config + +TIMESTAMP = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc) + + +def _collection_config(**entry: Any): + config = { + "collections": { + "adc": {"peer": "hsfei.adc", "keywords": ["%"], **entry}, + }, + } + return parse_config(config).collections[0] + + +class _FakeClock: + """Manually advanced clock, so refresh windows need no real time.""" + + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + """Move the clock forward.""" + self.now += seconds + + +class _FakeClient: + """Stands in for Client, serving a canned listing and canned reads.""" + + # Timeout and chunking arguments are accepted to match Client's signatures + # pylint: disable=unused-argument + + def __init__( + self, + names: Sequence[str] = ("positionvalue", "isconnected"), + services: Sequence[str] = ("keys.read",), + values: Optional[Dict[str, Dict[str, Any]]] = None, + timeouts_after: Optional[int] = None, + ) -> None: + self.names = list(names) + self.services = list(services) + self.values = values or {} + self.timeouts_after = timeouts_after + self.read_calls: List[Tuple[str, ...]] = [] + self.show_calls: List[str] = [] + self.listings = 0 + + def listing(self, pattern: str, *, timeout_s: float = 3.0) -> KeyListing: + """Return the canned listing, qualified like the real client does.""" + self.listings += 1 + prefix = pattern.rsplit(".", 1)[0] + return KeyListing( + names=tuple(f"{prefix}.{name}" for name in self.names), + services=tuple(self.services), + ) + + def read(self, names: Sequence[str], *, timeout_s: float = 3.0, + chunk_size: int = 100) -> Dict[str, Dict[str, Any]]: + """Record the batch and answer from the canned values.""" + self.read_calls.append(tuple(names)) + return {name: self._value(name) for name in names} + + def show(self, name: str, *, timeout_s: float = 3.0) -> Dict[str, Any]: + """Answer one keyword, timing out once past the configured point.""" + self.show_calls.append(name) + if self.timeouts_after is not None and len(self.show_calls) > self.timeouts_after: + raise LibbyTimeout(f"{name}: request timed out") + return self._value(name) + + def _value(self, qualified: str) -> Dict[str, Any]: + keyword = qualified.rsplit(".", 1)[-1] + return self.values.get(keyword, {"ok": True, "value": 1.0, "units": "mm"}) + + +class ResolveTests(unittest.TestCase): + """Keyword selection and capability detection at resolve time.""" + + def test_resolve_selects_and_detects_bulk_read(self): + """Take names and the bulk-read capability from one keys.list.""" + collection = Collection(_collection_config()) + names = collection.resolve(_FakeClient()) + self.assertEqual(names, ("isconnected", "positionvalue")) + self.assertTrue(collection.bulk_read) + + def test_peer_without_the_service_falls_back(self): + """Read one keyword at a time when the peer omits keys.read.""" + client = _FakeClient(services=()) + collection = Collection(_collection_config()) + collection.resolve(client) + self.assertFalse(collection.bulk_read) + + collection.tick(client, TIMESTAMP) + self.assertEqual(client.read_calls, []) + self.assertEqual(len(client.show_calls), 2) + + def test_resolve_is_due_before_the_first_read(self): + """Resolve once before reading anything.""" + self.assertTrue(Collection(_collection_config()).needs_resolve()) + + def test_resolve_is_not_due_again_until_refresh(self): + """Hold the selection until the refresh window elapses.""" + clock = _FakeClock() + collection = Collection(_collection_config(refresh_s=300.0, + interval_s=10.0), + clock=clock) + collection.resolve(_FakeClient()) + self.assertFalse(collection.needs_resolve()) + + clock.advance(299.0) + self.assertFalse(collection.needs_resolve()) + clock.advance(2.0) + self.assertTrue(collection.needs_resolve()) + + +class TickTests(unittest.TestCase): + """What one read of a collection produces.""" + + def _ticked(self, client: _FakeClient): + collection = Collection(_collection_config()) + collection.resolve(client) + return collection.tick(client, TIMESTAMP) + + def test_one_request_covers_the_whole_tick(self): + """Read a peer with a single bulk request.""" + client = _FakeClient() + self._ticked(client) + self.assertEqual(len(client.read_calls), 1) + self.assertEqual(client.read_calls[0], + ("hsfei.adc.isconnected", "hsfei.adc.positionvalue")) + + def test_samples_carry_group_peer_units_and_timestamp(self): + """Tag each sample with its own peer and the tick's read time.""" + result = self._ticked(_FakeClient()) + sample = next(s for s in result.samples if s.keyword == "positionvalue") + self.assertEqual(sample.group, "hsfei") + self.assertEqual(sample.peer, "adc") + self.assertEqual(sample.units, "mm") + self.assertEqual(sample.timestamp, TIMESTAMP) + + def test_every_sample_in_a_tick_shares_one_timestamp(self): + """Stamp a tick once, so a dashboard can correlate its values.""" + result = self._ticked(_FakeClient()) + self.assertEqual({s.timestamp for s in result.samples}, {TIMESTAMP}) + + def test_failed_read_is_counted_not_raised(self): + """Report a broken getter as an error without losing the tick.""" + client = _FakeClient(values={ + "positionvalue": {"ok": False, "error": "hardware unreachable"}, + }) + result = self._ticked(client) + self.assertEqual(result.read_errors, 1) + self.assertEqual([s.keyword for s in result.samples], ["isconnected"]) + + def test_null_value_is_kept_as_a_sample(self): + """Leave a null for the sink to judge, so a backend may store it.""" + client = _FakeClient(values={ + "positionvalue": {"ok": True, "value": None}, + }) + result = self._ticked(client) + self.assertEqual(result.read_errors, 0) + sample = next(s for s in result.samples if s.keyword == "positionvalue") + self.assertIsNone(sample.value) + + def test_tick_without_a_resolve_reads_nothing(self): + """Do nothing until the selection is known.""" + client = _FakeClient() + result = Collection(_collection_config()).tick(client, TIMESTAMP) + self.assertEqual(result, type(result)((), 0)) + self.assertEqual(client.read_calls, []) + + def test_fallback_abandons_the_tick_after_a_timeout(self): + """Stop reading a peer that has stopped answering. + + Paying timeout_s per keyword would overrun the interval many times + over, so the rest of the tick is abandoned after the first timeout. + """ + client = _FakeClient( + names=("a", "b", "c", "d"), services=(), timeouts_after=1) + collection = Collection(_collection_config()) + collection.resolve(client) + result = collection.tick(client, TIMESTAMP) + + self.assertEqual(len(client.show_calls), 2) # one ok, one timeout + self.assertEqual(result.read_errors, 4 - 1) + self.assertEqual(len(result.samples), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_keygrabber_config.py b/tests/test_keygrabber_config.py new file mode 100644 index 0000000..005ee3e --- /dev/null +++ b/tests/test_keygrabber_config.py @@ -0,0 +1,223 @@ +"""Unit tests for keygrabber config parsing, validation and keyword selection.""" +from __future__ import annotations + +import os +import unittest +from typing import Any, Dict +from unittest import mock + +from libby.config import ConfigError +from libby.keygrabber import parse_config, select_keywords +from libby.keygrabber.config import DEFAULT_EXCLUDE, build_sink + +PEER_KEYWORDS = [ + "isconnected", "ismoving", "isreferenced", "lasterror", + "positionvalue1", "positionvalue2", "softmax", "uptime", +] + + +def _config(**overrides: Any) -> Dict[str, Any]: + config: Dict[str, Any] = { + "peer_id": "keygrabber", + "group_id": "hispec", + "defaults": {"interval_s": 10.0, "timeout_s": 2.0, "refresh_s": 300.0}, + "collections": { + "adc": {"peer": "hsfei.adc", "keywords": ["positionvalue%"]}, + }, + } + config.update(overrides) + return config + + +class ParseConfigTests(unittest.TestCase): + """Validation of the collections and defaults sections.""" + + def test_defaults_are_inherited(self): + """Fall back to the defaults section for unset cadences.""" + collection = parse_config(_config()).collections[0] + self.assertEqual(collection.interval_s, 10.0) + self.assertEqual(collection.timeout_s, 2.0) + self.assertEqual(collection.refresh_s, 300.0) + + def test_collection_overrides_defaults(self): + """Let one collection set its own cadence.""" + config = _config(collections={ + "adc": {"peer": "hsfei.adc", "keywords": ["%"], "interval_s": 60.0}, + }) + self.assertEqual(parse_config(config).collections[0].interval_s, 60.0) + + def test_peer_is_split_into_group_and_daemon(self): + """Keep group and daemon apart, since they become separate tags.""" + collection = parse_config(_config()).collections[0] + self.assertEqual((collection.group, collection.daemon), ("hsfei", "adc")) + self.assertEqual(collection.peer, "hsfei.adc") + + def test_collections_are_ordered_by_name(self): + """Parse deterministically, so two runs agree.""" + config = _config(collections={ + "zzz": {"peer": "hsfei.z", "keywords": ["%"]}, + "aaa": {"peer": "hsfei.a", "keywords": ["%"]}, + }) + names = [c.name for c in parse_config(config).collections] + self.assertEqual(names, ["aaa", "zzz"]) + + def test_interval_must_clear_the_timeout(self): + """Reject a cadence one unanswered read would overrun. + + bamboo waits timeout_s for the ACK and timeout_s/2 more for the reply, + so 2s of timeout can occupy a worker for 3s. + """ + config = _config(collections={ + "adc": {"peer": "hsfei.adc", "keywords": ["%"], + "interval_s": 3.0, "timeout_s": 2.0}, + }) + with self.assertRaises(ConfigError) as caught: + parse_config(config) + self.assertIn("interval_s", str(caught.exception)) + + def test_interval_just_above_the_headroom_is_accepted(self): + """Allow a cadence that clears 1.5x the timeout.""" + config = _config(collections={ + "adc": {"peer": "hsfei.adc", "keywords": ["%"], + "interval_s": 3.1, "timeout_s": 2.0}, + }) + self.assertEqual(parse_config(config).collections[0].interval_s, 3.1) + + def test_collection_name_must_compose_into_a_keyword(self): + """Reject a name that could not become a control keyword prefix.""" + config = _config(collections={ + "ADC Rotator": {"peer": "hsfei.adc", "keywords": ["%"]}, + }) + with self.assertRaises(ConfigError): + parse_config(config) + + def test_peer_must_be_group_and_daemon(self): + """Reject a peer that is not exactly '.'.""" + for peer in ("adc", "hsfei.adc.extra", "hsfei.", "hsfei.%"): + with self.subTest(peer=peer): + config = _config(collections={ + "adc": {"peer": peer, "keywords": ["%"]}, + }) + with self.assertRaises(ConfigError): + parse_config(config) + + def test_keywords_are_required(self): + """Refuse a collection that selects nothing.""" + for keywords in (None, [], "positionvalue", [""], [1]): + with self.subTest(keywords=keywords): + entry: Dict[str, Any] = {"peer": "hsfei.adc"} + if keywords is not None: + entry["keywords"] = keywords + with self.assertRaises(ConfigError): + parse_config(_config(collections={"adc": entry})) + + def test_non_numeric_cadence_is_rejected(self): + """Reject a cadence that is not a number.""" + config = _config(collections={ + "adc": {"peer": "hsfei.adc", "keywords": ["%"], "interval_s": "soon"}, + }) + with self.assertRaises(ConfigError): + parse_config(config) + + def test_workers_must_be_a_positive_integer(self): + """Reject a pool that could never run a tick.""" + for workers in (0, -1, 1.5, True, "four"): + with self.subTest(workers=workers): + with self.assertRaises(ConfigError): + parse_config(_config(workers=workers)) + + def test_retry_bounds_are_reported_as_config_errors(self): + """Surface a bad retry policy as a config error, not a ValueError.""" + with self.assertRaises(ConfigError): + parse_config(_config(retry={"max_batches": 0})) + + def test_missing_collections_section_is_allowed(self): + """Let a daemon start with nothing configured yet.""" + config = _config() + del config["collections"] + self.assertEqual(parse_config(config).collections, ()) + + +class SelectKeywordsTests(unittest.TestCase): + """Include and exclude patterns applied to a peer's keyword list.""" + + def _selected(self, keywords, exclude=None): + collection = parse_config(_config(collections={ + "adc": {"peer": "hsfei.adc", "keywords": keywords, + **({"exclude": exclude} if exclude else {})}, + })).collections[0] + return select_keywords(collection, PEER_KEYWORDS) + + def test_wildcard_selects_matching_names(self): + """Match a trailing wildcard within a single name.""" + self.assertEqual(self._selected(["positionvalue%"]), + ("positionvalue1", "positionvalue2")) + + def test_several_patterns_are_unioned_without_duplicates(self): + """Union patterns, keeping each name once.""" + selected = self._selected(["is%", "isconnected"]) + self.assertEqual(selected, ("isconnected", "ismoving", "isreferenced")) + + def test_explicit_exclude_is_applied_after_includes(self): + """Drop an excluded name a wildcard would otherwise have matched.""" + self.assertEqual(self._selected(["is%"], exclude=["ismoving"]), + ("isconnected", "isreferenced")) + + def test_noisy_keywords_are_excluded_by_default(self): + """Keep uptime and lasterror out of a select-everything collection. + + uptime changes every second and says nothing a timestamp does not, and + lasterror is null most of the time. + """ + selected = self._selected(["%"]) + for name in DEFAULT_EXCLUDE: + self.assertNotIn(name, selected) + self.assertIn("positionvalue1", selected) + + def test_naming_a_default_exclusion_opts_back_in(self): + """Honour an explicit request for a keyword the default would drop.""" + self.assertEqual(self._selected(["uptime"]), ("uptime",)) + + def test_no_matches_selects_nothing(self): + """Return empty rather than failing when a pattern matches nothing.""" + self.assertEqual(self._selected(["nosuch%"]), ()) + + +class BuildSinkTests(unittest.TestCase): + """Sink construction from the config's sink section.""" + + def test_token_in_the_config_is_refused(self): + """Refuse an inline credential so it cannot reach a config file.""" + with self.assertRaises(ConfigError) as caught: + build_sink({"type": "influxdb", "url": "http://x", "org": "o", + "bucket": "b", "token": "secret"}) + self.assertIn("token_env", str(caught.exception)) + + def test_unset_token_env_is_reported(self): + """Name the missing variable rather than failing at write time.""" + with mock.patch.dict(os.environ, {}, clear=True): + with self.assertRaises(ConfigError) as caught: + build_sink({"type": "influxdb", "url": "http://x", "org": "o", + "bucket": "b", "token_env": "NOT_SET_ANYWHERE"}) + self.assertIn("NOT_SET_ANYWHERE", str(caught.exception)) + + def test_missing_type_is_reported(self): + """Require an explicit sink type.""" + with self.assertRaises(ConfigError): + build_sink({}) + + def test_unknown_type_is_reported(self): + """Name the unsupported backend.""" + with self.assertRaises(ConfigError) as caught: + build_sink({"type": "sqlite"}) + self.assertIn("sqlite", str(caught.exception)) + + def test_missing_connection_fields_are_reported(self): + """Require url, org and bucket before reaching the client.""" + with mock.patch.dict(os.environ, {"TOKEN": "t"}, clear=True): + with self.assertRaises(ConfigError): + build_sink({"type": "influxdb", "token_env": "TOKEN"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_keygrabber_daemon.py b/tests/test_keygrabber_daemon.py new file mode 100644 index 0000000..8cfd085 --- /dev/null +++ b/tests/test_keygrabber_daemon.py @@ -0,0 +1,191 @@ +"""End-to-end test: the keygrabber against a live peer, over ZMQ. + +Uses ZMQ so no broker is needed, and a recording sink so no database is. +""" +from __future__ import annotations + +import socket +import threading +import time +import unittest +from typing import List, Sequence + +from libby.daemon import LibbyDaemon +from libby.keygrabber import KeygrabberDaemon, Sample + +SETTLE_TIMEOUT_S = 15.0 + + +def _free_endpoint() -> str: + """Return a loopback ZMQ endpoint on a port the OS just reported free.""" + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return f"tcp://127.0.0.1:{probe.getsockname()[1]}" + + +class _FixtureDaemon(LibbyDaemon): + """Peer the keygrabber reads, with one keyword of each interesting shape.""" + + peer_id = "kgtarget" + group_id = "hsfei" + transport = "zmq" + discovery_enabled = False + + def on_start(self, libby) -> None: + self.keyword_registry.float("positionvalue", getter=lambda: 7.5, + units="mm") + self.keyword_registry.bool("isconnected", getter=lambda: True) + self.keyword_registry.string("status", getter=lambda: "Ready") + + +class _RecordingSink: + """Collects everything the keygrabber writes, in place of a database.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self.samples: List[Sample] = [] + + def connect(self) -> None: + """Nothing to open.""" + + def is_connected(self) -> bool: + """Always reachable.""" + return True + + def write(self, samples: Sequence[Sample]) -> int: + """Record a batch and report it all stored.""" + with self._lock: + self.samples.extend(samples) + return len(samples) + + def close(self) -> None: + """Nothing to release.""" + + def keywords(self) -> List[str]: + """Return the keyword names recorded so far.""" + return [sample.keyword for sample in self.snapshot()] + + def snapshot(self) -> List[Sample]: + """Return a copy of the recorded samples.""" + with self._lock: + return list(self.samples) + + +class _TestKeygrabber(KeygrabberDaemon): + """Keygrabber writing to a recording sink instead of InfluxDB. + + ``from_config`` constructs a daemon with no arguments, so the sink is + assigned after construction and before ``start``. + """ + + sink: _RecordingSink + + def make_sink(self) -> _RecordingSink: + """Return the injected sink.""" + return self.sink + + +class KeygrabberEndToEndTests(unittest.TestCase): + """The daemon reads a real peer on a cadence and writes what it read.""" + + @classmethod + def setUpClass(cls): + cls.target_endpoint = _free_endpoint() + cls.target = _FixtureDaemon() + cls.target.bind = cls.target_endpoint + cls.target.start() + + @classmethod + def tearDownClass(cls): + cls.target.stop() + + def setUp(self) -> None: + self.sink = _RecordingSink() + self.grabber = _TestKeygrabber.from_config({ + "peer_id": "keygrabber", + "group_id": "hispec", + "transport": "zmq", + "bind": _free_endpoint(), + "address_book": {"hsfei.kgtarget": self.target_endpoint}, + "discovery_enabled": False, + "sink": {"type": "recording"}, + "workers": 2, + "collections": { + "target": { + "peer": "hsfei.kgtarget", + "keywords": ["%"], + "interval_s": 0.5, + "timeout_s": 0.3, + "refresh_s": 60.0, + }, + }, + }) + self.grabber.sink = self.sink + + def tearDown(self) -> None: + self.grabber.stop() + + def _await_samples(self, minimum: int = 3) -> None: + """Block until enough samples arrive, rather than sleeping a guess.""" + deadline = time.monotonic() + SETTLE_TIMEOUT_S + while time.monotonic() < deadline: + if len(self.sink.keywords()) >= minimum: + return + time.sleep(0.05) + self.fail(f"only {len(self.sink.keywords())} samples within " + f"{SETTLE_TIMEOUT_S}s: {self.sink.keywords()}") + + def test_collects_the_selected_keywords(self): + """Read every selected keyword and write it to the sink.""" + self.grabber.start() + self._await_samples() + self.assertEqual(set(self.sink.keywords()), + {"positionvalue", "isconnected", "status"}) + + def test_default_exclusions_are_not_collected(self): + """Keep uptime and lasterror out of a select-everything collection.""" + self.grabber.start() + self._await_samples() + recorded = set(self.sink.keywords()) + self.assertNotIn("uptime", recorded) + self.assertNotIn("lasterror", recorded) + + def test_samples_carry_their_tags_and_values(self): + """Carry group, peer, units and value through to the sink.""" + self.grabber.start() + self._await_samples() + sample = next(s for s in self.sink.snapshot() + if s.keyword == "positionvalue") + self.assertEqual((sample.group, sample.peer), ("hsfei", "kgtarget")) + self.assertEqual(sample.value, 7.5) + self.assertEqual(sample.units, "mm") + + def test_points_written_is_counted(self): + """Report what the sink actually stored.""" + self.grabber.start() + self._await_samples() + self.assertGreater(self.grabber.counters.points_written, 0) + self.assertEqual(self.grabber.counters.write_errors, 0) + + def test_repeats_on_the_configured_cadence(self): + """Read again on the next interval rather than once at startup.""" + self.grabber.start() + self._await_samples(minimum=3) + first = len(self.sink.keywords()) + deadline = time.monotonic() + SETTLE_TIMEOUT_S + while time.monotonic() < deadline: + if len(self.sink.keywords()) > first: + return + time.sleep(0.05) + self.fail("the collection never ran a second tick") + + def test_stop_is_clean_with_nothing_left_queued(self): + """Drain on the way out, so the last tick is not dropped.""" + self.grabber.start() + self._await_samples() + self.grabber.stop() + self.assertEqual(self.grabber.counters.dropped_batches, 0) + + +if __name__ == "__main__": + unittest.main() From 1d57a859e397d90c2e442f075415301132c2d884 Mon Sep 17 00:00:00 2001 From: Mike Langmayr <1809691+mikelangmayr@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:58:00 -0700 Subject: [PATCH 2/3] Run the keygrabber end-to-end tests on both transports --- tests/test_keygrabber_daemon.py | 266 +++++++++++++++++++++----------- 1 file changed, 177 insertions(+), 89 deletions(-) diff --git a/tests/test_keygrabber_daemon.py b/tests/test_keygrabber_daemon.py index 8cfd085..1d95537 100644 --- a/tests/test_keygrabber_daemon.py +++ b/tests/test_keygrabber_daemon.py @@ -1,6 +1,11 @@ -"""End-to-end test: the keygrabber against a live peer, over ZMQ. +"""End-to-end tests: the keygrabber against a live peer, on each transport. -Uses ZMQ so no broker is needed, and a recording sink so no database is. +Every case runs once per transport. The ZMQ cases need no external service and +always run; the RabbitMQ cases skip unless a broker is reachable. RabbitMQ +matters most here: it is the default transport, and the one where the worker +pool issues concurrent requests over a single pika connection. + +A recording sink stands in for the database, so no server is needed either. """ from __future__ import annotations @@ -12,8 +17,20 @@ from libby.daemon import LibbyDaemon from libby.keygrabber import KeygrabberDaemon, Sample +from libby.rabbitmq_transport import RabbitMQTransport SETTLE_TIMEOUT_S = 15.0 +RABBITMQ_URL = "amqp://localhost" + + +def _broker_available() -> bool: + try: + probe = RabbitMQTransport(peer_id="keygrabber-test-probe", + rabbitmq_url=RABBITMQ_URL) + probe.stop() + return True + except Exception: # pylint: disable=broad-exception-caught + return False def _free_endpoint() -> str: @@ -26,9 +43,7 @@ def _free_endpoint() -> str: class _FixtureDaemon(LibbyDaemon): """Peer the keygrabber reads, with one keyword of each interesting shape.""" - peer_id = "kgtarget" group_id = "hsfei" - transport = "zmq" discovery_enabled = False def on_start(self, libby) -> None: @@ -38,6 +53,21 @@ def on_start(self, libby) -> None: self.keyword_registry.string("status", getter=lambda: "Ready") +class _ZmqFixtureDaemon(_FixtureDaemon): + """Target peer served over ZMQ; ``bind`` is assigned per instance.""" + + peer_id = "kgtargetzmq" + transport = "zmq" + + +class _RabbitFixtureDaemon(_FixtureDaemon): + """Target peer served over RabbitMQ.""" + + peer_id = "kgtargetrmq" + transport = "rabbitmq" + rabbitmq_url = RABBITMQ_URL + + class _RecordingSink: """Collects everything the keygrabber writes, in place of a database.""" @@ -85,106 +115,164 @@ def make_sink(self) -> _RecordingSink: return self.sink -class KeygrabberEndToEndTests(unittest.TestCase): - """The daemon reads a real peer on a cadence and writes what it read.""" +# Nested inside a plain class so unittest's loader, which collects every +# module-level TestCase subclass, does not run the bases on their own. +class _Bases: # pylint: disable=too-few-public-methods + """Namespace for the transport-agnostic cases each transport subclasses.""" + + class KeygrabberCases(unittest.TestCase): + """The daemon reads a real peer on a cadence and writes what it read. + + Concrete subclasses start their own target peer and supply the matching + keygrabber config. + """ + + target_peer: str + + def grabber_config(self) -> dict: + """Return the transport-specific keygrabber config.""" + raise NotImplementedError + + def setUp(self) -> None: # pylint: disable=invalid-name + self.sink = _RecordingSink() + self.grabber = _TestKeygrabber.from_config(self.grabber_config()) + self.grabber.sink = self.sink + + def tearDown(self) -> None: # pylint: disable=invalid-name + self.grabber.stop() + + def _await_samples(self, minimum: int = 3) -> None: + """Block until enough samples arrive, rather than sleeping a guess.""" + deadline = time.monotonic() + SETTLE_TIMEOUT_S + while time.monotonic() < deadline: + if len(self.sink.keywords()) >= minimum: + return + time.sleep(0.05) + self.fail(f"only {len(self.sink.keywords())} samples within " + f"{SETTLE_TIMEOUT_S}s: {self.sink.keywords()}") + + def test_collects_the_selected_keywords(self): + """Read every selected keyword and write it to the sink.""" + self.grabber.start() + self._await_samples() + self.assertEqual(set(self.sink.keywords()), + {"positionvalue", "isconnected", "status"}) + + def test_default_exclusions_are_not_collected(self): + """Keep uptime and lasterror out of a select-everything collection.""" + self.grabber.start() + self._await_samples() + recorded = set(self.sink.keywords()) + self.assertNotIn("uptime", recorded) + self.assertNotIn("lasterror", recorded) + + def test_samples_carry_their_tags_and_values(self): + """Carry group, peer, units and value through to the sink.""" + self.grabber.start() + self._await_samples() + sample = next(s for s in self.sink.snapshot() + if s.keyword == "positionvalue") + group, daemon = self.target_peer.split(".") + self.assertEqual((sample.group, sample.peer), (group, daemon)) + self.assertEqual(sample.value, 7.5) + self.assertEqual(sample.units, "mm") + + def test_points_written_is_counted(self): + """Report what the sink actually stored.""" + self.grabber.start() + self._await_samples() + self.assertGreater(self.grabber.counters.points_written, 0) + self.assertEqual(self.grabber.counters.write_errors, 0) + + def test_repeats_on_the_configured_cadence(self): + """Read again on the next interval rather than once at startup.""" + self.grabber.start() + self._await_samples(minimum=3) + first = len(self.sink.keywords()) + deadline = time.monotonic() + SETTLE_TIMEOUT_S + while time.monotonic() < deadline: + if len(self.sink.keywords()) > first: + return + time.sleep(0.05) + self.fail("the collection never ran a second tick") + + def test_stop_is_clean_with_nothing_left_queued(self): + """Drain on the way out, so the last tick is not dropped.""" + self.grabber.start() + self._await_samples() + self.grabber.stop() + self.assertEqual(self.grabber.counters.dropped_batches, 0) + + +def _collections(peer: str) -> dict: + """Return a one-collection config reading everything on ``peer``.""" + return { + "target": { + "peer": peer, + "keywords": ["%"], + "interval_s": 0.5, + "timeout_s": 0.3, + "refresh_s": 60.0, + }, + } + + +class ZmqKeygrabberTests(_Bases.KeygrabberCases): + """Keygrabber cases over ZMQ.""" + + target_peer = f"hsfei.{_ZmqFixtureDaemon.peer_id}" @classmethod def setUpClass(cls): - cls.target_endpoint = _free_endpoint() - cls.target = _FixtureDaemon() - cls.target.bind = cls.target_endpoint + cls.endpoint = _free_endpoint() + cls.target = _ZmqFixtureDaemon() + cls.target.bind = cls.endpoint cls.target.start() @classmethod def tearDownClass(cls): cls.target.stop() - def setUp(self) -> None: - self.sink = _RecordingSink() - self.grabber = _TestKeygrabber.from_config({ - "peer_id": "keygrabber", + def grabber_config(self) -> dict: + return { + "peer_id": "keygrabberzmq", "group_id": "hispec", "transport": "zmq", "bind": _free_endpoint(), - "address_book": {"hsfei.kgtarget": self.target_endpoint}, + "address_book": {self.target_peer: self.endpoint}, + "discovery_enabled": False, + "sink": {"type": "recording"}, + "workers": 2, + "collections": _collections(self.target_peer), + } + + +@unittest.skipUnless(_broker_available(), "no RabbitMQ broker reachable at amqp://localhost") +class RabbitMQKeygrabberTests(_Bases.KeygrabberCases): + """Keygrabber cases over RabbitMQ, the default transport.""" + + target_peer = f"hsfei.{_RabbitFixtureDaemon.peer_id}" + + @classmethod + def setUpClass(cls): + cls.target = _RabbitFixtureDaemon() + cls.target.start() + + @classmethod + def tearDownClass(cls): + cls.target.stop() + + def grabber_config(self) -> dict: + return { + "peer_id": "keygrabberrmq", + "group_id": "hispec", + "transport": "rabbitmq", + "rabbitmq_url": RABBITMQ_URL, "discovery_enabled": False, "sink": {"type": "recording"}, "workers": 2, - "collections": { - "target": { - "peer": "hsfei.kgtarget", - "keywords": ["%"], - "interval_s": 0.5, - "timeout_s": 0.3, - "refresh_s": 60.0, - }, - }, - }) - self.grabber.sink = self.sink - - def tearDown(self) -> None: - self.grabber.stop() - - def _await_samples(self, minimum: int = 3) -> None: - """Block until enough samples arrive, rather than sleeping a guess.""" - deadline = time.monotonic() + SETTLE_TIMEOUT_S - while time.monotonic() < deadline: - if len(self.sink.keywords()) >= minimum: - return - time.sleep(0.05) - self.fail(f"only {len(self.sink.keywords())} samples within " - f"{SETTLE_TIMEOUT_S}s: {self.sink.keywords()}") - - def test_collects_the_selected_keywords(self): - """Read every selected keyword and write it to the sink.""" - self.grabber.start() - self._await_samples() - self.assertEqual(set(self.sink.keywords()), - {"positionvalue", "isconnected", "status"}) - - def test_default_exclusions_are_not_collected(self): - """Keep uptime and lasterror out of a select-everything collection.""" - self.grabber.start() - self._await_samples() - recorded = set(self.sink.keywords()) - self.assertNotIn("uptime", recorded) - self.assertNotIn("lasterror", recorded) - - def test_samples_carry_their_tags_and_values(self): - """Carry group, peer, units and value through to the sink.""" - self.grabber.start() - self._await_samples() - sample = next(s for s in self.sink.snapshot() - if s.keyword == "positionvalue") - self.assertEqual((sample.group, sample.peer), ("hsfei", "kgtarget")) - self.assertEqual(sample.value, 7.5) - self.assertEqual(sample.units, "mm") - - def test_points_written_is_counted(self): - """Report what the sink actually stored.""" - self.grabber.start() - self._await_samples() - self.assertGreater(self.grabber.counters.points_written, 0) - self.assertEqual(self.grabber.counters.write_errors, 0) - - def test_repeats_on_the_configured_cadence(self): - """Read again on the next interval rather than once at startup.""" - self.grabber.start() - self._await_samples(minimum=3) - first = len(self.sink.keywords()) - deadline = time.monotonic() + SETTLE_TIMEOUT_S - while time.monotonic() < deadline: - if len(self.sink.keywords()) > first: - return - time.sleep(0.05) - self.fail("the collection never ran a second tick") - - def test_stop_is_clean_with_nothing_left_queued(self): - """Drain on the way out, so the last tick is not dropped.""" - self.grabber.start() - self._await_samples() - self.grabber.stop() - self.assertEqual(self.grabber.counters.dropped_batches, 0) + "collections": _collections(self.target_peer), + } if __name__ == "__main__": From 4849cfae4992aba2f9ee30e6fe5d9bf3dd95d1ca Mon Sep 17 00:00:00 2001 From: Mike Langmayr <1809691+mikelangmayr@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:20:54 -0700 Subject: [PATCH 3/3] Release a claim when the pool is gone, and let the writer drain itself --- libby/keygrabber/daemon.py | 55 +++++++++++++++++++++++++++------ tests/test_keygrabber_daemon.py | 18 +++++++++++ 2 files changed, 63 insertions(+), 10 deletions(-) diff --git a/libby/keygrabber/daemon.py b/libby/keygrabber/daemon.py index ea5a510..4ad573c 100644 --- a/libby/keygrabber/daemon.py +++ b/libby/keygrabber/daemon.py @@ -34,6 +34,11 @@ # into a SIGKILL. DRAIN_DEADLINE_S = 5.0 +# How long ``on_stop`` waits for a thread to finish. Outlasts the drain +# deadline, so a writer that is draining gets to finish rather than being +# abandoned mid-batch. +STOP_JOIN_TIMEOUT_S = DRAIN_DEADLINE_S + WRITER_POLL_S * 2 + class Counters: # pylint: disable=too-few-public-methods """Tallies the daemon reports, guarded for cross-thread increments. @@ -81,6 +86,7 @@ def __init__(self) -> None: self._queue: "queue.Queue[Tuple[Sample, ...]]" = queue.Queue() self._pool: Optional[ThreadPoolExecutor] = None self._threads: List[threading.Thread] = [] + self._writer_thread: Optional[threading.Thread] = None self._halt = threading.Event() self._in_flight: set[str] = set() self._in_flight_lock = threading.Lock() @@ -100,24 +106,36 @@ def on_start(self, libby: Libby) -> None: self._pool = ThreadPoolExecutor(max_workers=self._settings.workers, thread_name_prefix="keygrabber-read") self._halt.clear() - self._spawn(self._write_loop, "keygrabber-write") + self._writer_thread = self._spawn(self._write_loop, "keygrabber-write") self._spawn(self._schedule_loop, "keygrabber-schedule") self.logger.info("collecting %d collections with %d workers", len(self._collections), self._settings.workers) def on_stop(self, libby: Optional[Libby] = None) -> None: - """Stop the threads, then give the retry queue a bounded chance to drain.""" + """Stop the threads and close the sink once the writer has drained it. + + The draining happens on the writer thread rather than here, so only one + thread is ever inside the sink. This method just waits for it. + """ self._halt.set() if self._pool is not None: self._pool.shutdown(wait=True) self._pool = None for thread in self._threads: - thread.join(timeout=WRITER_POLL_S * 8) + thread.join(timeout=STOP_JOIN_TIMEOUT_S) self._threads = [] - if self._writer is not None: - self._drain() - self._writer.close() - self._writer = None + + writer, self._writer = self._writer, None + if writer is None: + return + if self._writer_thread is not None and self._writer_thread.is_alive(): + # Closing while the writer is still inside a sink call would put + # two threads in one sink; leave it to process exit instead + self.logger.error( + "writer thread did not finish within %.1fs; leaving the sink open", + STOP_JOIN_TIMEOUT_S) + return + writer.close() def make_sink(self) -> Sink: """Build the configured sink. @@ -129,10 +147,11 @@ def make_sink(self) -> Sink: raise LibbyError("keygrabber config has not been parsed yet") return build_sink(self._settings.sink) - def _spawn(self, target: Callable[[], None], name: str) -> None: + def _spawn(self, target: Callable[[], None], name: str) -> threading.Thread: thread = threading.Thread(target=target, name=name, daemon=True) thread.start() self._threads.append(thread) + return thread ### Scheduling @@ -175,10 +194,24 @@ def _submit(self, collection: Collection) -> None: return self._in_flight.add(collection.name) + if not self._dispatch(collection): + # Shut down between the claim and the submit, so give the claim + # back rather than leaving the collection marked busy for good + with self._in_flight_lock: + self._in_flight.discard(collection.name) + + def _dispatch(self, collection: Collection) -> bool: + """Submit a tick, returning False once the pool can no longer take one.""" pool = self._pool if pool is None: - return - pool.submit(self._run_tick, collection) + return False + try: + pool.submit(self._run_tick, collection) + return True + except RuntimeError: + # ThreadPoolExecutor.submit raises once shutdown() has been called, + # which races the scheduler thread on the way out + return False def _run_tick(self, collection: Collection) -> None: """Resolve if due, read once, and hand the samples to the writer.""" @@ -218,6 +251,8 @@ def _write_loop(self) -> None: continue self._write(batch) self._flush() + # Halted: drain here, on the one thread allowed to touch the sink + self._drain() def _write(self, batch: Tuple[Sample, ...]) -> None: writer = self._writer diff --git a/tests/test_keygrabber_daemon.py b/tests/test_keygrabber_daemon.py index 1d95537..e2834a9 100644 --- a/tests/test_keygrabber_daemon.py +++ b/tests/test_keygrabber_daemon.py @@ -184,6 +184,24 @@ def test_points_written_is_counted(self): self.assertGreater(self.grabber.counters.points_written, 0) self.assertEqual(self.grabber.counters.write_errors, 0) + def test_submit_to_a_shut_down_pool_releases_the_claim(self): + """Survive the pool shutting down between a claim and its submit. + + ``ThreadPoolExecutor.submit`` raises once ``shutdown`` has been + called, which races the scheduler thread on the way out. The claim + has to come back, or the collection stays marked busy and every + later tick is skipped. + """ + self.grabber.start() + self._await_samples() + # pylint: disable=protected-access + collection = self.grabber._collections[0] + self.grabber._pool.shutdown(wait=True) # left non-None on purpose + self.grabber._in_flight.discard(collection.name) + + self.grabber._submit(collection) # must not raise + self.assertNotIn(collection.name, self.grabber._in_flight) + def test_repeats_on_the_configured_cadence(self): """Read again on the next interval rather than once at startup.""" self.grabber.start()