diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index f429f6a2..274a883a 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/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, + }, + ) diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 9045ab04..f232f31f 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/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()