wip feat: add redis-based queue

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",

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)

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()

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()

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