Rework worker stats: configurable consumers, occupancy tracking

tasktiger/stats.py

@@ -1,71 +1,153 @@
 import threading
 import time
-from typing import TYPE_CHECKING, Optional
+from contextlib import contextmanager
+from typing import Any, Iterator, Optional
 
 from ._internal import g_fork_lock
 
-if TYPE_CHECKING:
-    from .worker import Worker
 
+class StatsConsumer:
+    """Receives worker measurements, configurable via STATS_CONSUMERS.
 
-class StatsThread(threading.Thread):
-    def __init__(self, tiger: "Worker") -> None:
+    Subclasses override the context managers below. Overrides must not
+    raise on enter and must not suppress exceptions on exit (both would
+    disrupt task processing). Instances are shared across all workers
+    created from one TaskTiger, so keep span state local to the generator
+    frame, not on the instance.
+
+    start()/stop() are called by the worker around its run loop.
+    """
+
+    def start(self) -> None:
+        pass
+
+    def stop(self) -> None:
+        pass
+
+    @contextmanager
+    def measure_task(self) -> Iterator[None]:
+        yield
+
+    @contextmanager
+    def measure_idle(self) -> Iterator[None]:
+        yield
+
+
+class StatsThread(threading.Thread, StatsConsumer):
+    """Periodically logs a "stats" event with worker time accounting.
+
+    Construct with a logger and an interval in seconds and add it to
+    STATS_CONSUMERS; the worker starts and stops it. Like any thread, an
+    instance can only be started once — use a fresh instance per worker
+    run when running multiple workers in one process.
+    """
+
+    def __init__(self, log: Any, interval: float) -> None:
         super(StatsThread, self).__init__()
-        self.tiger = tiger
+        self.log = log
+        self.interval = interval
         self._stop_event = threading.Event()
 
         self._task_running = False
         self._time_start = time.monotonic()
         self._time_busy: float = 0.0
         self._task_start_time: Optional[float] = None
+        self._time_idle: float = 0.0
+        self._idle_start_time: Optional[float] = None
         self.daemon = True  # Exit process if main thread exits unexpectedly
 
         # Lock that protects stats computations from interleaving. For example,
         # we don't want report_task_start() to run at the same time as
         # compute_stats(), as it might result in an inconsistent state.
+        # Timestamps are taken under the lock so spans can't cross a
+        # concurrent compute_stats() window boundary.
         self._computation_lock = threading.Lock()
 
     def report_task_start(self) -> None:
-        now = time.monotonic()
         with self._computation_lock:
-            self._task_start_time = now
+            self._task_start_time = time.monotonic()
             self._task_running = True
 
     def report_task_end(self) -> None:
-        now = time.monotonic()
         with self._computation_lock:
             assert self._task_start_time is not None
-            self._time_busy += now - self._task_start_time
+            self._time_busy += time.monotonic() - self._task_start_time
             self._task_running = False
             self._task_start_time = None
 
-    def compute_stats(self) -> None:
-        now = time.monotonic()
+    def report_idle_start(self) -> None:
+        with self._computation_lock:
+            assert self._idle_start_time is None
+            self._idle_start_time = time.monotonic()
+
+    def report_idle_end(self) -> None:
+        with self._computation_lock:
+            assert self._idle_start_time is not None
+            self._time_idle += time.monotonic() - self._idle_start_time
+            self._idle_start_time = None
+
+    @contextmanager
+    def measure_task(self) -> Iterator[None]:
+        self.report_task_start()
+        try:
+            yield
+        finally:
+            self.report_task_end()
 
+    @contextmanager
+    def measure_idle(self) -> Iterator[None]:
+        self.report_idle_start()
+        try:
+            yield
+        finally:
+            self.report_idle_end()
+
+    def compute_stats(self) -> None:
         with self._computation_lock:
+            now = time.monotonic()
             time_total = now - self._time_start
             time_busy = self._time_busy
+            time_idle = self._time_idle
             self._time_start = now
             self._time_busy = 0
+            self._time_idle = 0
             if self._task_running:
                 assert self._task_start_time is not None
                 time_busy += now - self._task_start_time
                 self._task_start_time = now
             else:
                 self._task_start_time = None
+            if self._idle_start_time is not None:
+                time_idle += now - self._idle_start_time
+                self._idle_start_time = now
 
         if time_total:
+            # time_busy: inside task code; time_idle: blocking waits for new
+            # work; time_overhead: everything else (dequeueing, queue
+            # scanning, locks, maintenance).
+            time_overhead = time_total - time_busy - time_idle
+            # A worker saturated with short tasks can show low utilization;
+            # for load/autoscaling decisions use occupancy: the busy share
+            # of busy + idle time, excluding overhead. Reads 0 with no
+            # tasks and ~100 when saturated, regardless of overhead.
             utilization = 100.0 / time_total * time_busy
+            time_attributable = time_busy + time_idle
+            occupancy = (
+                100.0 * time_busy / time_attributable if time_attributable else 0.0
+            )
             with g_fork_lock:
