feat(worker): Add metrics for tasks lost during shutdown

clients/python/src/taskbroker_client/worker/worker.py

@@ -1033,6 +1033,21 @@ def send_results(self, results: list[ProcessingResult], is_draining: bool = Fals
             if not is_draining:
                 for result in results:
                     self.put_result(result)
+            else:
+                # Draining drops the batch instead of requeueing it, so the broker will
+                # not hear about these tasks until their processing deadline lapses.
+                self._metrics.incr(
+                    "taskworker.worker.shutdown.results_dropped",
+                    value=len(results),
+                    tags={"processing_pool": self._processing_pool_name},
+                )
+                logger.warning(
+                    "taskworker.worker.shutdown.results_dropped",
+                    extra={
+                        "results": [result.task_id for result in results],
+                        "processing_pool": self._processing_pool_name,
+                    },
+                )
 
     def start_metrics_thread(self) -> None:
         """
@@ -1318,41 +1333,147 @@ def is_worker_full(self) -> bool:
     def put_result(self, result: ProcessingResult) -> None:
         self._processed_tasks.put(result)
 
+    def _record_shutdown_stage(self, stage: str, started_at: float) -> None:
+        """
+        Emit one stage's duration as it finishes, so a pod SIGKILLed part way
+        through shutdown still reports how far it got.
+        """
+        self._metrics.distribution(
+            "taskworker.worker.shutdown.stage_duration",
+            time.monotonic() - started_at,
+            tags={"processing_pool": self._processing_pool_name, "stage": stage},
+        )
+
+    def _drain_child_tasks(self) -> None:
+        """
+        Empty the child tasks queue, counting what shutdown discards. Nothing consumes
+        it once the children are gone, so these tasks were being dropped silently.
+        """
+        discarded = 0
+        while True:
+            try:
+                inflight = self._child_tasks.get_nowait()
+            except queue.Empty:
+                # Undercounts: `Empty` also means items are unflushed in the pipe.
+                break
+
+            discarded += 1
+            logger.info(
+                "taskworker.worker.shutdown.child_task_discarded",
+                extra={
+                    "task_id": inflight.activation.id,
+                    "namespace": inflight.activation.namespace,
+                    "taskname": inflight.activation.taskname,
+                    "processing_pool": self._processing_pool_name,
+                },
+            )
+
+        self._metrics.incr(
+            "taskworker.worker.shutdown.child_tasks_discarded",
+            value=discarded,
+            tags={"processing_pool": self._processing_pool_name},
+        )
+
     def shutdown(self) -> None:
         """
         Shutdown cleanly
         Activate the shutdown event and drain results before terminating children.
         """
         logger.info("taskworker.worker.shutdown.start")
+        shutdown_start = time.monotonic()
         self._shutdown_event.set()
 
         logger.info("taskworker.worker.shutdown.spawn_children")
         if self._spawn_children_thread:
             self._spawn_children_thread.join()
+        self._record_shutdown_stage("spawn_children", shutdown_start)
 
         logger.info("taskworker.worker.shutdown.children")
+        children_start = time.monotonic()
         with self._children_lock:
             children = [tracked_child.process for tracked_child in self._children.values()]
 
         for child in children:
             child.terminate()
+        killed = 0
         for child in children:
             child.join(WORKER_CHILD_JOIN_TIMEOUT_SEC)
             if child.is_alive():
                 child.kill()
                 child.join()
+                killed += 1
+
+        # A child that needed SIGKILL was still running; if it held a task, that task
+        # produces no result and waits out its processing deadline on the broker.
+        self._metrics.incr(
+            "taskworker.worker.shutdown.children_killed",
+            value=killed,
+            tags={"processing_pool": self._processing_pool_name},
+        )
+        if killed:
+            logger.warning(
+                "taskworker.worker.shutdown.children_killed",
+                extra={
+                    "killed": killed,
+                    "children": len(children),
+                    "processing_pool": self._processing_pool_name,
+                },
+            )
+        self._record_shutdown_stage("children", children_start)
 
         logger.info("taskworker.worker.shutdown.result")
+        result_thread_start = time.monotonic()
         if self._result_thread:
             # Use a timeout as sometimes this thread can deadlock on the Event.
             self._result_thread.join(timeout=5)
 
+        # A result thread that did not join is abandoned at process exit, along with
+        # any status updates its executor still had in flight.
+        self._metrics.incr(
+            "taskworker.worker.shutdown.result_thread",
+            tags={
+                "processing_pool": self._processing_pool_name,
+                "outcome": (
+                    "timeout"
+                    if self._result_thread is not None and self._result_thread.is_alive()
+                    else "joined"
+                ),
+            },
+        )
+        self._record_shutdown_stage("result_thread", result_thread_start)
+
         # Drain any remaining results synchronously
+        drain_start = time.monotonic()
+        drained = 0
         while True:
             try:
                 result = self._processed_tasks.get_nowait()
-                self.send_results([result], True)
             except queue.Empty:
                 break
 
-        logger.info("taskworker.worker.shutdown.complete")
+            drained += 1
+            self.send_results([result], True)
+
+        self._metrics.incr(
+            "taskworker.worker.shutdown.results_drained",
+            value=drained,
+            tags={"processing_pool": self._processing_pool_name},
+        )
+        self._record_shutdown_stage("drain_results", drain_start)
+
+        child_tasks_start = time.monotonic()
+        self._drain_child_tasks()
+        self._record_shutdown_stage("drain_child_tasks", child_tasks_start)
+
+        self._metrics.distribution(
+            "taskworker.worker.shutdown.duration",
+            time.monotonic() - shutdown_start,
+            tags={"processing_pool": self._processing_pool_name},
+        )
+        logger.info(
+            "taskworker.worker.shutdown.complete",
+            extra={
+                "duration": time.monotonic() - shutdown_start,
+                "processing_pool": self._processing_pool_name,
+            },
+        )

clients/python/tests/worker/test_worker.py

@@ -45,6 +45,7 @@
 from taskbroker_client.types import InflightTaskActivation, ProcessingResult
 from taskbroker_client.worker.worker import (
     PushTaskWorker,
+    RequeueException,
     ShutdownSignal,
     TaskWorker,
     TaskWorkerProcessingPool,
@@ -2837,3 +2838,201 @@ def recording_sleep(seconds: float) -> None:
     # frequency for every iteration.
     assert idle_sleeps, "future-checking thread never slept while idle"
     assert all(seconds == configured_frequency for seconds in idle_sleeps)
+
+
+def _incr_calls(metrics: mock.Mock, name: str) -> list[Any]:
+    return [c for c in metrics.incr.call_args_list if c.args[0] == name]
+
+
+def _distribution_calls(metrics: mock.Mock, name: str) -> list[Any]:
+    return [c for c in metrics.distribution.call_args_list if c.args[0] == name]
+
+
+def _stage_names(metrics: mock.Mock) -> list[str]:
+    return [
+        c.kwargs["tags"]["stage"]
+        for c in _distribution_calls(metrics, "taskworker.worker.shutdown.stage_duration")
+    ]
+
+
+class _SurvivingProcess(_FakeProcess):
+    """A child that ignores SIGTERM and only dies when killed."""
+
+    def terminate(self) -> None:
+        self.terminated = True
+
+    def join(self, timeout: float | None = None) -> None:
+        self.join_calls.append(timeout)
+
+
+def test_shutdown_counts_child_tasks_it_discards() -> None:
+    fake_context = _FakeContext()
+    pool = _make_fake_context_pool(fake_context)
+    pool._metrics = mock.Mock()
+
+    # Tasks the broker already handed us that no child ever picked up.
+    pool._child_tasks.put(SIMPLE_TASK)
+    pool._child_tasks.put(RETRY_TASK)
+
+    pool.shutdown()
+
+    calls = _incr_calls(pool._metrics, "taskworker.worker.shutdown.child_tasks_discarded")
+    assert len(calls) == 1
+    assert calls[0].kwargs["value"] == 2
+    assert calls[0].kwargs["tags"] == {"processing_pool": "test"}
+    # The queue is emptied as a side effect of counting it.
+    assert pool._child_tasks.empty()
+
+
+def test_shutdown_reports_zero_when_no_child_tasks_are_lost() -> None:
+    fake_context = _FakeContext()
+    pool = _make_fake_context_pool(fake_context)
+    pool._metrics = mock.Mock()
+
+    pool.shutdown()
+
+    # An explicit zero, so dashboards can tell "nothing lost" from "no data".
+    calls = _incr_calls(pool._metrics, "taskworker.worker.shutdown.child_tasks_discarded")
+    assert len(calls) == 1
+    assert calls[0].kwargs["value"] == 0
+
+
+def test_shutdown_counts_children_it_had_to_kill() -> None:
+    fake_context = _FakeContext()
+    pool = _make_fake_context_pool(fake_context)
+    pool._metrics = mock.Mock()
+
+    stubborn = _SurvivingProcess(name="stubborn", target=lambda: None, args=())
+    stubborn.start()
+    compliant = _FakeProcess(name="compliant", target=lambda: None, args=())
+    compliant.start()
+
+    with pool._children_lock:
+        pool._children[uuid4()] = TrackedChild(
+            process=stubborn, state="running", release=threading.Event()  # type: ignore[arg-type]
+        )
+        pool._children[uuid4()] = TrackedChild(
+            process=compliant, state="running", release=threading.Event()  # type: ignore[arg-type]
+        )
+
+    pool.shutdown()
+
+    assert stubborn.killed is True
+    assert compliant.killed is False
+
+    calls = _incr_calls(pool._metrics, "taskworker.worker.shutdown.children_killed")
+    assert len(calls) == 1
+    assert calls[0].kwargs["value"] == 1
+
+
+def test_shutdown_counts_results_it_drains() -> None:
+    capture = _SendResultCapture()
+    fake_context = _FakeContext()
+    pool = _make_fake_context_pool(fake_context)
+    pool._send_result_fn = capture
+    pool._metrics = mock.Mock()
+
+    for task_id in ("one", "two", "three"):
+        pool._processed_tasks.put(
+            ProcessingResult(
+                task_id=task_id,
+                status=TASK_ACTIVATION_STATUS_COMPLETE,
+                host="localhost:50051",
+                receive_timestamp=0,
+            )
+        )
+
+    pool.shutdown()
+
+    calls = _incr_calls(pool._metrics, "taskworker.worker.shutdown.results_drained")
+    assert len(calls) == 1
+    assert calls[0].kwargs["value"] == 3
+    # Every drained result is sent with is_draining=True so no new work is fetched.
+    assert len(capture.send_calls) == 3
+    assert all(is_draining for _, is_draining in capture.send_calls)
+
+
+def test_send_results_counts_results_dropped_while_draining() -> None:
+    fake_context = _FakeContext()
+    pool = _make_fake_context_pool(fake_context)
+    pool._metrics = mock.Mock()
+
+    def explode(results: list[ProcessingResult], is_draining: bool) -> None:
+        raise RequeueException("broker is down")
+
+    pool._send_result_fn = explode
+
+    result = ProcessingResult(
+        task_id="lost",
+        status=TASK_ACTIVATION_STATUS_COMPLETE,
+        host="localhost:50051",
+        receive_timestamp=0,
+    )
+    pool.send_results([result], is_draining=True)
+
+    calls = _incr_calls(pool._metrics, "taskworker.worker.shutdown.results_dropped")
+    assert len(calls) == 1
+    assert calls[0].kwargs["value"] == 1
+    # Draining must not requeue, or shutdown would spin on the same failing batch.
+    assert pool._processed_tasks.empty()
+
+
+def test_send_results_requeues_rather_than_counting_a_drop_when_not_draining() -> None:
+    fake_context = _FakeContext()
+    pool = _make_fake_context_pool(fake_context)
+    pool._metrics = mock.Mock()
+
+    def explode(results: list[ProcessingResult], is_draining: bool) -> None:
+        raise RequeueException("broker is down")
+
+    pool._send_result_fn = explode
+
+    result = ProcessingResult(
+        task_id="retried",
+        status=TASK_ACTIVATION_STATUS_COMPLETE,
+        host="localhost:50051",
+        receive_timestamp=0,
+    )
+    pool.send_results([result], is_draining=False)
+
+    assert _incr_calls(pool._metrics, "taskworker.worker.shutdown.results_dropped") == []
+    assert pool._processed_tasks.get_nowait().task_id == "retried"
+
+
+def test_shutdown_records_a_duration_for_every_stage() -> None:
+    fake_context = _FakeContext()
+    pool = _make_fake_context_pool(fake_context)
+    pool._metrics = mock.Mock()
+
+    pool.shutdown()
+
+    # Stages are emitted in order as each completes, so a pod that is killed part
+    # way through still reports how far it got.
+    assert _stage_names(pool._metrics) == [
+        "spawn_children",
+        "children",
+        "result_thread",
+        "drain_results",
+        "drain_child_tasks",
+    ]
+    assert len(_distribution_calls(pool._metrics, "taskworker.worker.shutdown.duration")) == 1
+
+
+def test_shutdown_reports_whether_the_result_thread_joined() -> None:
+    fake_context = _FakeContext()
+    pool = _make_fake_context_pool(fake_context)
+    pool._metrics = mock.Mock()
+
+    stuck = threading.Event()
+    pool._result_thread = threading.Thread(target=stuck.wait, daemon=True)
+    pool._result_thread.start()
+
+    try:
+        with mock.patch.object(pool._result_thread, "join"):
+            pool.shutdown()
+
+        calls = _incr_calls(pool._metrics, "taskworker.worker.shutdown.result_thread")
+        assert len(calls) == 1
+        assert calls[0].kwargs["tags"]["outcome"] == "timeout"
+    finally:
+        stuck.set()