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..fa026692 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,22 +583,30 @@ 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 - pass + # Poll so a signal handler that only flipped a bool still gets us + # out, while also noticing a server that terminated on its own. + # + # 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: 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 +615,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 +694,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 +709,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 +784,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 +810,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: @@ -761,13 +839,32 @@ 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: """ 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..165773c6 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,94 @@ 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_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 + + +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 +543,93 @@ 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_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", + 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 +1008,226 @@ 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, 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 + if calls == 3: + handlers[signal.SIGTERM](signal.SIGTERM, None) + return True + + fake_server.wait_for_termination.side_effect = wait_for_termination + 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" + ): + 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. 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() + + +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() + + polls = 0 + + def wait_for_termination(timeout: float | None = None) -> bool: + nonlocal polls + polls += 1 + # 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 + ), + ): + exitcode = taskworker.start() + + assert exitcode == 0 + # We kept looping instead of bailing out on the first poll. + assert polls == 5 + 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. 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, + 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 polls == 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() + + 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 +1240,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(