diff --git a/clients/python/src/taskbroker_client/constants.py b/clients/python/src/taskbroker_client/constants.py index dbbd2128..17930050 100644 --- a/clients/python/src/taskbroker_client/constants.py +++ b/clients/python/src/taskbroker_client/constants.py @@ -83,12 +83,6 @@ This can be mutated by application test harnesses to run tasks without Kafka. """ -WORKER_CHILD_JOIN_TIMEOUT_SEC = 5 -""" -How long the parent worker process should allow child processes -to drain pending produce futures on shutdown before sending SIGKILL. -""" - SHUTDOWN_POLL_INTERVAL_SEC = 0.5 """ How often blocking waits in the worker check whether a signal handler diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index cdd71f71..36282253 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -38,7 +38,6 @@ DEFAULT_WORKER_WARMUP_TIMEOUT_SEC, MAX_BACKOFF_SECONDS_WHEN_HOST_UNAVAILABLE, SHUTDOWN_POLL_INTERVAL_SEC, - WORKER_CHILD_JOIN_TIMEOUT_SEC, ) from taskbroker_client.metrics import MetricsBackend from taskbroker_client.types import InflightTaskActivation, ProcessingResult @@ -1506,10 +1505,22 @@ def _drain_child_tasks(self) -> None: tags={"processing_pool": self._processing_pool_name}, ) + def _drain_results(self) -> int: + drained = 0 + while True: + try: + result = self._processed_tasks.get_nowait() + except queue.Empty: + return drained + + drained += 1 + self.send_results([result], True) + def shutdown(self) -> None: """ Shutdown cleanly - Activate the shutdown event and drain results before terminating children. + Stop children fetching, discard unstarted tasks, then wait for running tasks to + finish while sending their results. """ logger.info("taskworker.worker.shutdown.start") shutdown_start = time.monotonic() @@ -1521,38 +1532,15 @@ def shutdown(self) -> None: self._spawn_children_thread.join() self._record_shutdown_stage("spawn_children", shutdown_start) - logger.info("taskworker.worker.shutdown.children") - children_start = time.monotonic() + child_tasks_start = time.monotonic() + self._drain_child_tasks() + self._record_shutdown_stage("drain_child_tasks", child_tasks_start) + 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() @@ -1575,18 +1563,19 @@ def shutdown(self) -> None: ) self._record_shutdown_stage("result_thread", result_thread_start) - # Drain any remaining results synchronously - drain_start = time.monotonic() + logger.info("taskworker.worker.shutdown.children") + children_start = time.monotonic() drained = 0 - while True: - try: - result = self._processed_tasks.get_nowait() - except queue.Empty: - break - - drained += 1 - self.send_results([result], True) + # Keep draining: a child blocks on a full result queue, and on flushing it at exit. + while any(child.is_alive() for child in children): + drained += self._drain_results() + time.sleep(SHUTDOWN_POLL_INTERVAL_SEC) + for child in children: + child.join() + self._record_shutdown_stage("children", children_start) + drain_start = time.monotonic() + drained += self._drain_results() self._metrics.incr( "taskworker.worker.shutdown.results_drained", value=drained, @@ -1594,10 +1583,6 @@ def shutdown(self) -> None: ) 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, diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 56079841..160d0ee6 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -3513,14 +3513,23 @@ def _stage_names(metrics: mock.Mock) -> list[str]: ] -class _SurvivingProcess(_FakeProcess): - """A child that ignores SIGTERM and only dies when killed.""" +class _SlowExitProcess(_FakeProcess): + """A child that keeps running its task for a few polls after SIGTERM.""" + + def __init__(self, *args: Any, polls_left: int, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.polls_left = polls_left + self.on_exit: Callable[[], None] = lambda: None def terminate(self) -> None: self.terminated = True - def join(self, timeout: float | None = None) -> None: - self.join_calls.append(timeout) + def is_alive(self) -> bool: + if self.alive and self.polls_left == 0: + self.on_exit() + self.alive = False + self.polls_left -= 1 + return self.alive def test_shutdown_counts_child_tasks_it_discards() -> None: @@ -3555,32 +3564,36 @@ def test_shutdown_reports_zero_when_no_child_tasks_are_lost() -> None: assert calls[0].kwargs["value"] == 0 -def test_shutdown_counts_children_it_had_to_kill() -> None: +def test_shutdown_waits_for_running_children_and_sends_their_results() -> None: + capture = _SendResultCapture() fake_context = _FakeContext() pool = _make_fake_context_pool(fake_context) + pool._send_result_fn = capture pool._metrics = mock.Mock() - stubborn = _SurvivingProcess(name="stubborn", target=lambda: None, args=()) - stubborn.start() - compliant = _FakeProcess(name="compliant", target=lambda: None, args=()) - compliant.start() + slow = _SlowExitProcess(name="slow", target=lambda: None, args=(), polls_left=2) + slow.start() + late = ProcessingResult( + task_id="late", + status=TASK_ACTIVATION_STATUS_COMPLETE, + host="localhost:50051", + receive_timestamp=0, + ) + slow.on_exit = lambda: pool._processed_tasks.put(late) with pool._children_lock: pool._children[uuid4()] = _make_tracked_child( - "running", process=stubborn, release=threading.Event() + "running", process=slow, release=threading.Event() ) - pool._children[uuid4()] = _make_tracked_child( - "running", process=compliant, release=threading.Event() - ) - - pool.shutdown() - assert stubborn.killed is True - assert compliant.killed is False + with mock.patch("taskbroker_client.worker.worker.SHUTDOWN_POLL_INTERVAL_SEC", 0): + pool.shutdown() - calls = _incr_calls(pool._metrics, "taskworker.worker.shutdown.children_killed") - assert len(calls) == 1 - assert calls[0].kwargs["value"] == 1 + assert slow.terminated is True + assert slow.killed is False + assert slow.join_calls == [None] + assert [r.task_id for results, _ in capture.send_calls for r in results] == ["late"] + assert all(is_draining for _, is_draining in capture.send_calls) def test_shutdown_counts_results_it_drains() -> None: @@ -3668,10 +3681,10 @@ def test_shutdown_records_a_duration_for_every_stage() -> None: # way through still reports how far it got. assert _stage_names(pool._metrics) == [ "spawn_children", - "children", + "drain_child_tasks", "result_thread", + "children", "drain_results", - "drain_child_tasks", ] assert len(_distribution_calls(pool._metrics, "taskworker.worker.shutdown.duration")) == 1