diff --git a/README.md b/README.md index 427854b..82ceeb6 100644 --- a/README.md +++ b/README.md @@ -145,12 +145,9 @@ uv add "threadmill[sentry]" | `threadmill.queue.depth` | gauge | `queue`, `status` (ready, running, deferred) | Tasks in a queue with status ready, running, or deferred | | `threadmill.queue.ingress` | counter | `queue` | Tasks that enter a queue | | `threadmill.queue.egress` | counter | `queue` | Tasks that finish an attempt | -| `threadmill.worker.count` | gauge | `pool` (process, thread) | Worker processes and threads in the pool | +| `threadmill.worker.count` | gauge | `pool` (process, thread) | Live worker processes and threads in the pool | | `threadmill.task.duration` | distribution | `task`, `outcome` (successful, failed), `attempt` | Duration of one task attempt in milliseconds | -A task duration is recorded in the worker that runs the attempt. -The SDK drops all metrics when it is not initialized. - ### Autoscaling Threadmill exposes queue stats as JSON for KEDA's `metrics-api` scaler. diff --git a/tests/test_executor.py b/tests/test_executor.py index d4a42d7..83936fd 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -422,6 +422,17 @@ def gauge(self, name, value, *, attributes=None) -> None: self.gauges.append((name, value, attributes)) +class StubWorker: + """Worker process double that reports liveness.""" + + def __init__(self, *, alive: bool) -> None: + self.alive = alive + + def is_alive(self) -> bool: + """Return the configured liveness.""" + return self.alive + + class StubShutdownEvent: """Report a shutdown request after the first loop iteration.""" @@ -779,8 +790,8 @@ def start_after_failed_attempt(process): assert executor.worker_processes[0].is_alive() executor.worker_processes[0].shutdown() - def test_report_metrics__samples_pool_and_queue_counts(self, monkeypatch): - """Emit pool and per-queue counts once per sample interval.""" + def test_report_queue_depth__reports_counts(self, monkeypatch): + """Emit per-queue counts once per sample interval.""" backend = StubQueueStatsBackend( BackendTelemetry( queues={ @@ -809,18 +820,16 @@ def test_report_metrics__samples_pool_and_queue_counts(self, monkeypatch): monkeypatch.setattr(executor_module, "sentry_metrics", sentry_metrics) executor.shutdown_requested = shutdown_requested - executor.report_metrics() + executor.report_queue_depth() assert sentry_metrics.gauges == [ - ("threadmill.worker.count", 2, {"pool": "process"}), - ("threadmill.worker.count", 6, {"pool": "thread"}), ("threadmill.queue.depth", 5, {"queue": "default", "status": "ready"}), ("threadmill.queue.depth", 2, {"queue": "default", "status": "running"}), ("threadmill.queue.depth", 1, {"queue": "default", "status": "deferred"}), ] assert shutdown_requested.timeouts == [7.0] - def test_report_metrics__zero_interval_disables_sampling(self, monkeypatch): + def test_report_queue_depth__zero_interval_disables_sampling(self, monkeypatch): """Skip sampling when the interval is zero.""" executor = TaskExecutor( backend=StubQueueStatsBackend(BackendTelemetry(queues={})), @@ -830,10 +839,50 @@ def test_report_metrics__zero_interval_disables_sampling(self, monkeypatch): sentry_metrics = StubSentryMetrics() monkeypatch.setattr(executor_module, "sentry_metrics", sentry_metrics) - executor.report_metrics() + executor.report_queue_depth() assert sentry_metrics.gauges == [] + def test_report_worker_counts__reports_live_processes(self, monkeypatch): + """Report only live worker processes and their threads.""" + executor = TaskExecutor( + backend=StubQueueStatsBackend(BackendTelemetry(queues={})), + workers=3, + threads=2, + queues=("default",), + ) + sentry_metrics = StubSentryMetrics() + monkeypatch.setattr(executor_module, "sentry_metrics", sentry_metrics) + executor.worker_processes = [ + StubWorker(alive=True), + StubWorker(alive=False), + StubWorker(alive=True), + ] + + executor.report_worker_counts() + + assert sentry_metrics.gauges == [ + ("threadmill.worker.count", 2, {"pool": "process"}), + ("threadmill.worker.count", 4, {"pool": "thread"}), + ] + + def test_shutdown__reports_zero_worker_counts(self, monkeypatch): + """Report zero workers when the pool shuts down.""" + executor = TaskExecutor( + backend=StubQueueStatsBackend(BackendTelemetry(queues={})), + queues=("default",), + ) + sentry_metrics = StubSentryMetrics() + monkeypatch.setattr(executor_module, "sentry_metrics", sentry_metrics) + executor.pool_maintainer = threading.current_thread() + + executor.shutdown() + + assert sentry_metrics.gauges == [ + ("threadmill.worker.count", 0, {"pool": "process"}), + ("threadmill.worker.count", 0, {"pool": "thread"}), + ] + class TestWorkerProcess: """Tests for the WorkerProcess class.""" diff --git a/threadmill/executor.py b/threadmill/executor.py index 53472a9..9673fb9 100644 --- a/threadmill/executor.py +++ b/threadmill/executor.py @@ -189,11 +189,14 @@ def run(self) -> None: self.pool_maintainer = threading.Thread( target=self.maintain_worker_pool, daemon=True ) + self.report_worker_counts() threads = [self.pool_maintainer] if sentry_metrics is not None: threads.append( threading.Thread( - target=self.report_metrics, name="threadmill-metrics", daemon=True + target=self.report_queue_depth, + name="threadmill-queue-depth", + daemon=True, ) ) if self.backend.broker_class: @@ -220,11 +223,14 @@ def shutdown(self) -> None: self.broker.shutdown() with ThreadPoolExecutor(max_workers=self.process_count) as executor: executor.map(lambda worker: worker.shutdown(), self.worker_processes) + if self.pool_maintainer is not None: + self.report_worker_counts() def maintain_worker_pool(self) -> None: """Restart worker processes that have exited, until shutdown is requested.""" while not self.shutdown_requested.is_set(): all_dead = True + replaced = False for index, worker in enumerate(self.worker_processes): if worker.is_alive(): all_dead = False @@ -239,25 +245,33 @@ def maintain_worker_pool(self) -> None: logger.warning( "Unable to start worker-%d. Retrying ...", index, exc_info=error ) + replaced = True + if replaced: + self.report_worker_counts() if all_dead and self.exit_empty: self.shutdown() return self.shutdown_requested.wait(1) - def report_metrics(self) -> None: - """Sample queue depth and worker counts into Sentry until shutdown.""" + def report_worker_counts(self) -> None: + """Record the live worker process and thread counts.""" + if sentry_metrics is not None: + process_count = sum(worker.is_alive() for worker in self.worker_processes) + sentry_metrics.gauge( + "threadmill.worker.count", + process_count, + attributes={"pool": "process"}, + ) + sentry_metrics.gauge( + "threadmill.worker.count", + process_count * self.thread_count, + attributes={"pool": "thread"}, + ) + + def report_queue_depth(self) -> None: + """Sample queue depth into Sentry until shutdown.""" if self.sample_interval > datetime.timedelta(0): while not self.shutdown_requested.is_set(): - sentry_metrics.gauge( - "threadmill.worker.count", - self.process_count, - attributes={"pool": "process"}, - ) - sentry_metrics.gauge( - "threadmill.worker.count", - self.process_count * self.thread_count, - attributes={"pool": "thread"}, - ) telemetry = asyncio.run(self.backend.queue_stats()) for queue_name, stats in telemetry.queues.items(): for status, depth in ( diff --git a/threadmill/management/commands/threadmill.py b/threadmill/management/commands/threadmill.py index fc6803b..ecb4e5d 100644 --- a/threadmill/management/commands/threadmill.py +++ b/threadmill/management/commands/threadmill.py @@ -113,8 +113,8 @@ def add_arguments(self, parser): type=self.parse_duration, default=datetime.timedelta(seconds=30), help=( - "Interval between queue depth and worker count samples in seconds." - " Defaults to 30. Set to 0 to disable sampling." + "Interval between queue depth samples in seconds." + " Defaults to 30. Set to 0 to disable queue depth sampling." ), ) parser.add_argument(