fix(python): Shut workers down without raising from signal handlers

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

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,15 +559,21 @@ 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.
             self._await_children_warm()
 
             # 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()