From 6f892d309e5d3dcfae9d2bc60dc72c6cc1ce772a Mon Sep 17 00:00:00 2001 From: Markus Unterwaditzer Date: Tue, 25 Aug 2026 10:51:23 +0200 Subject: [PATCH 1/4] fix(python): Shut workers down without raising from signal handlers Reland of #764 (reverted in e9d1494) with the bug that forced the revert fixed. Signal handlers now only flip a bool instead of raising KeyboardInterrupt. Raising unwinds at an arbitrary bytecode and can leave locks held by the interrupted code in a broken state, which is where the `ValueError: semaphore or lock released too many times` crashes came from. Anything that takes a lock, including `Event.set()` and `server.stop()`, is also unsafe to call from a handler, so `ShutdownSignal.request()` does nothing but assign. The previous attempt exited on its own shortly after startup because the serve loop read the return value of `grpc.Server.wait_for_termination(timeout=...)` as "the server terminated". It actually returns True when the timeout elapsed, i.e. while the server is healthy, and False once it has terminated - the inverse of `Event.wait()`. A healthy worker therefore broke out of the loop after one poll interval and shut down cleanly with exit code 0. The serve loop no longer looks at that return value at all. It sleeps on `ShutdownSignal.wait()`, so our own flag is the only thing that can end it. ref STREAM-1649 --- .../python/src/taskbroker_client/constants.py | 7 + .../src/taskbroker_client/worker/worker.py | 134 ++++++-- clients/python/tests/worker/test_worker.py | 316 +++++++++++++++++- 3 files changed, 411 insertions(+), 46 deletions(-) diff --git a/clients/python/src/taskbroker_client/constants.py b/clients/python/src/taskbroker_client/constants.py index dd4b4234..4e700d54 100644 --- a/clients/python/src/taskbroker_client/constants.py +++ b/clients/python/src/taskbroker_client/constants.py @@ -81,6 +81,13 @@ 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 +asked for shutdown. Signal handlers cannot wake those waits directly, +so this is the worst case delay before a SIGTERM is noticed. +""" + TASK_PRODUCER_MAX_PENDING_FUTURES = 10_000 """ Maximum number of pending futures that can be in the TaskProducer module's diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 30f7c298..c5e3e59c 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -35,6 +35,7 @@ DEFAULT_WORKER_QUEUE_SIZE, 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 @@ -60,6 +61,60 @@ WORKER_SERVICE_NAME = "sentry_protos.taskbroker.v1.WorkerService" +class ShutdownSignal: + """ + Shutdown state that a signal handler is allowed to flip. + + Python runs signal handlers on the main thread in between bytecodes, which + makes plain attribute assignment safe but anything that takes a lock unsafe: + if the interrupted code already holds that lock, the handler blocks forever + on the thread that would have released it. `threading.Event.set()` and + `multiprocessing.Event.set()` both take a non-reentrant lock, so neither may + be called from a handler. `request()` therefore only assigns a bool. + + The event is here so that a shutdown noticed on the main thread can wake + sleeps on the result thread, and is only ever set from normal code. + """ + + def __init__(self) -> None: + self._requested = False + self._event = threading.Event() + + def request(self) -> None: + """ + Ask for shutdown. This is the only method a signal handler may call. + """ + self._requested = True + + def set(self) -> None: + """ + Ask for shutdown and wake anything sleeping in `wait()`. + + Takes a lock, so this must never be called from a signal handler. + """ + self._requested = True + self._event.set() + + def is_set(self) -> bool: + return self._requested + + def wait(self, timeout: float) -> bool: + """ + Sleep up to `timeout` seconds, returning True if shutdown was requested. + + A handler can only flip the bool, so this polls rather than relying on + the event alone. + """ + deadline = time.monotonic() + timeout + while True: + if self._requested: + return True + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + self._event.wait(min(remaining, SHUTDOWN_POLL_INTERVAL_SEC)) + + class WorkerPrometheusMetrics: """ Owns the Prometheus registry, server, and metrics we expose for scraping. @@ -268,7 +323,7 @@ def __init__( ) self._metrics = app.metrics self._concurrency = concurrency - self._grpc_sync_event = self._mp_context.Event() + self._shutdown_signal = ShutdownSignal() self._health_check_sec_per_touch = ( None if health_check_file_path is None else health_check_sec_per_touch ) @@ -331,8 +386,8 @@ def _send_results( "processing_pool": self._processing_pool_name, }, ) - # Use the shutdown_event as a sleep mechanism - self._grpc_sync_event.wait(self._setstatus_backoff_seconds) + # Use the shutdown signal as a sleep mechanism + self._shutdown_signal.wait(self._setstatus_backoff_seconds) try: self.client.update_tasks(results) @@ -426,8 +481,8 @@ def _await_children_warm(self) -> None: }, ) break - # Sleep and break early if shutdown was requested via shutdown(). - if self._grpc_sync_event.wait(0.25): + # Sleep and break early if shutdown was requested. + if self._shutdown_signal.wait(0.25): break self._metrics.distribution( @@ -453,16 +508,17 @@ def start(self) -> int: self.worker_pool.start_result_thread() self.worker_pool.start_spawn_children_thread() - # Convert signals into KeyboardInterrupt. - # Running shutdown() within the signal handler can lead to deadlocks - server: grpc.Server | None = None + server_started = False health_servicer: health.HealthServicer | None = None + # Record the request and let the loop below act on it. Raising from a + # handler unwinds at an arbitrary bytecode and can leave the locks held + # by the interrupted code in a broken state; calling anything that takes + # a lock (server.stop(), Event.set()) can deadlock against the code the + # handler interrupted. See ShutdownSignal. def signal_handler(*args: Any) -> None: - if server: - server.stop(grace=5) - raise KeyboardInterrupt() + self._shutdown_signal.request() signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) @@ -503,7 +559,13 @@ def signal_handler(*args: Any) -> None: health_servicer.set(WORKER_SERVICE_NAME, health_pb2.HealthCheckResponse.NOT_SERVING) server.add_insecure_port(f"[::]:{self._grpc_port}") + + # Don't accept connections we are about to drop. + if self._shutdown_signal.is_set(): + return 0 + server.start() + server_started = True # Hold NOT_SERVING until children are warm so the pod stays out of # the NEG/readiness set while its child processes are still loading. @@ -511,7 +573,7 @@ def signal_handler(*args: Any) -> None: # If shutdown was requested during warmup, don't advertise SERVING. # Bail to the finally below, which sets NOT_SERVING and tears everything down. - if self._grpc_sync_event.is_set(): + if self._shutdown_signal.is_set(): return 0 # Indicate that the server is ready @@ -521,10 +583,15 @@ def signal_handler(*args: Any) -> None: logger.info("taskworker.grpc_server.started", extra={"port": self._grpc_port}) self._start_health_check_thread() - try: - server.wait_for_termination() - except KeyboardInterrupt: - # Signals are converted to KeyboardInterrupt, swallow for exit code 0 + # Poll so a signal handler that only flipped a bool still gets us out. + # + # Deliberately not `server.wait_for_termination(timeout=...)`: that + # returns True when the *timeout* was reached, i.e. while the server + # is healthy, and False once it has terminated. That is the inverse + # of `Event.wait()`, and reading it as "has terminated" is what made + # a previous version of this patch exit half a second after startup. + # The shutdown signal is the only thing allowed to end this loop. + while not self._shutdown_signal.wait(SHUTDOWN_POLL_INTERVAL_SEC): pass finally: @@ -532,11 +599,10 @@ def signal_handler(*args: Any) -> None: health_servicer.set("", health_pb2.HealthCheckResponse.NOT_SERVING) health_servicer.set(WORKER_SERVICE_NAME, health_pb2.HealthCheckResponse.NOT_SERVING) - if server is not None: + if server is not None and server_started: server.stop(grace=5) - self._stop_health_check_thread() - self.worker_pool.shutdown() + self.shutdown() return 0 @@ -545,7 +611,7 @@ def shutdown(self) -> None: Shutdown the worker. """ self._stop_health_check_thread() - self._grpc_sync_event.set() + self._shutdown_signal.set() self.worker_pool.shutdown() @@ -624,7 +690,7 @@ def __init__( ) self._metrics = app.metrics - self._grpc_sync_event = self._mp_context.Event() + self._shutdown_signal = ShutdownSignal() self._gettask_backoff_seconds = 0 self._setstatus_backoff_seconds = 0 @@ -639,20 +705,24 @@ def start(self) -> int: self.worker_pool.start_result_thread() self.worker_pool.start_spawn_children_thread() - # Convert signals into KeyboardInterrupt. - # Running shutdown() within the signal handler can lead to deadlocks + # Record the request and let the loop below act on it. Raising from a + # handler unwinds at an arbitrary bytecode and can leave the locks held + # by the interrupted code in a broken state; calling anything that takes + # a lock (Event.set()) can deadlock against the code the handler + # interrupted. See ShutdownSignal. def signal_handler(*args: Any) -> None: - raise KeyboardInterrupt() + self._shutdown_signal.request() signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) try: - while True: + while not self._shutdown_signal.is_set(): self.run_once() - except KeyboardInterrupt: + finally: self.shutdown() - raise + + return 0 def run_once(self) -> None: """Access point for tests to run a single worker loop""" @@ -710,7 +780,9 @@ def _send_update_task( }, ) - self._grpc_sync_event.wait(self._setstatus_backoff_seconds) + if self._shutdown_signal.wait(self._setstatus_backoff_seconds): + # Don't claim a task we won't be around to run. + fetch_next = None try: next_task = self.client.update_task(result, fetch_next) @@ -734,7 +806,9 @@ def _send_update_task( raise RequeueException(f"Failed to update task: {e}") def fetch_task(self) -> InflightTaskActivation | None: - self._grpc_sync_event.wait(self._gettask_backoff_seconds) + if self._shutdown_signal.wait(self._gettask_backoff_seconds): + return None + try: activation = self.client.get_task(self._namespace) except grpc.RpcError as e: @@ -767,7 +841,7 @@ def shutdown(self) -> None: """ Shutdown the worker. """ - self._grpc_sync_event.set() + self._shutdown_signal.set() self.worker_pool.shutdown() diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 556e8023..974bb795 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -30,6 +30,7 @@ TASK_ACTIVATION_STATUS_COMPLETE, TASK_ACTIVATION_STATUS_FAILURE, TASK_ACTIVATION_STATUS_RETRY, + FetchNextTask, PushTaskRequest, PushTaskResponse, RetryState, @@ -44,6 +45,7 @@ from taskbroker_client.types import InflightTaskActivation, ProcessingResult from taskbroker_client.worker.worker import ( PushTaskWorker, + ShutdownSignal, TaskWorker, TaskWorkerProcessingPool, TrackedChild, @@ -427,6 +429,72 @@ def _make_fake_context_pool( ) +def test_shutdown_signal_wait_times_out() -> None: + shutdown_signal = ShutdownSignal() + + start = time.monotonic() + assert shutdown_signal.wait(0.2) is False + assert time.monotonic() - start >= 0.2 + + +def test_shutdown_signal_wait_returns_immediately_when_requested() -> None: + shutdown_signal = ShutdownSignal() + shutdown_signal.request() + + start = time.monotonic() + assert shutdown_signal.wait(30) is True + assert time.monotonic() - start < 1 + + +def test_shutdown_signal_wait_wakes_on_request_from_another_thread() -> None: + shutdown_signal = ShutdownSignal() + threading.Timer(0.1, shutdown_signal.request).start() + + start = time.monotonic() + assert shutdown_signal.wait(30) is True + assert time.monotonic() - start < 5 + + +def test_shutdown_signal_wait_wakes_on_set_from_another_thread() -> None: + shutdown_signal = ShutdownSignal() + threading.Timer(0.1, shutdown_signal.set).start() + + start = time.monotonic() + assert shutdown_signal.wait(30) is True + assert time.monotonic() - start < 1 + + +def test_shutdown_signal_request_from_signal_handler_during_wait() -> None: + """ + A handler that fires while wait() is sleeping must not hang. Event.set() + can deadlock here because wait() may already hold the event's lock. + """ + shutdown_signal = ShutdownSignal() + previous = signal.signal(signal.SIGALRM, lambda *args: shutdown_signal.request()) + try: + signal.setitimer(signal.ITIMER_REAL, 0.1) + start = time.monotonic() + assert shutdown_signal.wait(30) is True + assert time.monotonic() - start < 5 + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + signal.signal(signal.SIGALRM, previous) + + +def _capture_signal_handlers( + handlers: dict[int, Callable[..., None]], +) -> Callable[[int, Callable[..., None]], None]: + """ + Stand-in for signal.signal that records handlers instead of installing them, + so tests can deliver a signal without touching the pytest process. + """ + + def install_handler(signum: int, handler: Callable[..., None]) -> None: + handlers[signum] = handler + + return install_handler + + def _wait_for(condition: Callable[[], bool], timeout: float = 5) -> None: start = time.time() while time.time() - start < timeout: @@ -453,6 +521,68 @@ def test_fetch_task(self) -> None: assert task assert task.activation.id == SIMPLE_TASK.activation.id + def test_fetch_task_skips_request_during_shutdown(self) -> None: + taskworker = TaskWorker( + app_module="examples.app:app", + broker_hosts=["127.0.0.1:50051"], + max_child_task_count=100, + process_type="fork", + ) + taskworker._shutdown_signal.request() + + with mock.patch.object(taskworker.client, "get_task") as mock_get: + task = taskworker.fetch_task() + + assert task is None + mock_get.assert_not_called() + + def test_send_update_task_does_not_fetch_next_during_shutdown(self) -> None: + taskworker = TaskWorker( + app_module="examples.app:app", + broker_hosts=["127.0.0.1:50051"], + max_child_task_count=100, + process_type="fork", + ) + taskworker._shutdown_signal.request() + result = _make_processing_result("completed") + + with mock.patch.object(taskworker.client, "update_task", return_value=None) as mock_update: + taskworker._send_update_task(result, FetchNextTask(namespace="examples")) + + mock_update.assert_called_once_with(result, None) + + def test_start_exits_cleanly_on_sigterm(self) -> None: + taskworker = TaskWorker( + app_module="examples.app:app", + broker_hosts=["127.0.0.1:50051"], + max_child_task_count=100, + process_type="fork", + ) + handlers: dict[int, Callable[..., None]] = {} + + def deliver_sigterm() -> None: + handlers[signal.SIGTERM](signal.SIGTERM, None) + + with ( + mock.patch( + "taskbroker_client.worker.worker.signal.signal", + side_effect=_capture_signal_handlers(handlers), + ), + mock.patch.object(taskworker.worker_pool, "start_metrics_thread"), + mock.patch.object(taskworker.worker_pool, "start_result_thread"), + mock.patch.object(taskworker.worker_pool, "start_spawn_children_thread"), + mock.patch.object(taskworker.worker_pool, "shutdown") as pool_shutdown, + mock.patch.object(taskworker, "run_once", side_effect=deliver_sigterm) as run_once, + ): + exitcode = taskworker.start() + + # The handler returns instead of raising, so the loop finishes the + # iteration it was in and then exits. + assert exitcode == 0 + assert run_once.call_count == 1 + assert taskworker._shutdown_signal.is_set() + pool_shutdown.assert_called_once_with() + def test_fetch_no_task(self) -> None: taskworker = TaskWorker( app_module="examples.app:app", @@ -831,29 +961,181 @@ def warm_up() -> None: assert timeout_calls == [] -def test_start_does_not_serve_when_shutdown_during_warmup() -> None: - from grpc_health.v1 import health_pb2 +@contextlib.contextmanager +def _push_worker_grpc_mocks( + taskworker: PushTaskWorker, + fake_server: mock.MagicMock, + fake_health: mock.MagicMock, + handlers: dict[int, Callable[..., None]] | None = None, +) -> Iterator[mock.MagicMock]: + """ + Patch out everything PushTaskWorker.start() touches apart from its own + shutdown handling. Yields the mocked pool shutdown. + """ + with contextlib.ExitStack() as stack: + if handlers is not None: + stack.enter_context( + mock.patch( + "taskbroker_client.worker.worker.signal.signal", + side_effect=_capture_signal_handlers(handlers), + ) + ) + for name in ( + "start_metrics_thread", + "start_result_thread", + "start_spawn_children_thread", + ): + stack.enter_context(mock.patch.object(taskworker.worker_pool, name)) + pool_shutdown = stack.enter_context(mock.patch.object(taskworker.worker_pool, "shutdown")) + stack.enter_context(mock.patch.object(taskworker, "_start_health_check_thread")) + stack.enter_context(mock.patch.object(taskworker, "_stop_health_check_thread")) + stack.enter_context( + mock.patch("taskbroker_client.worker.worker.grpc.server", return_value=fake_server) + ) + stack.enter_context( + mock.patch( + "taskbroker_client.worker.worker.health.HealthServicer", return_value=fake_health + ) + ) + stack.enter_context( + mock.patch( + "taskbroker_client.worker.worker.health_pb2_grpc.add_HealthServicer_to_server" + ) + ) + stack.enter_context( + mock.patch( + "taskbroker_client.worker.worker.taskbroker_pb2_grpc" + ".add_WorkerServiceServicer_to_server" + ) + ) + yield pool_shutdown + +def test_push_start_exits_cleanly_on_sigterm() -> None: taskworker = _make_push_worker(concurrency=2, warmup_timeout=5) - # Children never warm, and shutdown is requested before start() runs. - taskworker._grpc_sync_event.set() + handlers: dict[int, Callable[..., None]] = {} fake_health = mock.MagicMock() fake_server = mock.MagicMock() + # Deliver the signal from inside the poll sleep, which is where a real + # SIGTERM would land once the server is up and serving. + real_wait = taskworker._shutdown_signal.wait + calls = 0 + + def wait(timeout: float) -> bool: + nonlocal calls + calls += 1 + if calls == 1: + handlers[signal.SIGTERM](signal.SIGTERM, None) + return real_wait(timeout) + with ( - mock.patch.object(taskworker.worker_pool, "start_metrics_thread"), - mock.patch.object(taskworker.worker_pool, "start_result_thread"), - mock.patch.object(taskworker.worker_pool, "start_spawn_children_thread"), - mock.patch.object(taskworker.worker_pool, "shutdown"), - mock.patch("taskbroker_client.worker.worker.grpc.server", return_value=fake_server), - mock.patch( - "taskbroker_client.worker.worker.health.HealthServicer", return_value=fake_health + _push_worker_grpc_mocks(taskworker, fake_server, fake_health, handlers) as pool_shutdown, + mock.patch.object( + TaskWorkerProcessingPool, "ready_count", new_callable=mock.PropertyMock, return_value=2 ), - mock.patch("taskbroker_client.worker.worker.health_pb2_grpc.add_HealthServicer_to_server"), - mock.patch( - "taskbroker_client.worker.worker.taskbroker_pb2_grpc" - ".add_WorkerServiceServicer_to_server" + mock.patch.object(taskworker._shutdown_signal, "wait", side_effect=wait), + ): + exitcode = taskworker.start() + + assert exitcode == 0 + assert taskworker._shutdown_signal.is_set() + # The handler only flipped a bool, so we polled our way out rather than + # relying on the server being stopped from inside the handler. + assert calls == 1 + fake_server.stop.assert_called_once_with(grace=5) + pool_shutdown.assert_called_once_with() + + +def test_push_start_keeps_serving_while_server_is_healthy() -> None: + """ + A healthy server must not end the serve loop. + + `grpc.Server.wait_for_termination(timeout=...)` returns True when the + timeout elapsed, i.e. while the server is still up, and False once it has + terminated -- the inverse of `Event.wait()`. A previous version of this + patch treated that True as "terminated" and so exited one poll interval + after startup, taking down every worker. Guard against reintroducing it by + making the mock behave the way grpc really does. + """ + taskworker = _make_push_worker(concurrency=2, warmup_timeout=5) + handlers: dict[int, Callable[..., None]] = {} + + fake_health = mock.MagicMock() + fake_server = mock.MagicMock() + # Healthy server: the timeout is always what elapses. + fake_server.wait_for_termination.return_value = True + + polls = 0 + + def wait(timeout: float) -> bool: + nonlocal polls + polls += 1 + # Survive several polls, then ask to stop so the test terminates. + if polls < 5: + return False + taskworker._shutdown_signal.request() + return True + + with ( + _push_worker_grpc_mocks(taskworker, fake_server, fake_health, handlers) as pool_shutdown, + mock.patch.object( + TaskWorkerProcessingPool, "ready_count", new_callable=mock.PropertyMock, return_value=2 + ), + mock.patch.object(taskworker._shutdown_signal, "wait", side_effect=wait), + ): + exitcode = taskworker.start() + + assert exitcode == 0 + # We kept looping instead of bailing out on the first poll. + assert polls == 5 + # The serve loop must not consult the server's return value at all. + fake_server.wait_for_termination.assert_not_called() + fake_server.stop.assert_called_once_with(grace=5) + pool_shutdown.assert_called_once_with() + + +def test_push_start_does_not_start_server_when_shutdown_requested_first() -> None: + taskworker = _make_push_worker(concurrency=2, warmup_timeout=5) + taskworker._shutdown_signal.request() + + fake_health = mock.MagicMock() + fake_server = mock.MagicMock() + + with _push_worker_grpc_mocks(taskworker, fake_server, fake_health) as pool_shutdown: + exitcode = taskworker.start() + + assert exitcode == 0 + fake_server.start.assert_not_called() + # Stopping a server that was never started is not something grpc promises + # to handle. + fake_server.stop.assert_not_called() + fake_server.wait_for_termination.assert_not_called() + pool_shutdown.assert_called_once_with() + + +def test_push_start_does_not_serve_when_shutdown_during_warmup() -> None: + from grpc_health.v1 import health_pb2 + + taskworker = _make_push_worker(concurrency=2, warmup_timeout=5) + handlers: dict[int, Callable[..., None]] = {} + + fake_health = mock.MagicMock() + fake_server = mock.MagicMock() + + # Children never warm; SIGTERM arrives while we are waiting on them. + def ready_count() -> int: + handlers[signal.SIGTERM](signal.SIGTERM, None) + return 0 + + with ( + _push_worker_grpc_mocks(taskworker, fake_server, fake_health, handlers) as pool_shutdown, + mock.patch.object( + TaskWorkerProcessingPool, + "ready_count", + new_callable=mock.PropertyMock, + side_effect=ready_count, ), ): exitcode = taskworker.start() @@ -866,8 +1148,10 @@ def test_start_does_not_serve_when_shutdown_during_warmup() -> None: if c.args[1] == health_pb2.HealthCheckResponse.SERVING ] assert serving_calls == [] - # We never reached server.wait_for_termination() (returned before it). + fake_server.start.assert_called_once() + fake_server.stop.assert_called_once_with(grace=5) fake_server.wait_for_termination.assert_not_called() + pool_shutdown.assert_called_once_with() def _make_tracked_child( From 91a3ef6d4f073fd7dd42fd79fe4f300c867c6f6d Mon Sep 17 00:00:00 2001 From: Markus Unterwaditzer Date: Tue, 25 Aug 2026 11:10:44 +0200 Subject: [PATCH 2/4] fix(python): Address review on shutdown handling - Notice a gRPC server that terminated on its own again. Polling only the shutdown flag meant an internally failed server left the parent running with live children and a green health check. Read wait_for_termination(timeout=...) with its real semantics: True means the timeout elapsed and the server is still up, False means it terminated. - Drop a pull-mode task claimed while shutting down. get_task() blocks with no deadline, so SIGTERM can land mid-RPC; handing the activation to a child claims work we won't run, which then has to expire on the broker before anyone else picks it up. --- .../src/taskbroker_client/worker/worker.py | 41 +++++++-- clients/python/tests/worker/test_worker.py | 87 +++++++++++++++---- 2 files changed, 101 insertions(+), 27 deletions(-) diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index c5e3e59c..fa026692 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -583,16 +583,20 @@ def signal_handler(*args: Any) -> None: logger.info("taskworker.grpc_server.started", extra={"port": self._grpc_port}) self._start_health_check_thread() - # Poll so a signal handler that only flipped a bool still gets us out. + # Poll so a signal handler that only flipped a bool still gets us + # out, while also noticing a server that terminated on its own. # - # Deliberately not `server.wait_for_termination(timeout=...)`: that - # returns True when the *timeout* was reached, i.e. while the server - # is healthy, and False once it has terminated. That is the inverse - # of `Event.wait()`, and reading it as "has terminated" is what made - # a previous version of this patch exit half a second after startup. - # The shutdown signal is the only thing allowed to end this loop. - while not self._shutdown_signal.wait(SHUTDOWN_POLL_INTERVAL_SEC): - pass + # Mind the return value of `wait_for_termination(timeout=...)`: it + # is True when the *timeout* elapsed, i.e. while the server is still + # healthy, and False once the server has terminated. That is the + # inverse of `Event.wait()`, and reading it as "has terminated" is + # what made a previous version of this patch exit half a second + # after startup and take down every worker. + while not self._shutdown_signal.is_set(): + still_running = server.wait_for_termination(timeout=SHUTDOWN_POLL_INTERVAL_SEC) + if not still_running: + logger.warning("taskworker.grpc_server.terminated_unexpectedly") + break finally: if health_servicer is not None: @@ -835,6 +839,25 @@ def fetch_task(self) -> InflightTaskActivation | None: return None self._gettask_backoff_seconds = 0 + + # get_task() blocks with no deadline, so a SIGTERM can arrive while it + # is in flight. Re-check before handing the activation to a child: + # claiming work we are not going to run means waiting for it to expire + # on the broker before anyone else picks it up. + if self._shutdown_signal.is_set(): + self._metrics.incr( + "taskworker.worker.fetch_task.dropped_during_shutdown", + tags={"processing_pool": self._processing_pool_name}, + ) + logger.info( + "taskworker.fetch_task.dropped_during_shutdown", + extra={ + "task_id": activation.activation.id, + "processing_pool": self._processing_pool_name, + }, + ) + return None + return activation def shutdown(self) -> None: diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 974bb795..96567077 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -536,6 +536,31 @@ def test_fetch_task_skips_request_during_shutdown(self) -> None: assert task is None mock_get.assert_not_called() + def test_fetch_task_drops_task_claimed_during_shutdown(self) -> None: + """ + get_task() blocks with no deadline, so SIGTERM can land mid-RPC. + + Handing the activation to a child anyway claims work we will not run, + which then has to expire on the broker before anyone else picks it up. + """ + taskworker = TaskWorker( + app_module="examples.app:app", + broker_hosts=["127.0.0.1:50051"], + max_child_task_count=100, + process_type="fork", + ) + + def get_task(namespace: str | None = None) -> InflightTaskActivation: + # Shutdown requested while the RPC was in flight. + taskworker._shutdown_signal.request() + return SIMPLE_TASK + + with mock.patch.object(taskworker.client, "get_task", side_effect=get_task) as mock_get: + task = taskworker.fetch_task() + + mock_get.assert_called_once() + assert task is None + def test_send_update_task_does_not_fetch_next_during_shutdown(self) -> None: taskworker = TaskWorker( app_module="examples.app:app", @@ -1018,24 +1043,24 @@ def test_push_start_exits_cleanly_on_sigterm() -> None: fake_health = mock.MagicMock() fake_server = mock.MagicMock() - # Deliver the signal from inside the poll sleep, which is where a real - # SIGTERM would land once the server is up and serving. - real_wait = taskworker._shutdown_signal.wait + # Deliver the signal from inside the poll, which is where a real SIGTERM + # would land once the server is up and serving. True keeps the server + # "healthy", so only the flipped bool can end the loop. calls = 0 - def wait(timeout: float) -> bool: + def wait_for_termination(timeout: float | None = None) -> bool: nonlocal calls calls += 1 - if calls == 1: - handlers[signal.SIGTERM](signal.SIGTERM, None) - return real_wait(timeout) + handlers[signal.SIGTERM](signal.SIGTERM, None) + return True + + fake_server.wait_for_termination.side_effect = wait_for_termination with ( _push_worker_grpc_mocks(taskworker, fake_server, fake_health, handlers) as pool_shutdown, mock.patch.object( TaskWorkerProcessingPool, "ready_count", new_callable=mock.PropertyMock, return_value=2 ), - mock.patch.object(taskworker._shutdown_signal, "wait", side_effect=wait), ): exitcode = taskworker.start() @@ -1064,38 +1089,64 @@ def test_push_start_keeps_serving_while_server_is_healthy() -> None: fake_health = mock.MagicMock() fake_server = mock.MagicMock() - # Healthy server: the timeout is always what elapses. - fake_server.wait_for_termination.return_value = True polls = 0 - def wait(timeout: float) -> bool: + def wait_for_termination(timeout: float | None = None) -> bool: nonlocal polls polls += 1 - # Survive several polls, then ask to stop so the test terminates. - if polls < 5: - return False - taskworker._shutdown_signal.request() + # Survive several polls, then SIGTERM so the test terminates. + if polls >= 5: + handlers[signal.SIGTERM](signal.SIGTERM, None) + # Healthy server: the timeout is always what elapses. return True + fake_server.wait_for_termination.side_effect = wait_for_termination + with ( _push_worker_grpc_mocks(taskworker, fake_server, fake_health, handlers) as pool_shutdown, mock.patch.object( TaskWorkerProcessingPool, "ready_count", new_callable=mock.PropertyMock, return_value=2 ), - mock.patch.object(taskworker._shutdown_signal, "wait", side_effect=wait), ): exitcode = taskworker.start() assert exitcode == 0 # We kept looping instead of bailing out on the first poll. assert polls == 5 - # The serve loop must not consult the server's return value at all. - fake_server.wait_for_termination.assert_not_called() fake_server.stop.assert_called_once_with(grace=5) pool_shutdown.assert_called_once_with() +def test_push_start_exits_when_server_terminates_unexpectedly() -> None: + """ + A server that goes away on its own must end the serve loop. + + Otherwise the parent keeps its children alive and keeps touching the health + check file while no longer accepting tasks. + """ + taskworker = _make_push_worker(concurrency=2, warmup_timeout=5) + handlers: dict[int, Callable[..., None]] = {} + + fake_health = mock.MagicMock() + fake_server = mock.MagicMock() + # False means the server terminated, without anyone asking it to. + fake_server.wait_for_termination.return_value = False + + with ( + _push_worker_grpc_mocks(taskworker, fake_server, fake_health, handlers) as pool_shutdown, + mock.patch.object( + TaskWorkerProcessingPool, "ready_count", new_callable=mock.PropertyMock, return_value=2 + ), + ): + exitcode = taskworker.start() + + assert exitcode == 0 + # Noticed on the first poll, rather than spinning forever. + assert fake_server.wait_for_termination.call_count == 1 + pool_shutdown.assert_called_once_with() + + def test_push_start_does_not_start_server_when_shutdown_requested_first() -> None: taskworker = _make_push_worker(concurrency=2, warmup_timeout=5) taskworker._shutdown_signal.request() From a7f685ab4fbc8268b17cf20892cfc339c391a433 Mon Sep 17 00:00:00 2001 From: Markus Unterwaditzer Date: Tue, 25 Aug 2026 11:21:59 +0200 Subject: [PATCH 3/4] test(python): Make the shutdown tests actually observe the exit condition Verified by reintroducing the inverted boolean and re-running: previously one test failed, one hung until timeout, and the SIGTERM test passed regardless. Now all three fail fast. - test_push_start_exits_cleanly_on_sigterm fired SIGTERM from the first poll, so the loop exited after one iteration whichever way the exit condition was read. That is how the inverted boolean got through review. Deliver the signal on the third poll instead, so surviving to it proves the condition is being exercised. - test_push_start_exits_when_server_terminates_unexpectedly spun forever when the loop ignored a terminated server. Bail out after a few polls so it fails with a message instead of hanging CI. --- clients/python/tests/worker/test_worker.py | 31 +++++++++++++++++----- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 96567077..1b2d3e73 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -1046,12 +1046,18 @@ def test_push_start_exits_cleanly_on_sigterm() -> None: # Deliver the signal from inside the poll, which is where a real SIGTERM # would land once the server is up and serving. True keeps the server # "healthy", so only the flipped bool can end the loop. + # + # Deliberately not on the first poll: a loop that exits after one iteration + # no matter what its condition says would pass either way, which is exactly + # how the inverted-boolean bug got through review. Surviving to poll 3 means + # the exit condition is actually being exercised. calls = 0 def wait_for_termination(timeout: float | None = None) -> bool: nonlocal calls calls += 1 - handlers[signal.SIGTERM](signal.SIGTERM, None) + if calls == 3: + handlers[signal.SIGTERM](signal.SIGTERM, None) return True fake_server.wait_for_termination.side_effect = wait_for_termination @@ -1067,8 +1073,9 @@ def wait_for_termination(timeout: float | None = None) -> bool: assert exitcode == 0 assert taskworker._shutdown_signal.is_set() # The handler only flipped a bool, so we polled our way out rather than - # relying on the server being stopped from inside the handler. - assert calls == 1 + # relying on the server being stopped from inside the handler. Three polls + # means we kept serving until the signal, then left promptly. + assert calls == 3 fake_server.stop.assert_called_once_with(grace=5) pool_shutdown.assert_called_once_with() @@ -1130,8 +1137,20 @@ def test_push_start_exits_when_server_terminates_unexpectedly() -> None: fake_health = mock.MagicMock() fake_server = mock.MagicMock() - # False means the server terminated, without anyone asking it to. - fake_server.wait_for_termination.return_value = False + + # False means the server terminated, without anyone asking it to. Bail out + # after a few polls so a loop that ignores this fails the assertion below + # instead of hanging until CI times out. + polls = 0 + + def wait_for_termination(timeout: float | None = None) -> bool: + nonlocal polls + polls += 1 + if polls > 3: + raise AssertionError("serve loop ignored a terminated server") + return False + + fake_server.wait_for_termination.side_effect = wait_for_termination with ( _push_worker_grpc_mocks(taskworker, fake_server, fake_health, handlers) as pool_shutdown, @@ -1143,7 +1162,7 @@ def test_push_start_exits_when_server_terminates_unexpectedly() -> None: assert exitcode == 0 # Noticed on the first poll, rather than spinning forever. - assert fake_server.wait_for_termination.call_count == 1 + assert polls == 1 pool_shutdown.assert_called_once_with() From c61dcfbc7ee8080f792a54a1d5f47bd358192447 Mon Sep 17 00:00:00 2001 From: Markus Unterwaditzer Date: Tue, 25 Aug 2026 11:31:59 +0200 Subject: [PATCH 4/4] test(python): Pin the property ShutdownSignal actually exists for Mutation check: making request() also call _event.set() -- which destroys the whole point of the class, since Event.set() takes the lock a signal handler must not touch -- passed all 18 shutdown tests. The two wakeup tests look like they cover this but cannot: one allows <5s, loose enough to pass whether request() polls or wakes instantly, and the other goes through set(), which the mutation makes identical to request(). Assert the contract directly instead: request() flips the bool and leaves the event alone, set() sets it. The mutation now fails. --- clients/python/tests/worker/test_worker.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 1b2d3e73..165773c6 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -446,12 +446,34 @@ def test_shutdown_signal_wait_returns_immediately_when_requested() -> None: assert time.monotonic() - start < 1 +def test_shutdown_signal_request_does_not_touch_the_event() -> None: + """ + `request()` must not take a lock, so it must not touch the event. + + This is the entire reason the class exists: it is the only method a signal + handler may call, and `Event.set()` takes a non-reentrant lock that can + deadlock against whatever the handler interrupted. A refactor that "helpfully" + sets the event here would be silently unsafe, and the wakeup timings asserted + below are too loose to notice. + """ + shutdown_signal = ShutdownSignal() + shutdown_signal.request() + + assert shutdown_signal.is_set() is True + assert shutdown_signal._event.is_set() is False + + # ...whereas set() is allowed to, and does. + shutdown_signal.set() + assert shutdown_signal._event.is_set() is True + + def test_shutdown_signal_wait_wakes_on_request_from_another_thread() -> None: shutdown_signal = ShutdownSignal() threading.Timer(0.1, shutdown_signal.request).start() start = time.monotonic() assert shutdown_signal.wait(30) is True + # request() cannot wake the event, so this is the poll interval, not instant. assert time.monotonic() - start < 5