From a2e29a455d547c63fecc502e2089d78e4dc455fa Mon Sep 17 00:00:00 2001 From: David Perl Date: Fri, 11 Sep 2026 17:28:51 +0200 Subject: [PATCH] wip feat: add redis-based queue --- pyproject.toml | 2 +- src/aare/queue/__init__.py | 0 src/aare/queue/client.py | 101 +++++++++ src/aare/queue/models.py | 101 +++++++++ tests/unit/queue/test_queue_client.py | 287 ++++++++++++++++++++++++++ tests/unit/queue/test_queue_models.py | 131 ++++++++++++ 6 files changed, 621 insertions(+), 1 deletion(-) create mode 100644 src/aare/queue/__init__.py create mode 100644 src/aare/queue/client.py create mode 100644 src/aare/queue/models.py create mode 100644 tests/unit/queue/test_queue_client.py create mode 100644 tests/unit/queue/test_queue_models.py diff --git a/pyproject.toml b/pyproject.toml index 198f79a5..df6433af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ name = "aaredaq" version = "0.31.1" description = "AareDAQ (with GUI)" readme = "README.md" -requires-python = ">=3.11" +requires-python = ">=3.12" dependencies = [ "uv", "aarecommon>=0.9.0", diff --git a/src/aare/queue/__init__.py b/src/aare/queue/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/aare/queue/client.py b/src/aare/queue/client.py new file mode 100644 index 00000000..cdc2bafc --- /dev/null +++ b/src/aare/queue/client.py @@ -0,0 +1,101 @@ +from collections.abc import Callable +from threading import Event, RLock, Thread +from uuid import uuid4 + +from aarecommon.config.logger import setup_logger +from aarecommon.models.beamline import MXBeamline +from redis.client import Redis + +from aare.queue.models import ( + QUEUE_KEY_SUFFIX, + QUEUE_NOTIFICATION_CHANNEL_SUFFIX, + AareQueueItem, + parse_queue_item, +) + +type CallbackType = Callable[[], None] + + +class AareQueueClient: + def __init__(self, beamline: MXBeamline, redis: Redis, logger_name="AareDAQ") -> None: + self.logger = setup_logger(logger_name) + self._redis = redis + self._key = str(beamline) + QUEUE_KEY_SUFFIX + self._stop_event = Event() + self._callback_event = Event() + self._callback_lock = RLock() + self._callbacks: dict[str, CallbackType] = {} + self._notification_key = str(beamline) + QUEUE_NOTIFICATION_CHANNEL_SUFFIX + self._notification_sub = self._redis.pubsub(ignore_subscribe_messages=True) + self._notification_sub.subscribe( + **{self._notification_key: lambda _: self._callback_event.set()} + ) + self._notification_thread = self._notification_sub.run_in_thread( + sleep_time=0.01, daemon=True + ) + self._callback_thread = Thread( + target=self._queue_notification_callback_thread, + name="AareQueueClient-callbacks", + daemon=True, + ) + self._callback_thread.start() + + def _try_run_callback(self, cb: CallbackType): + try: + cb() + except Exception: + self.logger.exception(f"Exception during queue update callback excecution: {cb}.") + + def _queue_notification_callback_thread(self): + while not self._stop_event.is_set(): + if not self._callback_event.wait(0.05): + continue + self._callback_event.clear() + with self._callback_lock: + # Snapshot: a callback may deregister itself, and the RLock is + # reentrant so that would mutate the dict mid-iteration. + callbacks = list(self._callbacks.values()) + for cb in callbacks: + self._try_run_callback(cb) + + def add_update_callback(self, callback: CallbackType) -> str: + """Register a callback to be called when the queue is updated. Returns an ID which can be used to deregister. + + The callback runs on this client's dispatch thread, not the caller's. Qt consumers must propgate to the GUI + thread via signal rather than touching widgets directly. + """ + with self._callback_lock: + cb_id = str(uuid4()) + self._callbacks[cb_id] = callback + return cb_id + + def remove_update_callback(self, id: str): + """Provide the ID returned from registration to deregister a callback.""" + with self._callback_lock: + if id not in self._callbacks: + raise ValueError( + f"No callback with the ID {id} has been registered for queue updates." + ) + del self._callbacks[id] + + def get_all(self) -> list[AareQueueItem]: + """Read the whole queue in order. + + Entries with an unrecognised ``kind`` come back as UnknownInstruction + rather than failing the read. Entries that are not even a JSON object + are skipped and logged: one corrupt write should not blank the queue + for every consumer. + """ + items: list[AareQueueItem] = [] + for raw in self._redis.lrange(self._key, 0, -1): # pyright: ignore[reportGeneralTypeIssues] + try: + items.append(parse_queue_item(raw)) + except (ValueError, TypeError): + self.logger.exception(f"Skipping unreadable queue entry: {raw!r}") + return items + + def shutdown(self): + self._stop_event.set() + self._notification_thread.stop() + self._notification_thread.join(timeout=2.0) + self._callback_thread.join(timeout=2.0) diff --git a/src/aare/queue/models.py b/src/aare/queue/models.py new file mode 100644 index 00000000..b211e183 --- /dev/null +++ b/src/aare/queue/models.py @@ -0,0 +1,101 @@ +import json +from enum import StrEnum +from typing import Annotated, Any, Literal +from uuid import UUID, uuid4 + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError + +QUEUE_KEY_SUFFIX = ":SampleQueue" +QUEUE_NOTIFICATION_CHANNEL_SUFFIX = ":SampleQueueNotification" + + +class QueueItemKind(StrEnum): + RASTER = "raster" + SCREENING = "screening" + ROTATION = "rotation" + + +class QueueInstructionBase(BaseModel): + """Fields shared by every queue instruction.""" + + id: UUID = Field(default_factory=uuid4) + sample_id: int + + +class RasterInstruction(QueueInstructionBase): + kind: Literal[QueueItemKind.RASTER] = QueueItemKind.RASTER + + +class ScreeningInstruction(QueueInstructionBase): + kind: Literal[QueueItemKind.SCREENING] = QueueItemKind.SCREENING + + +class RotationInstruction(QueueInstructionBase): + kind: Literal[QueueItemKind.ROTATION] = QueueItemKind.ROTATION + + +KnownInstruction = Annotated[ + RasterInstruction | ScreeningInstruction | RotationInstruction, Field(discriminator="kind") +] + + +class UnknownInstruction(QueueInstructionBase): + """An instruction this process is too old (or too foreign) to understand. + + Fields we do not know about are kept as pydantic extras, so they serialize + back out untouched and a reader that cannot interpret the instruction can + still reorder the queue without destroying it. + + ``id`` drops the base's default: a stored instruction always carries one + already, and minting a fresh one on read would silently break the identity + that remove/reorder depend on. + """ + + model_config = ConfigDict(extra="allow") + + id: UUID # pyright: ignore[reportGeneralTypeIssues] + kind: str + + +AareQueueItem = KnownInstruction | UnknownInstruction + +_KNOWN_ADAPTER: TypeAdapter[Any] = TypeAdapter(KnownInstruction) +_KNOWN_KINDS = frozenset(kind.value for kind in QueueItemKind) + + +def _as_unknown(raw: str | bytes, known_kind_error: Exception) -> UnknownInstruction: + try: + payload = json.loads(raw) + except json.JSONDecodeError as e: + raise ValueError(f"Queue entry is not valid JSON: {e}") from e + if not isinstance(payload, dict): + raise TypeError(f"Queue entry is not a JSON object: {type(payload).__name__}") + if payload.get("kind") in _KNOWN_KINDS: + raise ValueError( + f"Malformed {payload['kind']!r} instruction: {known_kind_error}" + ) from known_kind_error + try: + return UnknownInstruction.model_validate(payload) + except ValidationError as e: + raise ValueError(f"Queue entry is not a usable instruction: {e}") from e + + +def parse_queue_item(raw: str | bytes) -> AareQueueItem: + """Decode one stored queue entry, falling back to UnknownInstruction. + + Raises: + ValueError: the entry is not valid JSON, carries no usable id/kind, or + is a malformed instance of a kind we do recognise -- all corrupt + rather than merely unrecognised. + TypeError: the entry is valid JSON but not an object, so there is no + payload a forward-compatible reader could preserve. + """ + try: + return _KNOWN_ADAPTER.validate_json(raw) + except ValidationError as e: + return _as_unknown(raw, e) + + +def dump_queue_item(item: AareQueueItem) -> str: + """Encode a queue item for storage. Unknown items re-emit their original payload.""" + return item.model_dump_json() diff --git a/tests/unit/queue/test_queue_client.py b/tests/unit/queue/test_queue_client.py new file mode 100644 index 00000000..93108863 --- /dev/null +++ b/tests/unit/queue/test_queue_client.py @@ -0,0 +1,287 @@ +import gc +import json +import threading +from uuid import UUID, uuid4 + +import fakeredis +import pytest +from aarecommon.models.beamline import MXBeamline + +from aare.queue.client import AareQueueClient +from aare.queue.models import ( + QUEUE_KEY_SUFFIX, + QUEUE_NOTIFICATION_CHANNEL_SUFFIX, + RasterInstruction, + RotationInstruction, + UnknownInstruction, + dump_queue_item, +) + +BEAMLINE = MXBeamline.X10SA +QUEUE_KEY = str(BEAMLINE) + QUEUE_KEY_SUFFIX +NOTIFY_CHANNEL = str(BEAMLINE) + QUEUE_NOTIFICATION_CHANNEL_SUFFIX + +# Callbacks run on the client's dispatch thread, so tests wait on an Event +# rather than sleeping. Generous enough for a loaded CI box, but only ever +# paid in full when something is actually broken. +DISPATCH_TIMEOUT_S = 2.0 +# Waiting for something NOT to happen has to be a real sleep; keep it short. +QUIET_PERIOD_S = 0.3 + + +class Recorder: + """A callback that records its calls, as a bound method. + + Bound methods are the realistic case (a Qt panel registering + ``self.on_queue_update``) and the one that a weak callback registry drops. + """ + + def __init__(self): + self.calls = 0 + self.called = threading.Event() + + def on_queue_update(self): + self.calls += 1 + self.called.set() + + def wait(self, timeout=DISPATCH_TIMEOUT_S) -> bool: + return self.called.wait(timeout) + + +@pytest.fixture +def redis(): + return fakeredis.FakeRedis(decode_responses=True) + + +@pytest.fixture +def client(redis): + c = AareQueueClient(BEAMLINE, redis) + yield c + c.shutdown() + + +@pytest.fixture +def recorder(client): + r = Recorder() + client.add_update_callback(r.on_queue_update) + return r + + +def notify(redis): + """Stand in for whatever writes the queue and announces the change.""" + redis.publish(NOTIFY_CHANNEL, "1") + + +# --------------------------------------------------------------------------- +# construction and notification plumbing +# --------------------------------------------------------------------------- + + +def test_client_constructs_with_a_subscribed_handler(client): + """Handlers must be passed to subscribe() as kwargs. A dict passed + positionally is read as a list of channel names, leaving the channel + handler-less and making run_in_thread raise PubSubError.""" + assert NOTIFY_CHANNEL in client._notification_sub.channels + assert client._notification_sub.channels[NOTIFY_CHANNEL] is not None + + +def test_callback_runs_when_a_notification_arrives(redis, recorder): + notify(redis) + assert recorder.wait(), "callback never ran" + assert recorder.calls == 1 + + +def test_callback_does_not_rerun_without_a_new_notification(redis, recorder): + """The dispatch event must be cleared after waking. Left set, the loop + spins and replays every callback continuously.""" + notify(redis) + assert recorder.wait() + + calls_after_first = recorder.calls + threading.Event().wait(QUIET_PERIOD_S) + assert recorder.calls == calls_after_first + + +def test_each_notification_dispatches_again(redis, recorder): + for expected in (1, 2, 3): + recorder.called.clear() + notify(redis) + assert recorder.wait(), f"no dispatch for notification {expected}" + assert recorder.calls == expected + + +# --------------------------------------------------------------------------- +# callback registry +# --------------------------------------------------------------------------- + + +def test_bound_method_callback_survives_registration(client): + """A bound method is a fresh object per attribute access, so a weak + registry collects it the instant registration returns.""" + recorder = Recorder() + cb_id = client.add_update_callback(recorder.on_queue_update) + gc.collect() + assert cb_id in client._callbacks + + +def test_lambda_callback_survives_registration(client): + cb_id = client.add_update_callback(lambda: None) + gc.collect() + assert cb_id in client._callbacks + + +def test_every_registered_callback_is_dispatched(redis, client): + recorders = [Recorder() for _ in range(3)] + for r in recorders: + client.add_update_callback(r.on_queue_update) + + notify(redis) + + for i, r in enumerate(recorders): + assert r.wait(), f"callback {i} never ran" + + +def test_removed_callback_stops_being_dispatched(redis, client): + removed, kept = Recorder(), Recorder() + removed_id = client.add_update_callback(removed.on_queue_update) + client.add_update_callback(kept.on_queue_update) + + client.remove_update_callback(removed_id) + notify(redis) + + assert kept.wait(), "surviving callback never ran" + assert removed.calls == 0 + + +def test_removing_an_unregistered_callback_raises(client): + with pytest.raises(ValueError): + client.remove_update_callback(str(uuid4())) + + +def test_callback_may_deregister_itself_during_dispatch(redis, client): + """The dispatch loop iterates a snapshot; mutating the live dict under the + reentrant lock would raise RuntimeError mid-dispatch.""" + done = threading.Event() + ids = {} + + def self_removing(): + client.remove_update_callback(ids["self"]) + done.set() + + ids["self"] = client.add_update_callback(self_removing) + later = Recorder() + client.add_update_callback(later.on_queue_update) + + notify(redis) + + assert done.wait(DISPATCH_TIMEOUT_S), "self-removing callback never ran" + assert later.wait(), "dispatch aborted before reaching the next callback" + + +def test_a_raising_callback_does_not_stop_the_others(redis, client): + def boom(): + raise RuntimeError("callback blew up") + + client.add_update_callback(boom) + survivor = Recorder() + client.add_update_callback(survivor.on_queue_update) + + notify(redis) + + assert survivor.wait(), "one bad callback killed the dispatch" + + +def test_dispatch_survives_a_raising_callback_across_notifications(redis, client): + """An exception must not take the dispatch thread down with it.""" + + def boom(): + raise RuntimeError("callback blew up") + + client.add_update_callback(boom) + survivor = Recorder() + client.add_update_callback(survivor.on_queue_update) + + notify(redis) + assert survivor.wait() + + survivor.called.clear() + notify(redis) + assert survivor.wait(), "dispatch thread died on the first exception" + + +# --------------------------------------------------------------------------- +# reads +# --------------------------------------------------------------------------- + + +def test_get_all_is_empty_for_an_untouched_queue(client): + assert client.get_all() == [] + + +def test_get_all_preserves_insertion_order(redis, client): + items = [RasterInstruction(sample_id=i) for i in range(4)] + redis.rpush(QUEUE_KEY, *[dump_queue_item(i) for i in items]) + + assert [i.id for i in client.get_all()] == [i.id for i in items] + + +def test_get_all_keeps_unrecognised_instructions(redis, client): + future = {"id": str(uuid4()), "kind": "xrf_map", "sample_id": 9, "dwell_ms": 50} + redis.rpush(QUEUE_KEY, json.dumps(future)) + + (item,) = client.get_all() + assert isinstance(item, UnknownInstruction) + assert item.kind == "xrf_map" + + +def test_get_all_skips_corrupt_entries_without_losing_the_rest(redis, client): + """One bad write must not blank the queue for every consumer.""" + good = RotationInstruction(sample_id=1) + redis.rpush( + QUEUE_KEY, + dump_queue_item(good), + "NOT JSON AT ALL", + "[1, 2, 3]", + json.dumps({"kind": "xrf_map", "sample_id": 9}), # no id + dump_queue_item(RasterInstruction(sample_id=2)), + ) + + items = client.get_all() + + assert len(items) == 2 + assert items[0].id == good.id + + +def test_every_item_read_back_is_targetable_by_id(redis, client): + """Identity ops (remove, reorder) rely on this holding for unknown items too.""" + redis.rpush( + QUEUE_KEY, + dump_queue_item(RasterInstruction(sample_id=1)), + json.dumps({"id": str(uuid4()), "kind": "xrf_map", "sample_id": 2}), + ) + + items = client.get_all() + + assert len(items) == 2, "an unknown instruction was dropped, making this vacuous" + assert all(isinstance(item.id, UUID) for item in items) + + +# --------------------------------------------------------------------------- +# shutdown +# --------------------------------------------------------------------------- + + +def test_shutdown_stops_both_threads(redis): + client = AareQueueClient(BEAMLINE, redis) + client.shutdown() + + assert not client._callback_thread.is_alive() + assert not client._notification_thread.is_alive() + + +def test_shutdown_is_safe_with_callbacks_registered(redis): + client = AareQueueClient(BEAMLINE, redis) + client.add_update_callback(Recorder().on_queue_update) + client.shutdown() + + assert not client._callback_thread.is_alive() diff --git a/tests/unit/queue/test_queue_models.py b/tests/unit/queue/test_queue_models.py new file mode 100644 index 00000000..8709f972 --- /dev/null +++ b/tests/unit/queue/test_queue_models.py @@ -0,0 +1,131 @@ +import json +from uuid import UUID, uuid4 + +import pytest + +from aare.queue.models import ( + QueueInstructionBase, + QueueItemKind, + RasterInstruction, + RotationInstruction, + ScreeningInstruction, + UnknownInstruction, + dump_queue_item, + parse_queue_item, +) + +KNOWN_INSTRUCTIONS = [ + (RasterInstruction, QueueItemKind.RASTER), + (ScreeningInstruction, QueueItemKind.SCREENING), + (RotationInstruction, QueueItemKind.ROTATION), +] + +# A plausible instruction from a producer newer than this reader: a kind we have +# never heard of, plus fields we know nothing about. +FUTURE_INSTRUCTION = { + "id": "11111111-2222-3333-4444-555555555555", + "kind": "xrf_map", + "sample_id": 9, + "grid_steps": [4, 4], + "dwell_ms": 50, +} + + +@pytest.mark.parametrize("cls,kind", KNOWN_INSTRUCTIONS, ids=lambda v: getattr(v, "value", None)) +def test_known_instruction_defaults_its_own_kind(cls, kind): + assert cls(sample_id=1).kind == kind + + +def test_ids_are_generated_and_distinct(): + ids = {RasterInstruction(sample_id=1).id for _ in range(5)} + assert len(ids) == 5 + + +@pytest.mark.parametrize("cls,kind", KNOWN_INSTRUCTIONS, ids=lambda v: getattr(v, "value", None)) +def test_known_instruction_round_trips_to_its_own_class(cls, kind): + original = cls(sample_id=3) + restored = parse_queue_item(dump_queue_item(original)) + assert type(restored) is cls + assert restored.sample_id == 3 + + +@pytest.mark.parametrize("cls,kind", KNOWN_INSTRUCTIONS, ids=lambda v: getattr(v, "value", None)) +def test_round_trip_preserves_the_id(cls, kind): + """Identity ops target the id, so a read must never regenerate it.""" + original = cls(sample_id=3) + assert parse_queue_item(dump_queue_item(original)).id == original.id + + +def test_unrecognised_kind_becomes_an_unknown_instruction(): + item = parse_queue_item(json.dumps(FUTURE_INSTRUCTION)) + assert isinstance(item, UnknownInstruction) + assert item.kind == "xrf_map" + assert item.id == UUID(FUTURE_INSTRUCTION["id"]) + + +def test_unknown_instruction_survives_a_read_write_cycle_verbatim(): + """Reordering rewrites every entry. A reader that cannot interpret an item + must not drop the fields it did not understand.""" + item = parse_queue_item(json.dumps(FUTURE_INSTRUCTION)) + assert json.loads(dump_queue_item(item)) == FUTURE_INSTRUCTION + + +def test_unknown_instruction_is_a_queue_instruction(): + """It shares the base, so consumers reach its common fields the same way + they would for a kind they understand.""" + item = parse_queue_item(json.dumps(FUTURE_INSTRUCTION)) + assert isinstance(item, QueueInstructionBase) + assert item.id == UUID(FUTURE_INSTRUCTION["id"]) + assert item.sample_id == FUTURE_INSTRUCTION["sample_id"] + + +def test_unknown_fields_are_kept_as_extras(): + """Held as pydantic extras rather than a raw blob, which is what lets them + serialize back out without a custom serializer.""" + item = parse_queue_item(json.dumps(FUTURE_INSTRUCTION)) + assert item.model_extra == { + "grid_steps": FUTURE_INSTRUCTION["grid_steps"], + "dwell_ms": FUTURE_INSTRUCTION["dwell_ms"], + } + + +def test_unknown_instruction_does_not_regenerate_a_missing_id(): + """The base defaults id, but a stored entry always has one already -- + minting a fresh one on read would silently break remove/reorder.""" + with pytest.raises(ValueError): + parse_queue_item(json.dumps({"kind": "xrf_map", "sample_id": 9})) + + +@pytest.mark.parametrize( + "payload", + [ + {"kind": "xrf_map", "sample_id": 9}, + {"id": "not-a-uuid", "kind": "xrf_map", "sample_id": 9}, + {"id": str(uuid4()), "sample_id": 9}, + {"id": str(uuid4()), "kind": "xrf_map"}, + ], + ids=["no-id", "non-uuid-id", "no-kind", "no-sample-id"], +) +def test_entries_missing_a_base_field_are_rejected(payload): + """Every instruction inherits id and sample_id and declares a kind, so an + entry missing any of them is corrupt rather than merely unrecognised -- it + must not parse as Unknown.""" + with pytest.raises(ValueError): + parse_queue_item(json.dumps(payload)) + + +def test_malformed_json_is_rejected(): + with pytest.raises(ValueError): + parse_queue_item("NOT JSON AT ALL") + + +def test_non_object_json_is_rejected(): + with pytest.raises(TypeError): + parse_queue_item("[1, 2, 3]") + + +def test_a_known_kind_with_a_bad_payload_is_not_smuggled_in_as_unknown(): + """Falling back on any ValidationError would turn a malformed raster item + into an opaque Unknown that silently round-trips broken data.""" + with pytest.raises(ValueError): + parse_queue_item(json.dumps({"id": str(uuid4()), "kind": "raster"})) # no sample_id