From dbdd7953c5cc9851a4a394ad11f43f03d15007e5 Mon Sep 17 00:00:00 2001 From: q Date: Sun, 18 Oct 2026 11:34:56 +0000 Subject: [PATCH] Wake the story completion loop on a native Pub/Sub scheduler hint Add the scheduler:wake hint contract and publisher, an opt-in event wait in periodic_loop that keeps the interval as backstop, and a supervised listener in scheduler-pipeline that routes valid story_completion hints. No business transition publishes a hint yet. --- ARCHITECTURE.md | 2 +- docs/CHANGELOG.md | 2 + docs/CONTRACTS.md | 5 + docs/contracts/scheduler-wake-hints.md | 50 +++ services/scheduler/src/pipeline.py | 13 +- services/scheduler/src/runtime.py | 25 +- .../src/tasks/scheduler_wake_listener.py | 69 ++++ .../src/tasks/story_completion_loop.py | 5 +- .../scheduler/tests/unit/_wake_doubles.py | 165 ++++++++++ .../unit/test_scheduler_pipeline_loops.py | 2 +- .../tests/unit/test_scheduler_wake.py | 302 ++++++++++++++++++ .../tests/unit/test_service_entrypoints.py | 79 ++++- shared/contracts/queues/scheduler_wake.py | 49 +++ shared/redis/scheduler_wake.py | 16 + shared/tests/test_scheduler_wake_contract.py | 50 +++ 15 files changed, 815 insertions(+), 19 deletions(-) create mode 100644 docs/contracts/scheduler-wake-hints.md create mode 100644 services/scheduler/src/tasks/scheduler_wake_listener.py create mode 100644 services/scheduler/tests/unit/_wake_doubles.py create mode 100644 services/scheduler/tests/unit/test_scheduler_wake.py create mode 100644 shared/contracts/queues/scheduler_wake.py create mode 100644 shared/redis/scheduler_wake.py create mode 100644 shared/tests/test_scheduler_wake_contract.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index f6e3009d6..542c3785a 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -74,7 +74,7 @@ taken from the default value. | `worker-broker` | The only service on both control-plane and worker networks. Authenticates per-worker credentials and brokers worker streams, sessions, status and Compose requests. | | `langgraph` | PO runtime (`po:input` consumer + ReactAgent). The same image also provides separate `engineering-worker`, `deploy-worker`, `qa-worker`, and `architect` container entrypoints; those roles are not the `langgraph` service itself. | | `architect` | Story→tasks LLM decomposition. Consumes `architect:queue`. A container of the `langgraph` image, not part of `scheduler` | -| `scheduler-pipeline` | Ten loops with independent failure boundaries: `task_dispatcher` (engineering admission/dispatch), `scaffold`, `story_completion`, `lifecycle_supervision`, `qa_routing`, `pr_ci` (merged-PR handling and CI-failure routing), `worker_reconciliation`, `temporary_access` cleanup, `owner_notifications` recovery and `story_supervision` (state-age watchdog, stage notices). A container of the shared `scheduler` image | +| `scheduler-pipeline` | Ten loops with independent failure boundaries: `task_dispatcher` (engineering admission/dispatch), `scaffold`, `story_completion`, `lifecycle_supervision`, `qa_routing`, `pr_ci` (merged-PR handling and CI-failure routing), `worker_reconciliation`, `temporary_access` cleanup, `owner_notifications` recovery and `story_supervision` (state-age watchdog, stage notices), plus the `scheduler_wake` Pub/Sub listener that wakes `story_completion` early. A container of the shared `scheduler` image | | `scheduler-infrastructure` | Fail-closed Time4VPS server sync, health checks, provisioner trigger and restart-safe result consumption. A container of the shared `scheduler` image | | `scheduler-maintenance` | GitHub project sync, analytics aggregation and queue cleanup. A container of the shared `scheduler` image | | `infra-service` | An Ansible runner and SSH operations | diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 758d23e4d..75a34b6ab 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -11,6 +11,8 @@ See the CHANGELOG rule in [AGENTS.md](../AGENTS.md). `USERS_GRANT_CAPABILITY` instead of redeploying the product, so QA access no longer costs a deploy. - A temporary-access operation takes its Run's dispatch claim before any external write, and every revoke re-derives the earlier-workflow stop from all the grant's Runs, so withdrawal or retry cannot skip either. +- scheduler-pipeline's `story_completion` loop also wakes on a native Pub/Sub hint on `scheduler:wake`; + hints are lossy and coalesced, the interval stays the backstop, and nothing publishes one yet. ## 2026-10-17 diff --git a/docs/CONTRACTS.md b/docs/CONTRACTS.md index 587b59815..4e3c13ed5 100644 --- a/docs/CONTRACTS.md +++ b/docs/CONTRACTS.md @@ -37,6 +37,7 @@ Feature-specific invariants live in focused guides. Read this index first, then | Generated-product template/package installation and QA | [Generated product kit and QA contracts](contracts/kit-template-and-qa.md) | | Runtime lifecycle/security invariants | [Lifecycle and security invariants](contracts/lifecycle-invariants.md) | | Managed target reconciliation/readiness | [Managed target readiness](contracts/managed-target-readiness.md) | +| Scheduler loop wake hints over native Pub/Sub | [Scheduler wake hints](contracts/scheduler-wake-hints.md) | ## Canonical vocabularies @@ -70,6 +71,10 @@ are repository-relative. | `po:response:{request_id}` | direct response | `queues/po.py` | PO consumer | telegram-bot | | `po:proactive` | `tg-bot-proactive` | `queues/po.py` | PO notification tools | telegram-bot | +The native Pub/Sub channel `scheduler:wake` (`queues/scheduler_wake.py`) is not a queue: its +lossy hints only shorten a scheduler-pipeline loop's wait, see +[Scheduler wake hints](contracts/scheduler-wake-hints.md). + ### Recipient and PO rules Telegram admits bounded concurrent updates through PTB's native update processor. diff --git a/docs/contracts/scheduler-wake-hints.md b/docs/contracts/scheduler-wake-hints.md new file mode 100644 index 000000000..3d5e845ea --- /dev/null +++ b/docs/contracts/scheduler-wake-hints.md @@ -0,0 +1,50 @@ +# Scheduler wake hints + +A wake hint asks one scheduler loop to run its next cycle now instead of at the end of its +interval. It is an optimisation of latency only: the loop's interval sweep stays the recovery +path, and the awakened cycle decides from the same durable Story/Task/Run state it always reads. + +## Wire contract + +- Transport: native Redis Pub/Sub (`PUBLISH`/`SUBSCRIBE`), channel `scheduler:wake`. It is not a + Redis Stream: nothing is stored, acknowledged, retried or replayed. +- Message: the JSON body of `SchedulerWakeHint` in `shared/contracts/queues/scheduler_wake.py`, + for example `{"target":"story_completion"}`. `target` is a `SchedulerWakeTarget`; no other + field is accepted. +- Targets: `story_completion` only. +- A hint carries routing only: no credentials, ids, correlation metadata or business state. It + grants no permission to transition a Story, Task or Run and never invokes a business action. +- Producers publish with `shared.redis.scheduler_wake.publish_scheduler_wake` after the durable + write the hint announces has committed. A publish failure loses only the hint. + +## Ownership + +| Target | Logical producer | Consumer | +|---|---|---| +| `story_completion` | none yet: no business transition publishes a hint | scheduler-pipeline `story_completion` loop | + +scheduler-pipeline runs one `scheduler_wake` listener beside its loops under the same +supervisor. It must hold an event for every target in the vocabulary, or it refuses to start. + +## Delivery and coalescing + +- Delivery is lossy: a hint published while no listener is subscribed, during a reconnect, or + before a restart is gone. The loop's `scheduler.dispatch_interval_seconds` sweep covers it. +- A hint sets the target loop's `asyncio.Event`. The loop clears the event when a cycle starts, + so any number of hints during a cycle, or between its end and the next wait, yield exactly one + prompt follow-up cycle. Cycles never overlap and hints queue no work or tasks. +- After a cycle the loop waits for its event or its interval, whichever comes first; the + interval is reread from configuration for every wait. There is no minimum delay on the wake + path. + +## Rejection and failure + +- A payload that is not text, exceeds 256 characters, is not JSON, names an unknown target or + carries an extra field is rejected and logged as `scheduler_wake_hint_rejected` with a reason + and validation shape only, never the payload. It wakes nothing and is never mapped to another + target. +- A Redis connect, subscribe or read failure, or an ended subscription, closes that Pub/Sub and + Redis client and resubscribes after one `scheduler.dispatch_interval_seconds`. The loops keep + their own interval deadlines and keep sweeping meanwhile. +- Cancellation propagates through the listener and the loop wait; each closes its Pub/Sub and + owned Redis client once and leaves no helper task running. diff --git a/services/scheduler/src/pipeline.py b/services/scheduler/src/pipeline.py index cd9983f76..420d010ab 100644 --- a/services/scheduler/src/pipeline.py +++ b/services/scheduler/src/pipeline.py @@ -1,9 +1,11 @@ """Pipeline scheduler service entry point.""" import asyncio +from functools import partial import structlog +from shared.contracts.queues.scheduler_wake import SchedulerWakeTarget from shared.log_config import setup_logging from . import runtime @@ -13,6 +15,7 @@ from .tasks.pr_ci_loop import pr_ci_loop from .tasks.qa_routing_loop import qa_routing_loop from .tasks.scaffold_loop import scaffold_loop +from .tasks.scheduler_wake_listener import scheduler_wake_listener from .tasks.story_completion_loop import story_completion_loop from .tasks.story_supervision_loop import story_supervision_loop from .tasks.task_dispatcher import task_dispatcher_loop @@ -26,11 +29,19 @@ async def main() -> None: setup_logging(service_name="scheduler-pipeline") logger.info("scheduler_pipeline_started") await runtime.initialize_configs(PIPELINE_REQUIRED_KEYS, service_name="scheduler-pipeline") + story_completion_wake = asyncio.Event() await runtime.run_workers( [ ("task_dispatcher", task_dispatcher_loop), ("scaffold", scaffold_loop), - ("story_completion", story_completion_loop), + ("story_completion", partial(story_completion_loop, wake=story_completion_wake)), + ( + "scheduler_wake", + partial( + scheduler_wake_listener, + {SchedulerWakeTarget.STORY_COMPLETION: story_completion_wake}, + ), + ), ("lifecycle_supervision", lifecycle_supervision_loop), ("qa_routing", qa_routing_loop), ("pr_ci", pr_ci_loop), diff --git a/services/scheduler/src/runtime.py b/services/scheduler/src/runtime.py index 61f961679..418bcbdf5 100644 --- a/services/scheduler/src/runtime.py +++ b/services/scheduler/src/runtime.py @@ -48,11 +48,14 @@ async def periodic_loop( stopped_event: str, redis_factory: Callable[[], RedisStreamClient] = RedisStreamClient, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + wake: asyncio.Event | None = None, ) -> None: """Run one Redis-backed periodic worker with a stable lifecycle. Each cycle starts with the message context unbound, so its messages never inherit - a correlation id bound before it, at startup or by the previous cycle. + a correlation id bound before it, at startup or by the previous cycle. With `wake`, + the wait after a cycle also ends when the event is set; a set during a cycle or + before the wait still runs one prompt follow-up cycle. """ redis_client = redis_factory() await redis_client.connect() @@ -60,18 +63,36 @@ async def periodic_loop( try: while True: unbind_message_context() + if wake is not None: + wake.clear() try: fields = await cycle(redis_client) if cycle_event is not None: logger.info(cycle_event, **dict(fields or {})) except Exception: logger.exception(error_event) - await sleep(interval()) + if wake is None: + await sleep(interval()) + else: + await _wait_for_wake(wake, sleep(interval())) finally: await redis_client.close() logger.info(stopped_event) +async def _wait_for_wake(wake: asyncio.Event, timer: Awaitable[None]) -> None: + """Return when `wake` is set or `timer` finishes, cancelling and awaiting the other.""" + waiters = {asyncio.ensure_future(wake.wait()), asyncio.ensure_future(timer)} + try: + done, _pending = await asyncio.wait(waiters, return_when=asyncio.FIRST_COMPLETED) + finally: + for waiter in waiters: + waiter.cancel() + await asyncio.gather(*waiters, return_exceptions=True) + for waiter in done: + waiter.result() + + async def run_workers( workers: Collection[tuple[str, Callable[[], Awaitable[None]]]], *, diff --git a/services/scheduler/src/tasks/scheduler_wake_listener.py b/services/scheduler/src/tasks/scheduler_wake_listener.py new file mode 100644 index 000000000..1887a57d5 --- /dev/null +++ b/services/scheduler/src/tasks/scheduler_wake_listener.py @@ -0,0 +1,69 @@ +"""Route native Pub/Sub wake hints to the scheduler loops that accept them.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Mapping + +from redis.exceptions import RedisError +import structlog + +from shared.contracts.queues.scheduler_wake import ( + SCHEDULER_WAKE_CHANNEL, + SchedulerWakeHintRejected, + SchedulerWakeTarget, + parse_scheduler_wake_hint, +) +from shared.redis import RedisStreamClient + +from .. import startup + +logger = structlog.get_logger(__name__) + + +def _reconnect_interval() -> int: + return startup.get_config().get_int("scheduler.dispatch_interval_seconds") + + +async def scheduler_wake_listener( + wakes: Mapping[SchedulerWakeTarget, asyncio.Event], + *, + redis_factory: Callable[[], RedisStreamClient] = RedisStreamClient, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, +) -> None: + """Set the event of each valid hint's target; Redis failures reconnect after one interval. + + Every target in the vocabulary needs an event, so a hint never lacks its receiver. + """ + if set(wakes) != set(SchedulerWakeTarget): + raise ValueError( + f"scheduler wake targets without an event: {set(SchedulerWakeTarget) - set(wakes)}" + ) + while True: + redis_client = redis_factory() + try: + await redis_client.connect() + pubsub = redis_client.redis.pubsub() + try: + await pubsub.subscribe(SCHEDULER_WAKE_CHANNEL) + logger.info("scheduler_wake_subscribed", channel=SCHEDULER_WAKE_CHANNEL) + async for message in pubsub.listen(): + if message["type"] == "message": + _route(message["data"], wakes) + logger.warning("scheduler_wake_subscription_ended") + finally: + await pubsub.aclose() + except (RedisError, OSError) as exc: + logger.warning("scheduler_wake_listener_disconnected", error_type=type(exc).__name__) + finally: + await redis_client.close() + await sleep(_reconnect_interval()) + + +def _route(data: object, wakes: Mapping[SchedulerWakeTarget, asyncio.Event]) -> None: + try: + hint = parse_scheduler_wake_hint(data) + except SchedulerWakeHintRejected as rejected: + logger.warning("scheduler_wake_hint_rejected", **rejected.diagnostics) + return + wakes[hint.target].set() diff --git a/services/scheduler/src/tasks/story_completion_loop.py b/services/scheduler/src/tasks/story_completion_loop.py index 10ac9818b..0b041ff26 100644 --- a/services/scheduler/src/tasks/story_completion_loop.py +++ b/services/scheduler/src/tasks/story_completion_loop.py @@ -18,8 +18,8 @@ def _story_completion_interval() -> int: return startup.get_config().get_int("scheduler.dispatch_interval_seconds") -async def story_completion_loop() -> None: - """Complete eligible stories from durable task/story state.""" +async def story_completion_loop(*, wake: asyncio.Event) -> None: + """Complete eligible stories from durable task/story state, early when `wake` is set.""" from ..clients.api import api_client async def cycle(redis_client: RedisStreamClient) -> dict[str, object]: @@ -35,4 +35,5 @@ async def cycle(redis_client: RedisStreamClient) -> dict[str, object]: stopped_event="story_completion_stopped", redis_factory=RedisStreamClient, sleep=asyncio.sleep, + wake=wake, ) diff --git a/services/scheduler/tests/unit/_wake_doubles.py b/services/scheduler/tests/unit/_wake_doubles.py new file mode 100644 index 000000000..122901696 --- /dev/null +++ b/services/scheduler/tests/unit/_wake_doubles.py @@ -0,0 +1,165 @@ +"""In-memory Redis Pub/Sub, controlled timers and a cycle probe for scheduler wake tests.""" + +from __future__ import annotations + +import asyncio + +from redis.exceptions import ConnectionError as RedisConnectionError + + +async def settle() -> None: + """Let every ready task run until the loop is idle again.""" + for _ in range(50): + await asyncio.sleep(0) + + +class Broker: + """Channel fan-out shared by every `StreamClient` a test hands out.""" + + def __init__(self) -> None: + self.subscriptions: list[PubSub] = [] + self.connect_failures = 0 + self.clients: list[StreamClient] = [] + + def client(self) -> StreamClient: + client = StreamClient(self) + self.clients.append(client) + return client + + def deliver(self, channel: str, data: object) -> None: + for pubsub in self.subscriptions: + if channel in pubsub.channels: + pubsub.inbox.put_nowait({"type": "message", "channel": channel, "data": data}) + + def drop_connections(self) -> None: + for pubsub in self.subscriptions: + pubsub.inbox.put_nowait(RedisConnectionError("connection reset")) + + async def subscribed(self) -> PubSub: + for _ in range(50): + live = [ + pubsub for pubsub in self.subscriptions if pubsub.channels and not pubsub.closed + ] + if live: + return live[-1] + await asyncio.sleep(0) + raise AssertionError("no live subscription") + + +class PubSub: + def __init__(self) -> None: + self.channels: set[str] = set() + self.inbox: asyncio.Queue[object] = asyncio.Queue() + self.closed = 0 + + async def subscribe(self, channel: str) -> None: + self.channels.add(channel) + self.inbox.put_nowait({"type": "subscribe", "channel": channel, "data": 1}) + + async def listen(self): + while True: + item = await self.inbox.get() + if isinstance(item, Exception): + raise item + yield item + + async def aclose(self) -> None: + self.closed += 1 + + +class Redis: + def __init__(self, broker: Broker) -> None: + self._broker = broker + + def pubsub(self) -> PubSub: + pubsub = PubSub() + self._broker.subscriptions.append(pubsub) + return pubsub + + async def publish(self, channel: str, data: str) -> int: + receivers = [ + p for p in self._broker.subscriptions if channel in p.channels and not p.closed + ] + self._broker.deliver(channel, data) + return len(receivers) + + +class StreamClient: + """Stands in for `RedisStreamClient`: owns one connection and counts its closes.""" + + def __init__(self, broker: Broker) -> None: + self._broker = broker + self._redis: Redis | None = None + self.closed = 0 + + async def connect(self) -> None: + if self._broker.connect_failures: + self._broker.connect_failures -= 1 + raise RedisConnectionError("redis unavailable") + self._redis = Redis(self._broker) + + async def close(self) -> None: + self.closed += 1 + self._redis = None + + @property + def redis(self) -> Redis: + if self._redis is None: + raise RuntimeError("not connected") + return self._redis + + +class Timers: + """A `sleep` that only finishes when the test fires it, recording each requested delay.""" + + def __init__(self) -> None: + self.requested: list[float] = [] + self._pending: list[asyncio.Event] = [] + self.cancelled = 0 + + async def sleep(self, seconds: float) -> None: + self.requested.append(seconds) + fired = asyncio.Event() + self._pending.append(fired) + try: + await fired.wait() + except asyncio.CancelledError: + self.cancelled += 1 + raise + finally: + self._pending.remove(fired) + + @property + def waiting(self) -> int: + return len(self._pending) + + def fire(self) -> None: + (pending,) = self._pending + pending.set() + + +class Cycles: + """A periodic-loop cycle that counts runs, detects overlap and can hold one run open.""" + + def __init__(self, *, fail_first: bool = False) -> None: + self.runs = 0 + self.active = 0 + self.max_active = 0 + self.fail_first = fail_first + self.hold: asyncio.Event | None = None + self.entered = asyncio.Event() + + async def __call__(self, _redis_client: object) -> dict[str, object]: + self.runs += 1 + self.active += 1 + self.max_active = max(self.max_active, self.active) + self.entered.set() + try: + if self.hold is not None: + hold, self.hold = self.hold, None + await hold.wait() + if self.fail_first and self.runs == 1: + raise RuntimeError("cycle failed") + return {"run": self.runs} + finally: + self.active -= 1 diff --git a/services/scheduler/tests/unit/test_scheduler_pipeline_loops.py b/services/scheduler/tests/unit/test_scheduler_pipeline_loops.py index 42851e83e..4c57b376e 100644 --- a/services/scheduler/tests/unit/test_scheduler_pipeline_loops.py +++ b/services/scheduler/tests/unit/test_scheduler_pipeline_loops.py @@ -67,7 +67,7 @@ async def test_story_completion_loop_continues_after_failure_and_owns_redis(monk ) with pytest.raises(asyncio.CancelledError): - await story_completion_loop.story_completion_loop() + await story_completion_loop.story_completion_loop(wake=asyncio.Event()) assert complete.await_count == 2 complete.assert_awaited_with(api, redis) diff --git a/services/scheduler/tests/unit/test_scheduler_wake.py b/services/scheduler/tests/unit/test_scheduler_wake.py new file mode 100644 index 000000000..218b414f1 --- /dev/null +++ b/services/scheduler/tests/unit/test_scheduler_wake.py @@ -0,0 +1,302 @@ +"""Wake hints shorten a periodic loop's wait; the interval stays the backstop.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +from _wake_doubles import Broker, Cycles, Timers, settle +import pytest + +from shared.contracts.queues.scheduler_wake import SCHEDULER_WAKE_CHANNEL, SchedulerWakeTarget +from shared.redis.scheduler_wake import publish_scheduler_wake + +STORY_COMPLETION = SchedulerWakeTarget.STORY_COMPLETION + + +def start_loop(cycles: Cycles, timers: Timers, wake: asyncio.Event, *, interval=lambda: 30): + from src import runtime + + client = MagicMock(connect=AsyncMock(), close=AsyncMock()) + log = MagicMock() + task = asyncio.create_task( + runtime.periodic_loop( + interval=interval, + cycle=cycles, + logger=log, + started_event="started", + cycle_event="cycle", + error_event="failed", + stopped_event="stopped", + redis_factory=lambda: client, + sleep=timers.sleep, + wake=wake, + ) + ) + return task, client, log + + +def start_listener(broker: Broker, timers: Timers, wakes, monkeypatch): + from src.tasks import scheduler_wake_listener + + monkeypatch.setattr(scheduler_wake_listener, "_reconnect_interval", lambda: 30) + return asyncio.create_task( + scheduler_wake_listener.scheduler_wake_listener( + wakes, redis_factory=broker.client, sleep=timers.sleep + ) + ) + + +async def stop(*tasks: asyncio.Task) -> None: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +async def publish(broker: Broker, target=STORY_COMPLETION) -> int: + publisher = broker.client() + await publisher.connect() + return await publish_scheduler_wake(publisher, target) + + +def other_tasks() -> set[asyncio.Task]: + return asyncio.all_tasks() - {asyncio.current_task()} + + +@pytest.mark.asyncio +async def test_a_wake_while_idle_runs_the_next_cycle_before_the_backstop(): + cycles, timers, wake = Cycles(), Timers(), asyncio.Event() + loop, _client, _log = start_loop(cycles, timers, wake) + await settle() + assert (cycles.runs, timers.waiting) == (1, 1) + + wake.set() + await settle() + + assert cycles.runs == 2 + assert timers.cancelled == 1 + assert timers.waiting == 1 + assert not wake.is_set() + await stop(loop) + + +@pytest.mark.asyncio +async def test_without_a_wake_the_backstop_runs_the_next_cycle_and_rereads_the_interval(): + cycles, timers, wake = Cycles(), Timers(), asyncio.Event() + intervals = iter([30, 5, 7]) + loop, _client, _log = start_loop(cycles, timers, wake, interval=lambda: next(intervals)) + await settle() + assert timers.requested == [5] + + timers.fire() + await settle() + + assert cycles.runs == 2 + assert timers.requested == [5, 7] + await stop(loop) + + +@pytest.mark.asyncio +async def test_wakes_during_a_cycle_run_exactly_one_prompt_follow_up(): + cycles, timers, wake = Cycles(), Timers(), asyncio.Event() + cycles.hold = release = asyncio.Event() + loop, _client, _log = start_loop(cycles, timers, wake) + await settle() + assert cycles.active == 1 + + for _ in range(3): + wake.set() + await settle() + release.set() + await settle() + + assert cycles.runs == 2 + assert cycles.max_active == 1 + assert timers.waiting == 1 + assert timers.requested == [30, 30] + await stop(loop) + + +@pytest.mark.asyncio +async def test_a_wake_just_before_the_wait_registers_runs_one_follow_up(): + cycles, timers, wake = Cycles(), Timers(), asyncio.Event() + calls = 0 + + def interval() -> int: + nonlocal calls + calls += 1 + if calls == 2: + wake.set() + return 30 + + loop, _client, _log = start_loop(cycles, timers, wake, interval=interval) + await settle() + + assert cycles.runs == 2 + assert timers.waiting == 1 + await stop(loop) + + +@pytest.mark.asyncio +async def test_a_burst_of_wakes_coalesces_without_overlap_or_queued_tasks(): + cycles, timers, wake = Cycles(), Timers(), asyncio.Event() + loop, _client, _log = start_loop(cycles, timers, wake) + await settle() + baseline = len(other_tasks()) + + for _ in range(100): + wake.set() + await settle() + + assert cycles.runs == 2 + assert cycles.max_active == 1 + assert len(other_tasks()) == baseline + await stop(loop) + + +@pytest.mark.asyncio +async def test_a_failed_cycle_keeps_both_the_wake_and_the_backstop(): + cycles, timers, wake = Cycles(fail_first=True), Timers(), asyncio.Event() + loop, _client, log = start_loop(cycles, timers, wake) + await settle() + log.exception.assert_called_once_with("failed") + + wake.set() + await settle() + assert cycles.runs == 2 + + timers.fire() + await settle() + assert cycles.runs == 3 + await stop(loop) + + +@pytest.mark.asyncio +async def test_cancelling_a_waiting_loop_closes_redis_once_and_leaves_no_helper_task(): + cycles, timers, wake = Cycles(), Timers(), asyncio.Event() + loop, client, _log = start_loop(cycles, timers, wake) + await settle() + + loop.cancel() + with pytest.raises(asyncio.CancelledError): + await loop + + client.close.assert_awaited_once() + assert timers.cancelled == 1 + assert other_tasks() == set() + + +@pytest.mark.asyncio +async def test_a_published_story_completion_hint_wakes_its_loop(monkeypatch): + broker, wake = Broker(), asyncio.Event() + listener = start_listener(broker, Timers(), {STORY_COMPLETION: wake}, monkeypatch) + await broker.subscribed() + + assert await publish(broker) == 1 + await settle() + + assert wake.is_set() + await stop(listener) + + +@pytest.mark.asyncio +async def test_rejected_payloads_neither_wake_nor_stop_the_listener(monkeypatch): + from src.tasks import scheduler_wake_listener + + log = MagicMock() + monkeypatch.setattr(scheduler_wake_listener, "logger", log) + broker, wake = Broker(), asyncio.Event() + listener = start_listener(broker, Timers(), {STORY_COMPLETION: wake}, monkeypatch) + await broker.subscribed() + + for payload in ( + '{"target": "pr_ci"}', + '{"target": "story_completion", "story_id": "story-1"}', + "not json", + b'{"target": "story_completion"}', + ): + broker.deliver(SCHEDULER_WAKE_CHANNEL, payload) + await settle() + + assert not wake.is_set() + rejected = [ + c for c in log.warning.call_args_list if c.args == ("scheduler_wake_hint_rejected",) + ] + assert len(rejected) == 4 + broker.deliver(SCHEDULER_WAKE_CHANNEL, '{"target": "story_completion"}') + await settle() + assert wake.is_set() + await stop(listener) + + +@pytest.mark.asyncio +async def test_a_listener_outage_reconnects_after_the_interval_while_the_backstop_runs( + monkeypatch, +): + broker, wake = Broker(), asyncio.Event() + broker.connect_failures = 1 + listener_timers, cycles, loop_timers = Timers(), Cycles(), Timers() + listener = start_listener(broker, listener_timers, {STORY_COMPLETION: wake}, monkeypatch) + loop, _client, _log = start_loop(cycles, loop_timers, wake) + await settle() + + assert cycles.runs == 1 + assert listener_timers.requested == [30] + assert broker.clients[0].closed == 1 + loop_timers.fire() + await settle() + assert cycles.runs == 2 + assert loop_timers.requested == [30, 30] + + listener_timers.fire() + await broker.subscribed() + await publish(broker) + await settle() + assert cycles.runs == 3 + await stop(listener, loop) + + +@pytest.mark.asyncio +async def test_a_dropped_subscription_closes_its_resources_and_resubscribes(monkeypatch): + broker, wake = Broker(), asyncio.Event() + timers = Timers() + listener = start_listener(broker, timers, {STORY_COMPLETION: wake}, monkeypatch) + first = await broker.subscribed() + + broker.drop_connections() + await settle() + + assert first.closed == 1 + assert broker.clients[0].closed == 1 + assert timers.requested == [30] + timers.fire() + await settle() + second = await broker.subscribed() + assert second is not first + await publish(broker) + await settle() + assert wake.is_set() + await stop(listener) + + +@pytest.mark.asyncio +async def test_cancelling_the_listener_closes_pubsub_and_client_once(monkeypatch): + broker = Broker() + listener = start_listener(broker, Timers(), {STORY_COMPLETION: asyncio.Event()}, monkeypatch) + pubsub = await broker.subscribed() + + listener.cancel() + with pytest.raises(asyncio.CancelledError): + await listener + + assert pubsub.closed == 1 + assert [client.closed for client in broker.clients] == [1] + assert other_tasks() == set() + + +@pytest.mark.asyncio +async def test_the_listener_refuses_a_target_without_a_receiver(): + from src.tasks.scheduler_wake_listener import scheduler_wake_listener + + with pytest.raises(ValueError, match="story_completion"): + await scheduler_wake_listener({}) diff --git a/services/scheduler/tests/unit/test_service_entrypoints.py b/services/scheduler/tests/unit/test_service_entrypoints.py index 266d4db94..355a56c84 100644 --- a/services/scheduler/tests/unit/test_service_entrypoints.py +++ b/services/scheduler/tests/unit/test_service_entrypoints.py @@ -315,19 +315,73 @@ async def test_pipeline_starts_without_loki(monkeypatch): initialize.assert_awaited_once_with( pipeline.PIPELINE_REQUIRED_KEYS, service_name="scheduler-pipeline" ) - workers = run_workers.await_args.args[0] - assert [(name, worker) for name, worker in workers] == [ - ("task_dispatcher", pipeline.task_dispatcher_loop), - ("scaffold", pipeline.scaffold_loop), - ("story_completion", pipeline.story_completion_loop), - ("lifecycle_supervision", pipeline.lifecycle_supervision_loop), - ("qa_routing", pipeline.qa_routing_loop), - ("pr_ci", pipeline.pr_ci_loop), - ("worker_reconciliation", pipeline.worker_reconciliation_loop), - ("temporary_access", pipeline.temporary_access_loop), - ("owner_notifications", pipeline.owner_notification_loop), - ("story_supervision", pipeline.story_supervision_loop), + workers = dict(run_workers.await_args.args[0]) + assert list(workers) == [ + "task_dispatcher", + "scaffold", + "story_completion", + "scheduler_wake", + "lifecycle_supervision", + "qa_routing", + "pr_ci", + "worker_reconciliation", + "temporary_access", + "owner_notifications", + "story_supervision", ] + assert workers["task_dispatcher"] is pipeline.task_dispatcher_loop + assert workers["scaffold"] is pipeline.scaffold_loop + assert workers["lifecycle_supervision"] is pipeline.lifecycle_supervision_loop + assert workers["qa_routing"] is pipeline.qa_routing_loop + assert workers["pr_ci"] is pipeline.pr_ci_loop + assert workers["worker_reconciliation"] is pipeline.worker_reconciliation_loop + assert workers["temporary_access"] is pipeline.temporary_access_loop + assert workers["owner_notifications"] is pipeline.owner_notification_loop + assert workers["story_supervision"] is pipeline.story_supervision_loop + + +@pytest.mark.asyncio +async def test_pipeline_routes_a_story_completion_hint_to_its_loop(monkeypatch): + from _wake_doubles import Broker, Timers, settle + + from shared.contracts.queues.scheduler_wake import SchedulerWakeTarget + from shared.redis.scheduler_wake import publish_scheduler_wake + from src import pipeline + from src.tasks import scheduler_wake_listener, story_completion_loop + + monkeypatch.setenv("INTERNAL_API_KEY", "test-internal-key") + monkeypatch.setenv("API_BASE_URL", "http://127.0.0.1:9") + broker = Broker() + complete = AsyncMock(return_value=0) + run_workers = AsyncMock() + monkeypatch.setattr(pipeline.runtime, "initialize_configs", AsyncMock()) + monkeypatch.setattr(pipeline.runtime, "run_workers", run_workers) + monkeypatch.setattr(story_completion_loop, "RedisStreamClient", broker.client) + monkeypatch.setattr(story_completion_loop, "complete_stories", complete) + monkeypatch.setattr(story_completion_loop, "_story_completion_interval", lambda: 3600) + monkeypatch.setattr(scheduler_wake_listener, "_reconnect_interval", lambda: 3600) + await pipeline.main() + workers = dict(run_workers.await_args.args[0]) + + tasks = [ + asyncio.create_task(workers["story_completion"]()), + asyncio.create_task( + workers["scheduler_wake"](redis_factory=broker.client, sleep=Timers().sleep) + ), + ] + await broker.subscribed() + await settle() + assert complete.await_count == 1 + + publisher = broker.client() + await publisher.connect() + await publish_scheduler_wake(publisher, SchedulerWakeTarget.STORY_COMPLETION) + await settle() + + assert complete.await_count == 2 + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) @pytest.mark.asyncio @@ -364,6 +418,7 @@ async def capture(module, service_name): "temporary_access", "owner_notifications", "story_supervision", + "scheduler_wake", "server_sync", "health_checker", "provisioner_results", diff --git a/shared/contracts/queues/scheduler_wake.py b/shared/contracts/queues/scheduler_wake.py new file mode 100644 index 000000000..3e7ad2649 --- /dev/null +++ b/shared/contracts/queues/scheduler_wake.py @@ -0,0 +1,49 @@ +"""Nonpersistent scheduler wake hints carried over native Redis Pub/Sub. + +A hint only names the loop that should reread its durable state soon. Delivery is lossy and +coalescible, and a hint grants no permission to transition a Story, Task or Run. +""" + +from enum import StrEnum + +from pydantic import BaseModel, ConfigDict, ValidationError + +from shared.diagnostics import safe_validation_errors + +SCHEDULER_WAKE_CHANNEL = "scheduler:wake" +# Bounds the work and the diagnostics one hostile message can cost the listener. +MAX_WAKE_HINT_CHARS = 256 + + +class SchedulerWakeTarget(StrEnum): + """Scheduler loops that accept a wake hint.""" + + STORY_COMPLETION = "story_completion" + + +class SchedulerWakeHint(BaseModel): + """A routing-only request that one scheduler loop runs its next cycle early.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + target: SchedulerWakeTarget + + +class SchedulerWakeHintRejected(ValueError): + """A Pub/Sub payload that is not a supported wake hint; `diagnostics` never echoes it.""" + + def __init__(self, reason: str, **diagnostics: object) -> None: + super().__init__(reason) + self.diagnostics: dict[str, object] = {"reason": reason, **diagnostics} + + +def parse_scheduler_wake_hint(data: object) -> SchedulerWakeHint: + """Validate one Pub/Sub payload as a wake hint or raise `SchedulerWakeHintRejected`.""" + if not isinstance(data, str): + raise SchedulerWakeHintRejected("not_text", payload_type=type(data).__name__) + if len(data) > MAX_WAKE_HINT_CHARS: + raise SchedulerWakeHintRejected("oversized", chars=len(data)) + try: + return SchedulerWakeHint.model_validate_json(data) + except ValidationError as exc: + raise SchedulerWakeHintRejected("invalid", errors=safe_validation_errors(exc)) from None diff --git a/shared/redis/scheduler_wake.py b/shared/redis/scheduler_wake.py new file mode 100644 index 000000000..61c8696a2 --- /dev/null +++ b/shared/redis/scheduler_wake.py @@ -0,0 +1,16 @@ +"""Publish scheduler wake hints on the native Redis Pub/Sub channel.""" + +from shared.contracts.queues.scheduler_wake import ( + SCHEDULER_WAKE_CHANNEL, + SchedulerWakeHint, + SchedulerWakeTarget, +) +from shared.redis.client import RedisStreamClient + + +async def publish_scheduler_wake( + redis_client: RedisStreamClient, target: SchedulerWakeTarget +) -> int: + """Publish one lossy hint after the durable write it announces; returns the receiver count.""" + hint = SchedulerWakeHint(target=target) + return await redis_client.redis.publish(SCHEDULER_WAKE_CHANNEL, hint.model_dump_json()) diff --git a/shared/tests/test_scheduler_wake_contract.py b/shared/tests/test_scheduler_wake_contract.py new file mode 100644 index 000000000..67b08d381 --- /dev/null +++ b/shared/tests/test_scheduler_wake_contract.py @@ -0,0 +1,50 @@ +import json + +import pytest + +from shared.contracts.queues.scheduler_wake import ( + MAX_WAKE_HINT_CHARS, + SchedulerWakeHint, + SchedulerWakeHintRejected, + SchedulerWakeTarget, + parse_scheduler_wake_hint, +) + + +def test_a_published_hint_parses_back_to_its_target(): + wire = SchedulerWakeHint(target=SchedulerWakeTarget.STORY_COMPLETION).model_dump_json() + + assert json.loads(wire) == {"target": "story_completion"} + assert parse_scheduler_wake_hint(wire).target is SchedulerWakeTarget.STORY_COMPLETION + + +@pytest.mark.parametrize( + ("payload", "reason"), + [ + ('{"target": "pr_ci"}', "invalid"), + ('{"target": "STORY_COMPLETION"}', "invalid"), + ('{"target": 0}', "invalid"), + ("{}", "invalid"), + ('{"target": "story_completion", "story_id": "story-1"}', "invalid"), + ("not json", "invalid"), + (b'{"target": "story_completion"}', "not_text"), + (1, "not_text"), + ('{"target": "story_completion"' + " " * MAX_WAKE_HINT_CHARS + "}", "oversized"), + ], +) +def test_a_payload_outside_the_contract_is_rejected_not_rerouted(payload, reason): + with pytest.raises(SchedulerWakeHintRejected) as rejected: + parse_scheduler_wake_hint(payload) + + assert rejected.value.diagnostics["reason"] == reason + + +def test_rejection_diagnostics_never_echo_the_payload(): + secret = "ghp_" + "x" * 36 + payload = json.dumps({"target": secret}) + + with pytest.raises(SchedulerWakeHintRejected) as rejected: + parse_scheduler_wake_hint(payload) + + assert secret not in repr(rejected.value.diagnostics) + assert secret not in str(rejected.value)