feat(taskworker): Wait for running tasks on shutdown instead of killing children

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

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,29 +1563,26 @@ 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,
             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,

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