-                self.tiger.log.info(
+                self.log.info(
                     "stats",
                     time_total=time_total,
                     time_busy=time_busy,
+                    time_idle=time_idle,
+                    time_overhead=time_overhead,
                     utilization=utilization,
+                    occupancy=occupancy,
                 )
 
     def run(self) -> None:
-        while not self._stop_event.wait(self.tiger.config["STATS_INTERVAL"]):
+        while not self._stop_event.wait(self.interval):
             self.compute_stats()
 
     def stop(self) -> None:

tasktiger/tasktiger.py

@@ -191,8 +191,10 @@ def init(
             # subqueues will be automatically treated as batch queues, and the
             # batch value of the most specific subqueue name takes precedence.
             "BATCH_QUEUES": {},
-            # How often to print stats.
-            "STATS_INTERVAL": 60,
+            # List of StatsConsumer instances that receive worker task/idle
+            # measurements. Instances are shared across all workers created
+            # from this TaskTiger.
+            "STATS_CONSUMERS": [],
             # Single worker queues can reduce redis activity in some use cases
             # by locking at the queue level instead of just at the task or task
             # group level. These queues will only allow a single worker to

tasktiger/worker.py

@@ -8,10 +8,12 @@
 import time
 import uuid
 from collections import OrderedDict
+from contextlib import ExitStack, contextmanager
 from typing import (
     TYPE_CHECKING,
     Any,
     Collection,
+    Iterator,
     Dict,
     List,
     Literal,
@@ -41,7 +43,7 @@
 from .executor import Executor, ForkExecutor
 from .redis_semaphore import Semaphore
 from .runner import get_runner_class
-from .stats import StatsThread
+from .stats import StatsConsumer
 from .task import Task
 from .timeouts import JobTimeoutException
 from .utils import redis_glob_escape
@@ -80,7 +82,7 @@ def __init__(
         self._key = tiger._key
         self._did_work = True
         self._last_task_check = 0.0
-        self.stats_thread: Optional[StatsThread] = None
+        self.stats: List[StatsConsumer] = []
         self.id = str(uuid.uuid4())
 
         if executor_class is None:
@@ -214,6 +216,20 @@ def _worker_queue_scheduled_tasks(self) -> None:
                     self.connection.publish(self._key("activity"), queue)
                 self._did_work = True
 
+    @contextmanager
+    def _measure_idle(self) -> Iterator[None]:
+        with ExitStack() as stack:
+            for consumer in self.stats:
+                stack.enter_context(consumer.measure_idle())
+            yield
+
+    @contextmanager
+    def _measure_task(self) -> Iterator[None]:
+        with ExitStack() as stack:
+            for consumer in self.stats:
+                stack.enter_context(consumer.measure_task())
+            yield
+
     def _poll_for_queues(self) -> None:
         """
         Refresh list of queues.
@@ -223,7 +239,8 @@ def _poll_for_queues(self) -> None:
         This is only used when using polling to get queues with queued tasks.
         """
         if not self._did_work:
-            time.sleep(self.config["POLL_TASK_QUEUES_INTERVAL"])
+            with self._measure_idle():
+                time.sleep(self.config["POLL_TASK_QUEUES_INTERVAL"])
         self._refresh_queue_set()
 
     def _pubsub_for_queues(self, timeout: float = 0, batch_timeout: float = 0) -> None:
@@ -250,9 +267,10 @@ def _pubsub_for_queues(self, timeout: float = 0, batch_timeout: float = 0) -> No
                 pubsub_sleep = batch_exit - time.time()
             else:
                 pubsub_sleep = start_time + timeout - time.time()
-            message = self._pubsub.get_message(
-                timeout=0 if pubsub_sleep < 0 or self._did_work else pubsub_sleep
-            )
+            with self._measure_idle():
+                message = self._pubsub.get_message(
+                    timeout=0 if pubsub_sleep < 0 or self._did_work else pubsub_sleep
+                )
 
             # Pull remaining messages off of channel
             while message:
@@ -670,14 +688,10 @@ def _execute_task_group(
         if not ready_tasks:
             return True, []
 
-        if self.stats_thread:
-            self.stats_thread.report_task_start()
+        with self._measure_task():
+            self._prepare_execution(ready_tasks)
 
-        self._prepare_execution(ready_tasks)
-
-        success = self.executor.execute(queue, ready_tasks, log, locks, queue_lock)
-        if self.stats_thread:
-            self.stats_thread.report_task_end()
+            success = self.executor.execute(queue, ready_tasks, log, locks, queue_lock)
 
         for lock in locks:
             try:
@@ -974,11 +988,6 @@ def run(
             # executing pipelines.
             self.log.warning("using old Redis version")
 
-        if self.config["STATS_INTERVAL"]:
-            stats_thread = StatsThread(self)
-            self.stats_thread = stats_thread
-            stats_thread.start()
-
         # Queue any periodic tasks that are not queued yet.
         self._queue_periodic_tasks()
 
@@ -995,7 +1004,11 @@ def run(
 
         self._refresh_queue_set()
 
+        self.stats = list(self.config["STATS_CONSUMERS"])
         try:
+            for consumer in self.stats:
+                consumer.start()
+
             while True:
                 # Update the queue set on every iteration so we don't get stuck
                 # on processing a specific queue.
@@ -1028,9 +1041,9 @@ def run(
             raise
 
         finally:
-            if self.stats_thread:
-                self.stats_thread.stop()
-                self.stats_thread = None
+            for consumer in self.stats:
+                consumer.stop()
+            self.stats = []
 
             # Free up Redis connection
             if self._pubsub: