diff --git a/AGENTS.md b/AGENTS.md index aa12e04..f62f14c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -37,6 +37,7 @@ Run a single test by node ID, e.g. `uv run pytest tests/test_command.py::TestCom ```bash curl -sSL https://raw.githubusercontent.com/codingjoe/naming-things/refs/heads/main/README.md ``` +- Use Google-style docstrings for mkdocs. NEVER use RST directives (`:param:`, `:returns:`, `:meth:`). - NEVER format code, tests, or docs manually. MUST USE `prek run --all-files`. ## Coverage diff --git a/tests/backends/test_base.py b/tests/backends/test_base.py index b72deda..950e08a 100644 --- a/tests/backends/test_base.py +++ b/tests/backends/test_base.py @@ -81,10 +81,10 @@ def test_peek__raise_not_implemented_error(self) -> None: ) ) - def test_telemetry__raise_not_implemented_error(self) -> None: - """Raise NotImplementedError for backend telemetry API.""" + async def test_queue_stats__raise_not_implemented_error(self) -> None: + """Raise NotImplementedError for backend queue_stats API.""" with pytest.raises(NotImplementedError): - BackendDouble(alias="default", params={}).telemetry() + await BackendDouble(alias="default", params={}).queue_stats() def test_dequeue__raise_not_implemented_error(self) -> None: """Raise NotImplementedError for backend dequeue API.""" diff --git a/tests/backends/test_redis.py b/tests/backends/test_redis.py index aba9206..76c2fbf 100644 --- a/tests/backends/test_redis.py +++ b/tests/backends/test_redis.py @@ -61,54 +61,10 @@ def test_error_path__maintain_continues_after_exception(self, caplog): patch.object( broker, "_reap_running_queue", side_effect=RuntimeError("reaper") ), - patch.object( - default_task_backend, - "_trim_telemetry", - side_effect=RuntimeError("trim"), - ), ): broker.main() assert "Mover error for queue" in caplog.text assert "Running reaper error for queue" in caplog.text - assert "Telemetry trim error for queue" in caplog.text - - def test_main__trims_stale_telemetry(self): - """main() trims every telemetry series older than result_ttl.""" - backend = RedisTaskBackend( - "broker_trim_test", - { - "QUEUES": ["default"], - "REDIS_URL": "redis://localhost:6379/0", - "OPTIONS": { - "result_ttl": datetime.timedelta(seconds=60), - }, - }, - ) - try: - task_result = backend.enqueue(echo, args=[1]) - ingress_key = backend.INGRESS_KEY.format( - prefix=backend.key_prefix, queue_name="default" - ) - successful_key = backend._segment_key( - TaskResultStatus.SUCCESSFUL, "default" - ) - failed_key = backend._segment_key(TaskResultStatus.FAILED, "default") - old = (timezone.now() - datetime.timedelta(seconds=120)).timestamp() * 1000 - backend.client.zadd(ingress_key, {task_result.id: old}) - backend.client.zadd(successful_key, {task_result.id: old}) - backend.client.zadd(failed_key, {task_result.id: old}) - assert backend.client.zcard(ingress_key) == 1 - assert backend.client.zcard(successful_key) == 1 - assert backend.client.zcard(failed_key) == 1 - - broker = RedisBroker(backend) - broker.main() - - assert backend.client.zcard(ingress_key) == 0 - assert backend.client.zcard(successful_key) == 0 - assert backend.client.zcard(failed_key) == 0 - finally: - backend.close() class TestRedisTaskBackend: @@ -184,7 +140,7 @@ def test_acquire__sets_last_attempted_at(self): finally: backend.close() - def test_running_reaper__fails_expired_tasks(self): + async def test_running_reaper__fails_expired_tasks(self): """Running reaper creates FAILED results for tasks with expired lease.""" backend = RedisTaskBackend( "running_reaper_test", @@ -217,9 +173,10 @@ def test_running_reaper__fails_expired_tasks(self): assert len(result.errors) == 1 assert "AcknowledgementTimeout" in result.errors[0].exception_class_path - # Reaping an expired task must count as egress and failed - stats = backend.telemetry().queues["default"] - assert stats.rates.egress == 1 + # Reaping an expired task records a failed result. Live egress + # now arrives via pub/sub, so backend.queue_stats() rates are zero. + stats = (await backend.queue_stats()).queues["default"] + assert stats.rates.egress == 0 assert stats.counts.failed == 1 assert stats.counts.successful == 0 finally: @@ -267,8 +224,8 @@ def test_stale_acknowledge__is_noop(self): finally: backend.close() - def test_telemetry__empty_backend(self): - """Telemetry returns zero counts for an empty backend.""" + async def test_queue_stats__empty_backend(self): + """queue_stats returns zero counts for an empty backend.""" backend = RedisTaskBackend( "telemetry_empty_test", { @@ -280,13 +237,13 @@ def test_telemetry__empty_backend(self): }, ) try: - telemetry = backend.telemetry() + telemetry = await backend.queue_stats() assert telemetry == BackendTelemetry(queues={"default": _stats()}) finally: backend.close() - def test_telemetry__counts_tasks(self): - """Telemetry reflects windowed ingress/egress rates and per-status counts.""" + async def test_queue_stats__counts_tasks(self): + """queue_stats reports per-status counts; rates come from pub/sub, not polling.""" backend = RedisTaskBackend( "telemetry_counts_test", { @@ -325,18 +282,16 @@ def test_telemetry__counts_tasks(self): ) ) - telemetry = backend.telemetry() + telemetry = await backend.queue_stats() assert telemetry.queues["default"] == _stats( - ingress=2, - egress=2, successful=1, failed=1, ) finally: backend.close() - def test_telemetry__egress_equals_successful_plus_failed(self): - """Egress is derived as successful + failed over the display window.""" + async def test_queue_stats__counts_successful_and_failed(self): + """queue_stats counts successful and finished results; polling rates stay zero.""" backend = RedisTaskBackend( "telemetry_egress_test", { @@ -376,15 +331,16 @@ def test_telemetry__egress_equals_successful_plus_failed(self): ) ) - stats = backend.telemetry().queues["default"] - assert stats.rates.egress == stats.counts.successful + stats.counts.failed - assert stats.rates.egress == 3 + stats = (await backend.queue_stats()).queues["default"] + # Rates come from the pub/sub buffer, not from Redis polling. + assert stats.rates.ingress == 0 + assert stats.rates.egress == 0 assert stats.counts.successful == 2 assert stats.counts.failed == 1 finally: backend.close() - def test_telemetry__successful_failed_evicted_by_result_ttl(self): + async def test_queue_stats__successful_failed_evicted_by_result_ttl(self): """successful/failed segment counts drop when results age out of result_ttl.""" backend = RedisTaskBackend( "telemetry_eviction_test", @@ -431,16 +387,37 @@ def _ack(status: TaskResultStatus) -> str: assert backend.client.zcard(successful_key) == 1 assert backend.client.zcard(failed_key) == 1 - stats = backend.telemetry().queues["default"] + stats = (await backend.queue_stats()).queues["default"] assert stats.counts.successful == 1 assert stats.counts.failed == 1 finally: backend.close() - def test_telemetry__ingress_egress_age_out_of_window(self): - """Ingress and egress older than the display window are excluded from the rates.""" + @staticmethod + def _await_message(pubsub, expected: bytes, *, timeout: float = 2.0): + """Drain pubsub until a user message with the expected payload arrives.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + message = pubsub.get_message(timeout=0.1) + if message and message.get("type") == "message": + if message["data"] == expected: + return message + raise AssertionError(f"no pubsub message {expected!r} received") + + @staticmethod + def _drain_subscription(pubsub, *, timeout: float = 2.0): + """Block until the pubsub connection reports it is subscribed.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if pubsub.get_message(timeout=0.1): + if pubsub.subscribed: + return + raise AssertionError("pubsub never reached subscribed state") + + def test_enqueue__publishes_ingress_telemetry(self): + """enqueue() publishes an ingress event on the telemetry channel.""" backend = RedisTaskBackend( - "telemetry_window_test", + "telemetry_publish_ingress_test", { "QUEUES": ["default"], "REDIS_URL": "redis://localhost:6379/0", @@ -449,132 +426,86 @@ def test_telemetry__ingress_egress_age_out_of_window(self): }, }, ) + pubsub = backend.client.pubsub() try: - task_result = backend.enqueue(echo, args=[42]) - acquired = backend.acquire( - timeout=datetime.timedelta(seconds=1), worker="window-test" - ) - assert acquired is not None - backend.acknowledge( - dataclasses.replace( - acquired, - status=TaskResultStatus.SUCCESSFUL, - finished_at=timezone.now(), - ) - ) - - ingress_key = backend.INGRESS_KEY.format( - prefix=backend.key_prefix, queue_name="default" - ) - successful_results_key = backend._segment_key( - TaskResultStatus.SUCCESSFUL, "default" - ) - old = (timezone.now() - datetime.timedelta(seconds=120)).timestamp() * 1000 - backend.client.zadd(ingress_key, {task_result.id: old}) - backend.client.zadd(successful_results_key, {task_result.id: old}) - - stats = backend.telemetry().queues["default"] - # Older than the 60s display window: not counted as recent throughput. - assert stats.rates.ingress == 0 - assert stats.rates.egress == 0 - # Ingress is evicted by the read path (retention = result_ttl = 60s). - assert backend.client.zcard(ingress_key) == 0 - # The results segment is evicted by acknowledge/reaper, not by reads, - # so the aged result is still present in the segment count. - assert backend.client.zcard(successful_results_key) == 1 - assert stats.counts.successful == 1 + pubsub.subscribe(backend.telemetry_channel) + self._drain_subscription(pubsub) + backend.enqueue(echo, args=[1]) + message = self._await_message(pubsub, b"ingress:default") + assert message["channel"] == backend.telemetry_channel.encode() finally: + pubsub.unsubscribe(backend.telemetry_channel) + pubsub.close() backend.close() - def test_telemetry__interval_shrinks_window(self): - """A custom interval excludes events from the rate but never trims below result_ttl.""" + def test_acknowledge__publishes_egress_telemetry(self): + """acknowledge() publishes an egress event on the telemetry channel.""" backend = RedisTaskBackend( - "telemetry_interval_test", + "telemetry_publish_egress_test", { "QUEUES": ["default"], "REDIS_URL": "redis://localhost:6379/0", "OPTIONS": { + "lease_ttl": datetime.timedelta(hours=1), "result_ttl": datetime.timedelta(seconds=60), }, }, ) + pubsub = backend.client.pubsub() try: - task_result = backend.enqueue(echo, args=[42]) - ingress_key = backend.INGRESS_KEY.format( - prefix=backend.key_prefix, queue_name="default" + backend.enqueue(echo, args=[1]) + acquired = backend.acquire( + timeout=datetime.timedelta(seconds=1), worker="publish-test" ) - recent = (timezone.now() - datetime.timedelta(seconds=5)).timestamp() * 1000 - backend.client.zadd(ingress_key, {task_result.id: recent}) - - # The 30s-old event is inside the default 60s window but outside a 10s one. - stats = backend.telemetry(interval=datetime.timedelta(seconds=10)).queues[ - "default" - ] - assert stats.rates.ingress == 1 - - stale = (timezone.now() - datetime.timedelta(seconds=30)).timestamp() * 1000 - backend.client.zadd(ingress_key, {task_result.id: stale}) - stats = backend.telemetry(interval=datetime.timedelta(seconds=10)).queues[ - "default" - ] - # Outside the 10s count window, so not displayed... - assert stats.rates.ingress == 0 - # ...but still retained for result_ttl (60s): a narrower display - # window must not clobber the retention policy. - assert backend.client.zcard(ingress_key) == 1 + pubsub.subscribe(backend.telemetry_channel) + self._drain_subscription(pubsub) + backend.acknowledge( + dataclasses.replace( + acquired, + status=TaskResultStatus.SUCCESSFUL, + finished_at=timezone.now(), + ) + ) + message = self._await_message(pubsub, b"egress:default") + assert message["channel"] == backend.telemetry_channel.encode() finally: + pubsub.unsubscribe(backend.telemetry_channel) + pubsub.close() backend.close() - def test_telemetry__egress_window_follows_display_interval(self): - """Egress rates are windowed by the display interval, not by result_ttl.""" + def test_requeue__publishes_ingress_telemetry(self): + """requeue() publishes an ingress event for the re-queued task.""" backend = RedisTaskBackend( - "telemetry_egress_window_test", + "telemetry_publish_requeue_test", { "QUEUES": ["default"], "REDIS_URL": "redis://localhost:6379/0", "OPTIONS": { + "lease_ttl": datetime.timedelta(hours=1), "result_ttl": datetime.timedelta(seconds=60), }, }, ) + pubsub = backend.client.pubsub() try: backend.enqueue(echo, args=[1]) acquired = backend.acquire( - timeout=datetime.timedelta(seconds=1), worker="egress-window-test" + timeout=datetime.timedelta(seconds=1), worker="requeue-publish-test" ) assert acquired is not None - backend.acknowledge( - dataclasses.replace( - acquired, - status=TaskResultStatus.SUCCESSFUL, - finished_at=timezone.now(), - ) - ) - successful_key = backend._segment_key( - TaskResultStatus.SUCCESSFUL, "default" + failed = dataclasses.replace( + acquired, + status=TaskResultStatus.FAILED, + finished_at=timezone.now(), ) - - # Inside result_ttl (60s) but outside a 10s display window. - recent = ( - timezone.now() - datetime.timedelta(seconds=30) - ).timestamp() * 1000 - backend.client.zadd(successful_key, {acquired.id: recent}) - - # A narrow window excludes the 30s-old result from the egress rate... - narrow = backend.telemetry(interval=datetime.timedelta(seconds=10)).queues[ - "default" - ] - assert narrow.rates.egress == 0 - # ...but it is still a present segment member (ZCARD is time-independent). - assert narrow.counts.successful == 1 - - # A wide window (60s) includes the 30s-old result in the egress rate. - wide = backend.telemetry(interval=datetime.timedelta(seconds=60)).queues[ - "default" - ] - assert wide.rates.egress == 1 - assert wide.counts.successful == 1 + pubsub.subscribe(backend.telemetry_channel) + self._drain_subscription(pubsub) + backend.requeue(failed, timezone.now() + datetime.timedelta(seconds=10)) + message = self._await_message(pubsub, b"ingress:default") + assert message["channel"] == backend.telemetry_channel.encode() finally: + pubsub.unsubscribe(backend.telemetry_channel) + pubsub.close() backend.close() def _acknowledge(self, status: TaskResultStatus) -> str: diff --git a/tests/conftest.py b/tests/conftest.py index c6a49eb..698c782 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,7 +11,9 @@ def _flush_keys(backend) -> None: @pytest.fixture(autouse=True) def flush_default_backend(): - """Flush all threadmill keys before and after each test.""" + """Flush all threadmill keys and reset async client before and after each test.""" _flush_keys(default_task_backend) + default_task_backend._async_client = None yield _flush_keys(default_task_backend) + default_task_backend._async_client = None diff --git a/tests/test_inspector.py b/tests/test_inspector.py index 0b006d9..0c3b7cc 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -37,7 +37,7 @@ class FailingBackend(ThreadmillTaskBackend): - """Backend double whose peek and telemetry raise for error-path tests.""" + """Backend double whose peek and queue_stats raise for error-path tests.""" def enqueue(self, task, args, kwargs): raise NotImplementedError @@ -45,8 +45,8 @@ def enqueue(self, task, args, kwargs): def peek(self, *args, **kwargs): raise RuntimeError("peek failed") - def telemetry(self, *, interval=None): - raise RuntimeError("telemetry failed") + async def queue_stats(self, *, interval=None): + raise RuntimeError("queue_stats failed") class ErroringBackend(RedisTaskBackend): @@ -245,12 +245,12 @@ async def test_list_view_selected_sets_queue(self): task_list = app.query_one("#task-list", TaskList) assert task_list.queue_name == item.queue_name - async def test_telemetry_refresh_updates_and_prunes_queues(self): + async def test_queue_stats_refresh_updates_and_prunes_queues(self): """Telemetry refresh updates existing queue labels and drops stale queues.""" app = InspectorApp(backend=default_task_backend, auto_refresh=False) async with app.run_test() as pilot: await pilot.pause() - app._refresh_telemetry() + app.telemetry = await app.backend.queue_stats() await pilot.pause() queue_list = app.query_one("#queue-list", QueueList) assert set(queue_list._items) == set(default_task_backend.queues) @@ -258,7 +258,7 @@ async def test_telemetry_refresh_updates_and_prunes_queues(self): await pilot.pause() assert list(queue_list._items) == ["default"] - async def test_telemetry_counts_are_scoped_to_selected_queue(self): + async def test_queue_stats_counts_are_scoped_to_selected_queue(self): """Tab counts reflect the selected queue, not backend-wide totals.""" default_task_backend.enqueue(echo, args=[1]) stats = _stats(ready=1) @@ -457,17 +457,17 @@ async def test_successful_tab_lists_finished_task(self): assert table.row_count >= 1 async def test_auto_refresh_timer_armed(self): - """With auto-refresh on, a telemetry timer is armed on mount.""" + """With auto-refresh on, the sparkline timer is armed on mount.""" app = InspectorApp(backend=default_task_backend) async with app.run_test() as pilot: await pilot.pause() await pilot.pause() - assert app._telemetry_timer is not None - assert app._telemetry_timer.name == "telemetry-refresh" + assert app._sparkline_timer is not None + assert app._sparkline_timer.name == "sparkline-refresh" app = InspectorApp(backend=default_task_backend, auto_refresh=False) async with app.run_test() as pilot: await pilot.pause() - assert app._telemetry_timer is None + assert app._sparkline_timer is None async def test_initial_focus_on_queue_list(self): """The app opens with focus on the queue list, not the backend selector.""" @@ -478,7 +478,7 @@ async def test_initial_focus_on_queue_list(self): assert app.focused is not None assert app.focused.id == "queue-list" - async def test_telemetry_refresh_does_not_re_peek_task_list(self): + async def test_queue_stats_refresh_does_not_re_peek_task_list(self): """A telemetry refresh updates counts but leaves task rows stale until manual refresh.""" default_task_backend.enqueue(echo, args=[1]) default_task_backend.enqueue(echo, args=[2]) @@ -495,7 +495,7 @@ async def test_telemetry_refresh_does_not_re_peek_task_list(self): default_task_backend.acquire( timeout=datetime.timedelta(seconds=1), worker="stale-test" ) - app.telemetry = app.backend.telemetry() + app.telemetry = await app.backend.queue_stats() await pilot.pause() await pilot.pause() assert table.row_count == before diff --git a/threadmill/backends/base.py b/threadmill/backends/base.py index 313323a..fbab995 100644 --- a/threadmill/backends/base.py +++ b/threadmill/backends/base.py @@ -1,6 +1,7 @@ import collections.abc import dataclasses import datetime +import enum import json import threading import typing @@ -95,6 +96,21 @@ class BackendTelemetry: queues: dict[str, QueueStats] +class TelemetryDirection(enum.Enum): + """Direction of a telemetry event: a task entering or leaving a queue.""" + + INGRESS = "ingress" + EGRESS = "egress" + + +@dataclasses.dataclass(frozen=True, slots=True) +class TelemetryEvent: + """A single ingress/egress event published by a worker via pub/sub.""" + + direction: TelemetryDirection + queue_name: str + + class Broker(threading.Thread): """Backend maintenance thread launched by the task executor.""" @@ -242,14 +258,16 @@ def peek( """ raise NotImplementedError - def telemetry( + async def queue_stats( self, *, interval: datetime.timedelta = datetime.timedelta(seconds=60) ) -> BackendTelemetry: - """Return a snapshot of stats for all configured queues. + """Return per-queue task counts for all configured queues.""" + raise NotImplementedError - Args: - interval: The time window for rolling rates. - """ + async def worker_telemetry( + self, + ) -> collections.abc.AsyncIterator[TelemetryEvent]: + """Yield ingress/egress telemetry events from the backend's pub/sub stream.""" raise NotImplementedError def dequeue(self, task_result: TaskResult) -> None: diff --git a/threadmill/backends/lua/acknowledge.lua b/threadmill/backends/lua/acknowledge.lua index 9395947..896f02f 100644 --- a/threadmill/backends/lua/acknowledge.lua +++ b/threadmill/backends/lua/acknowledge.lua @@ -3,7 +3,8 @@ -- per-status results history set. Evicts results whose finish score falls -- outside the retention window (result_ttl). The per-status history sets are -- the time series the inspector counts over for telemetry, so no separate --- egress window or status counters are needed here. +-- egress window or status counters are needed here. A pub/sub event is +-- published so the inspector can track live egress throughput. -- -- KEYS[1] -- running set (ZSET) -- KEYS[2] -- result key (STRING, stores serialized TaskResult) @@ -15,6 +16,8 @@ -- ARGV[3] -- result TTL in seconds -- ARGV[4] -- finish timestamp in milliseconds (score for the history set) -- ARGV[5] -- status (SUCCESSFUL or FAILED) +-- ARGV[6] -- telemetry pub/sub channel +-- ARGV[7] -- queue name -- Returns: 1 on success, 0 if task was not in the running set local removed = redis.call('ZREM', KEYS[1], ARGV[1]) @@ -31,4 +34,5 @@ if ARGV[5] ~= 'SUCCESSFUL' then end redis.call('ZADD', results_key, finish, ARGV[1]) redis.call('ZREMRANGEBYSCORE', results_key, 0, cutoff) +redis.call('PUBLISH', ARGV[6], 'egress:' .. ARGV[7]) return 1 diff --git a/threadmill/backends/redis.py b/threadmill/backends/redis.py index 8029e30..f67989b 100644 --- a/threadmill/backends/redis.py +++ b/threadmill/backends/redis.py @@ -1,5 +1,6 @@ """Redis-backed durable priority queue backend for Django's task framework.""" +import collections.abc import dataclasses import datetime import logging @@ -10,6 +11,7 @@ from pathlib import Path import redis +import redis.asyncio from django.tasks import DEFAULT_TASK_QUEUE_NAME, TaskResult, TaskResultStatus from django.tasks.exceptions import TaskResultDoesNotExist from django.tasks.signals import task_enqueued @@ -21,6 +23,8 @@ QueueCounts, QueueRates, QueueStats, + TelemetryDirection, + TelemetryEvent, ThreadmillTaskBackend, ) @@ -88,7 +92,6 @@ def _reap_running_queue(self, queue_name: str) -> None: def main(self) -> None: """Run mover and running reaper passes for all queues.""" - now_ms = time.time() * 1000 for queue_name in self.backend.queues: try: self._move_queue(queue_name) @@ -100,11 +103,6 @@ def main(self) -> None: except Exception: # noqa: BLE001 logger.exception("Running reaper error for queue %r", queue_name) - try: - self.backend._trim_telemetry(queue_name, now_ms=now_ms) - except Exception: # noqa: BLE001 - logger.exception("Telemetry trim error for queue %r", queue_name) - class RedisTaskBackend(ThreadmillTaskBackend): """Redis-backed durable priority queue backend. @@ -126,7 +124,7 @@ class RedisTaskBackend(ThreadmillTaskBackend): SEGMENT_KEY = "{prefix}:{queue_name}:{status}" DEFERRED_KEY = "{prefix}:{queue_name}:deferred" - INGRESS_KEY = "{prefix}:{queue_name}:ingress:events" + TELEMETRY_CHANNEL = "{prefix}:telemetry" ACQUIRE_SCRIPT = _load_lua("acquire") """Pop the next task from a priority queue and move it directly to the running set.""" @@ -150,13 +148,23 @@ def __init__(self, alias: str, params: dict) -> None: f"REDIS_URL must be specified in your settings for the {type(self).__name__}." ) from e self.client = redis.from_url(redis_url) + self._async_client: redis.asyncio.Redis | None = None + self.redis_url = redis_url self.key_prefix = f"threadmill:{{{alias}}}" + self.telemetry_channel = self.TELEMETRY_CHANNEL.format(prefix=self.key_prefix) self.lease_ttl = self.options.get("lease_ttl", datetime.timedelta(hours=1)) self.result_ttl = self.options.get("result_ttl", datetime.timedelta(days=1)) self.batch_size = self.options.get("batch_size", 100) self._acquire_script = self.client.register_script(self.ACQUIRE_SCRIPT) self._acknowledge_script = self.client.register_script(self.ACKNOWLEDGE_SCRIPT) + @property + def async_client(self) -> redis.asyncio.Redis: + """Lazily-created async Redis client, reused across calls.""" + if self._async_client is None: + self._async_client = redis.asyncio.Redis.from_url(self.redis_url) + return self._async_client + def _compute_score(self, priority: int, enqueued_at: datetime.datetime) -> float: """Compute a ZSET score for priority-ordered FIFO queueing. @@ -196,15 +204,11 @@ def enqueue( ) score = self._compute_score(task.priority, enqueued_at) - enqueued_at_ms = enqueued_at.timestamp() * 1000 serialized = self.serialize_task_result(task_result) task_key = self.TASK_KEY.format(prefix=self.key_prefix, task_id=task_result.id) task_data_ttl = int( self.lease_ttl.total_seconds() * 3 + self.result_ttl.total_seconds() ) - ingress_key = self.INGRESS_KEY.format( - prefix=self.key_prefix, queue_name=task.queue_name - ) pipe = self.client.pipeline() pipe.hset( @@ -216,7 +220,7 @@ def enqueue( }, ) pipe.expire(task_key, task_data_ttl) - pipe.zadd(ingress_key, {task_result.id: enqueued_at_ms}) + pipe.publish(self.telemetry_channel, f"ingress:{task.queue_name}") if task.run_after is not None: deferred_key = self.DEFERRED_KEY.format( @@ -310,6 +314,8 @@ def acknowledge(self, task_result: TaskResult) -> None: str(int(self.result_ttl.total_seconds())), str(finish_score), task_result.status.name, + self.telemetry_channel, + task_result.task.queue_name, ], ) @@ -347,6 +353,7 @@ def requeue(self, task_result: TaskResult, run_after: datetime.datetime) -> None pipe.hset(task_key, mapping={"data": serialized, "score": str(score)}) pipe.expire(task_key, task_data_ttl) pipe.zadd(deferred_key, {task_result.id: run_after_ms}) + pipe.publish(self.telemetry_channel, f"ingress:{task_result.task.queue_name}") pipe.execute() def dequeue(self, task_result: TaskResult) -> None: @@ -416,84 +423,61 @@ def get_result(self, result_id: str) -> TaskResult: return self.deserialize_task_result(data) raise TaskResultDoesNotExist(f"Task result {result_id!r} does not exist.") - def _trim_telemetry(self, queue_name: str, *, now_ms: float) -> None: - cutoff = now_ms - self.result_ttl.total_seconds() * 1000 - pipe = self.client.pipeline() - pipe.zremrangebyscore( - self.INGRESS_KEY.format(prefix=self.key_prefix, queue_name=queue_name), - 0, - cutoff, - ) - pipe.zremrangebyscore( - self._segment_key(TaskResultStatus.SUCCESSFUL, queue_name), - 0, - cutoff, - ) - pipe.zremrangebyscore( - self._segment_key(TaskResultStatus.FAILED, queue_name), - 0, - cutoff, - ) - pipe.execute() - - def telemetry( + async def queue_stats( self, *, interval: datetime.timedelta = datetime.timedelta(seconds=60) ) -> BackendTelemetry: - now_ms = time.time() * 1000 - window_start_ms = now_ms - interval.total_seconds() * 1000 - exclusive_start = f"({window_start_ms}" - retention_start_ms = now_ms - self.result_ttl.total_seconds() * 1000 - pipe = self.client.pipeline() + client = self.async_client + pipe = client.pipeline() for queue_name in self.queues: pipe.zcard(self._segment_key(TaskResultStatus.READY, queue_name)) pipe.zcard(self._segment_key(TaskResultStatus.RUNNING, queue_name)) pipe.zcard( self.DEFERRED_KEY.format(prefix=self.key_prefix, queue_name=queue_name) ) - successful_results_key = self._segment_key( - TaskResultStatus.SUCCESSFUL, queue_name - ) - failed_results_key = self._segment_key(TaskResultStatus.FAILED, queue_name) - pipe.zcard(successful_results_key) - pipe.zcard(failed_results_key) - ingress_key = self.INGRESS_KEY.format( - prefix=self.key_prefix, queue_name=queue_name - ) - pipe.zcount(ingress_key, exclusive_start, now_ms) - pipe.zcount(successful_results_key, exclusive_start, now_ms) - pipe.zcount(failed_results_key, exclusive_start, now_ms) - pipe.zremrangebyscore(ingress_key, 0, retention_start_ms) - results = pipe.execute() + pipe.zcard(self._segment_key(TaskResultStatus.SUCCESSFUL, queue_name)) + pipe.zcard(self._segment_key(TaskResultStatus.FAILED, queue_name)) + results = await pipe.execute() + zero_rates = QueueRates(interval=interval, ingress=0, egress=0) queues: dict[str, QueueStats] = {} for index, queue_name in enumerate(self.queues): - base = index * 9 - ( - ready, - running, - deferred, - successful, - failed, - ingress, - successful_rate, - failed_rate, - _, - ) = (int(c or 0) for c in results[base : base + 9]) + base = index * 5 queues[queue_name] = QueueStats( counts=QueueCounts( - ready=ready, - running=running, - deferred=deferred, - successful=successful, - failed=failed, - ), - rates=QueueRates( - interval=interval, - ingress=ingress, - egress=successful_rate + failed_rate, + ready=int(results[base] or 0), + running=int(results[base + 1] or 0), + deferred=int(results[base + 2] or 0), + successful=int(results[base + 3] or 0), + failed=int(results[base + 4] or 0), ), + rates=zero_rates, ) return BackendTelemetry(queues=queues) + async def worker_telemetry( + self, + ) -> collections.abc.AsyncGenerator[TelemetryEvent]: + client = self.async_client + pubsub = client.pubsub() + await pubsub.subscribe(self.telemetry_channel) + pubsub.ignore_subscribe_messages = True + try: + async for message in pubsub.listen(): + if (data := message.get("data")) is not None: + payload = data.decode() if isinstance(data, bytes) else data + direction, _, queue_name = payload.partition(":") + try: + event = TelemetryEvent( + direction=TelemetryDirection(direction), + queue_name=queue_name, + ) + except ValueError: + continue + yield event + finally: + await pubsub.unsubscribe(self.telemetry_channel) + await pubsub.aclose() + def close(self) -> None: - """Close the Redis connection.""" + """Close the Redis connections.""" self.client.close() + self._async_client = None diff --git a/threadmill/inspector/app.py b/threadmill/inspector/app.py index 3b27ad9..4a7070a 100644 --- a/threadmill/inspector/app.py +++ b/threadmill/inspector/app.py @@ -1,8 +1,11 @@ """Textual app for the inspector TUI.""" +import asyncio +import dataclasses import datetime import logging import math +import time import typing from typing import Any @@ -13,6 +16,7 @@ TaskResultStatus, task_backends, ) +from textual import work from textual.app import App, ComposeResult from textual.binding import Binding from textual.containers import Horizontal, Vertical @@ -24,18 +28,27 @@ ListItem, ListView, Select, + Sparkline, Static, TabbedContent, TabPane, ) -from ..backends.base import BackendTelemetry, ThreadmillTaskBackend +from ..backends.base import ( + BackendTelemetry, + TelemetryDirection, + ThreadmillTaskBackend, +) from .screens import ConfirmScreen, PurgeScreen +from .telemetry import TelemetryBuffer logger = logging.getLogger(__name__) -TELEMETRY_INTERVAL_SECONDS = 2.0 -"""Seconds between automatic queue-stat refreshes; the task list stays manual.""" +TELEMETRY_INTERVAL = datetime.timedelta(seconds=2) +"""Between automatic queue-stat refreshes; the task list stays manual.""" + +SPARKLINE_INTERVAL = datetime.timedelta(seconds=1) +"""Between sparkline refreshes from the rolling telemetry buffer.""" TAB_STATUSES: list[tuple[str, TaskResultStatus | None]] = [ ("Running", TaskResultStatus.RUNNING), @@ -230,7 +243,7 @@ def switch_tab(self, tab_id: str) -> None: def action_refresh(self) -> None: """Refresh telemetry and re-fetch the task list.""" - self.app._refresh_telemetry() + self.app._poll_queue_stats() self.refresh_tasks() def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None: @@ -275,15 +288,14 @@ def action_requeue(self) -> None: ) def _do_requeue(self, confirmed: bool, task: TaskResult) -> None: - if not confirmed: - return - try: - self.backend.requeue(task, datetime.datetime.now(datetime.UTC)) - except Exception: # noqa: BLE001 - logger.exception("Requeue failed for task %s", task.id) - self.app.notify("Requeue failed.", severity="error") - return - self.app._after_action("Task requeued.") + if confirmed: + try: + self.backend.requeue(task, datetime.datetime.now(datetime.UTC)) + except Exception: # noqa: BLE001 + logger.exception("Requeue failed for task %s", task.id) + self.app.notify("Requeue failed.", severity="error") + else: + self.app._after_action("Task requeued.") def action_dequeue(self) -> None: """Remove the selected task from its current segment after confirmation.""" @@ -298,15 +310,14 @@ def action_dequeue(self) -> None: ) def _do_dequeue(self, confirmed: bool, task: TaskResult) -> None: - if not confirmed: - return - try: - self.backend.dequeue(task) - except Exception: # noqa: BLE001 - logger.exception("Delete failed for task %s", task.id) - self.app.notify("Delete failed.", severity="error") - return - self.app._after_action("Task deleted.") + if confirmed: + try: + self.backend.dequeue(task) + except Exception: # noqa: BLE001 + logger.exception("Delete failed for task %s", task.id) + self.app.notify("Delete failed.", severity="error") + else: + self.app._after_action("Task deleted.") def _refresh_data(self) -> None: """Fetch and display tasks for the current queue and active tab.""" @@ -413,24 +424,22 @@ def _notify_selection(self, queue_name: str) -> None: def action_purge(self) -> None: item = self.highlighted_child - if not isinstance(item, QueueItem): - return - queue_name = item.queue_name - self.app.push_screen( - PurgeScreen(queue_name), - lambda confirmed: self._do_purge(confirmed, queue_name), - ) + if isinstance(item, QueueItem): + queue_name = item.queue_name + self.app.push_screen( + PurgeScreen(queue_name), + lambda confirmed: self._do_purge(confirmed, queue_name), + ) def _do_purge(self, confirmed: bool, queue_name: str) -> None: - if not confirmed: - return - try: - self.app.backend.purge(queue_name) - except Exception: # noqa: BLE001 - logger.exception("Purge failed for %r", queue_name) - self.app.notify("Purge failed.", severity="error") - return - self.app._after_action(f"Purged queue {queue_name}.") + if confirmed: + try: + self.app.backend.purge(queue_name) + except Exception: # noqa: BLE001 + logger.exception("Purge failed for %r", queue_name) + self.app.notify("Purge failed.", severity="error") + else: + self.app._after_action(f"Purged queue {queue_name}.") class InspectorApp(App): @@ -459,8 +468,10 @@ def __init__( self._task_list: TaskList | None = None self._task_detail: TaskDetail | None = None self._options_static: Static | None = None + self._sparkline_timer = None self._telemetry_timer = None self._auto_refresh = auto_refresh + self.telemetry_buffer = TelemetryBuffer() self.set_reactive(InspectorApp.backend, backend) def compose(self) -> ComposeResult: @@ -481,6 +492,9 @@ def compose(self) -> ComposeResult: telemetry=InspectorApp.telemetry ) with Vertical(id="right-pane"): + with Horizontal(id="sparkline-row"): + yield Sparkline(id="ingress-sparkline") + yield Sparkline(id="egress-sparkline") yield TaskList(id="task-list").data_bind( backend=InspectorApp.backend, telemetry=InspectorApp.telemetry, @@ -496,13 +510,21 @@ def on_mount(self) -> None: self._task_detail = self.query_one("#task-detail", TaskDetail) self._options_static = self.query_one("#backend-options", Static) self._refresh_options() - self._refresh_telemetry() + self.query_one("#ingress-sparkline", Sparkline).border_title = "Ingress" + self.query_one("#egress-sparkline", Sparkline).border_title = "Egress" + self._poll_queue_stats() + self._listen_telemetry() if self._auto_refresh: self._telemetry_timer = self.set_interval( - TELEMETRY_INTERVAL_SECONDS, - self._refresh_telemetry, + TELEMETRY_INTERVAL.total_seconds(), + self._poll_queue_stats, name="telemetry-refresh", ) + self._sparkline_timer = self.set_interval( + SPARKLINE_INTERVAL.total_seconds(), + self._refresh_sparklines, + name="sparkline-refresh", + ) self._queue_list.focus() def action_quit(self) -> None: @@ -512,7 +534,7 @@ def action_quit(self) -> None: def _after_action(self, message: str) -> None: """Refresh telemetry and tasks after a mutating queue action.""" self.notify(message) - self._refresh_telemetry() + self._poll_queue_stats() self._task_list.refresh_tasks() def action_switch_tab(self, tab_id: str) -> None: @@ -521,7 +543,9 @@ def action_switch_tab(self, tab_id: str) -> None: def watch_backend(self) -> None: self._refresh_options() - self._refresh_telemetry() + self.telemetry_buffer = TelemetryBuffer() + self._poll_queue_stats() + self._listen_telemetry() def on_select_changed(self, event: Select.Changed) -> None: """Switch backend when the user selects a different alias.""" @@ -548,9 +572,53 @@ def _refresh_options(self) -> None: ] self._options_static.update(" ".join(parts) or "No options") - def _refresh_telemetry(self) -> None: - """Poll the backend for fresh telemetry.""" + @work(exclusive=True, group="queue-stats") + async def _poll_queue_stats(self) -> None: + """Fetch queue counts and merge live pub/sub rates into the telemetry.""" try: - self.telemetry = self.backend.telemetry() + telemetry = await self.backend.queue_stats() except Exception: # noqa: BLE001 - logger.exception("Failed to refresh telemetry") + logger.exception("Failed to refresh queue stats") + else: + for queue_name, stats in telemetry.queues.items(): + live_rates = self.telemetry_buffer.rates_for(queue_name) + telemetry.queues[queue_name] = dataclasses.replace( + stats, rates=live_rates + ) + self.telemetry = telemetry + + def _refresh_sparklines(self) -> None: + """Push the latest 60-second rolling series into the ingress/egress sparklines.""" + queue_name = self._task_list.queue_name + ingress = self.query_one("#ingress-sparkline", Sparkline) + egress = self.query_one("#egress-sparkline", Sparkline) + now = time.monotonic() + ingress.data = ( + self.telemetry_buffer.series( + queue_name, TelemetryDirection.INGRESS, now=now + ) + if queue_name + else [0.0] * 60 + ) + egress.data = ( + self.telemetry_buffer.series(queue_name, TelemetryDirection.EGRESS, now=now) + if queue_name + else [0.0] * 60 + ) + + @work(exclusive=True, group="telemetry") + async def _listen_telemetry(self) -> None: + """Feed the telemetry buffer from the backend's pub/sub stream.""" + try: + stream = self.backend.worker_telemetry() + if asyncio.iscoroutine(stream): + stream = await stream + except NotImplementedError: + pass + else: + if stream is not None: + try: + async for event in stream: + self.telemetry_buffer.record(event, now=time.monotonic()) + except asyncio.CancelledError: + raise diff --git a/threadmill/inspector/inspector.scss b/threadmill/inspector/inspector.scss index 067f694..330a806 100644 --- a/threadmill/inspector/inspector.scss +++ b/threadmill/inspector/inspector.scss @@ -42,6 +42,29 @@ Select { width: 42; } +#sparkline-row { + height: 5; + margin-bottom: 0; +} + +#ingress-sparkline, +#egress-sparkline { + width: 1fr; + height: 1fr; + padding: 0; + margin: 0; + border: round $accent 40%; + border-title-color: $primary; +} + +#ingress-sparkline > .sparkline--max-color { + color: $success; +} + +#egress-sparkline > .sparkline--max-color { + color: $error; +} + #backend-row { height: 3; } @@ -76,7 +99,7 @@ QueueList ListItem { } TaskList { - height: 60%; + height: 1fr; TabbedContent, ContentSwitcher, @@ -104,7 +127,7 @@ DataTable { } TaskDetail { - height: 40%; + height: 1fr; width: 100%; overflow: scroll; padding: 0 1; diff --git a/threadmill/inspector/telemetry.py b/threadmill/inspector/telemetry.py new file mode 100644 index 0000000..efe0ddb --- /dev/null +++ b/threadmill/inspector/telemetry.py @@ -0,0 +1,81 @@ +"""In-memory rolling telemetry buffer fed by backend pub/sub events.""" + +import collections +import datetime +import time + +from threadmill.backends.base import QueueRates, TelemetryDirection, TelemetryEvent + +WINDOW_SECONDS = 60 +"""Rolling window the inspector displays for ingress/egress sparklines.""" + + +class TelemetryBuffer: + """Per-(queue, direction) rolling 60-second count buffer.""" + + def __init__(self, *, window: int = WINDOW_SECONDS) -> None: + self._window = window + self._counts: dict[ + tuple[str, TelemetryDirection], collections.OrderedDict[int, int] + ] = collections.defaultdict(collections.OrderedDict) + + def record(self, event: TelemetryEvent, *, now: float | None = None) -> None: + """Increment the current-second bucket for the event's queue+direction.""" + second = int(now if now is not None else time.monotonic()) + buckets = self._counts[event.queue_name, event.direction] + buckets[second] = buckets.get(second, 0) + 1 + self._evict(buckets, second) + + def series( + self, + queue_name: str, + direction: TelemetryDirection, + *, + now: float | None = None, + ) -> list[float]: + """Return ``window`` per-second counts, oldest to newest, for the pair.""" + second = int(now if now is not None else time.monotonic()) + buckets = self._counts.get((queue_name, direction)) + if buckets is not None: + self._evict(buckets, second) + return [ + float(buckets.get(second - offset, 0)) if buckets else 0.0 + for offset in reversed(range(self._window)) + ] + + def rates_for( + self, + queue_name: str, + *, + interval: datetime.timedelta = datetime.timedelta(seconds=WINDOW_SECONDS), + now: float | None = None, + ) -> QueueRates: + """Sum the last ``interval`` of events into per-second ingress/egress rates.""" + second = int(now if now is not None else time.monotonic()) + span = min(int(interval.total_seconds()), self._window) + ingress = self._sum_direction( + queue_name, TelemetryDirection.INGRESS, second, span + ) + egress = self._sum_direction( + queue_name, TelemetryDirection.EGRESS, second, span + ) + return QueueRates( + interval=datetime.timedelta(seconds=span), ingress=ingress, egress=egress + ) + + def _sum_direction( + self, queue_name: str, direction: TelemetryDirection, second: int, span: int + ) -> int: + buckets = self._counts.get((queue_name, direction)) + if buckets is None: + return 0 + self._evict(buckets, second) + return sum(buckets.get(second - offset, 0) for offset in range(span)) + + def _evict( + self, buckets: collections.OrderedDict[int, int], now_second: int + ) -> None: + """Drop buckets older than the window (LRU pop from front).""" + cutoff = now_second - self._window + while buckets and next(iter(buckets)) <= cutoff: + buckets.popitem(last=False)