fix: Re-dispatch SIGTERM to the handler add_signal_handler displaces

python/restate/server.py

@@ -13,7 +13,7 @@
 import asyncio
 import logging
 import signal
-from typing import Dict, Set, TypedDict, Literal
+from typing import Any, Dict, Set, TypedDict, Literal
 
 from restate.discovery import compute_discovery_json
 from restate.endpoint import Endpoint
@@ -216,18 +216,34 @@ def asgi_app(endpoint: Endpoint) -> RestateAppT:
 
     active_channels: Set[ReceiveChannel] = set()
     sigterm_installed = False
+    displaced_sigterm_handler: Any = None
 
     def _on_sigterm() -> None:
         """Notify all active receive channels of graceful shutdown."""
         for ch in active_channels:
             ch.notify_shutdown()
+        # Re-dispatch to the handler we displaced when installing this one, so a
+        # host ASGI server still learns about SIGTERM and can start its own
+        # graceful shutdown. Called from the event loop rather than from signal
+        # context, which is why frame is None; handlers must accept that.
+        handler = displaced_sigterm_handler
+        if callable(handler):
+            handler(signal.SIGTERM, None)
 
     async def app(scope: Scope, receive: Receive, send: Send):
-        nonlocal sigterm_installed
+        nonlocal sigterm_installed, displaced_sigterm_handler
         if not sigterm_installed:
             loop = asyncio.get_running_loop()
             try:
+                # add_signal_handler installs a dummy handler via signal.signal
+                # underneath, silently displacing whatever the host ASGI server
+                # registered (uvicorn's Server.handle_exit, for one). Remember it
+                # so _on_sigterm can re-dispatch: without that the host never
+                # sees SIGTERM, never drains, and is force-killed when the
+                # supervisor's stop grace period expires.
+                displaced = signal.getsignal(signal.SIGTERM)
                 loop.add_signal_handler(signal.SIGTERM, _on_sigterm)
+                displaced_sigterm_handler = displaced
             except (NotImplementedError, RuntimeError, ValueError):
                 pass  # Windows or non-main thread
             sigterm_installed = True

tests/server.py

@@ -9,12 +9,47 @@
 #  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
 #
 import asyncio
+import signal
+from typing import Any, List, Optional
+from types import FrameType
 from unittest.mock import Mock
 
 import pytest
 
 from restate.endpoint import Endpoint
 
+HEALTH_SCOPE: Any = {
+    "type": "http",
+    "asgi": {"version": "3.0"},
+    "http_version": "1.1",
+    "method": "GET",
+    "scheme": "http",
+    "path": "/restate/health",
+    "raw_path": b"/restate/health",
+    "query_string": b"",
+    "headers": [],
+    "client": ("127.0.0.1", 1234),
+    "server": ("127.0.0.1", 9080),
+}
+
+
+async def drive_health_request(app: Any) -> List[Any]:
+    """Send one health request through the app, returning the messages it sent.
+
+    The SIGTERM handler is installed on the first request the app serves, so a
+    request is how a test reaches that code path.
+    """
+    sent: List[Any] = []
+
+    async def receive():
+        return {"type": "http.request", "body": b"", "more_body": False}
+
+    async def send(message):
+        sent.append(message)
+
+    await app(HEALTH_SCOPE, receive, send)
+    return sent
+
 
 @pytest.fixture(scope="session")
 def anyio_backend():
@@ -61,3 +96,61 @@ async def send(message):
 
     response_starts = [message for message in sent if message["type"] == "http.response.start"]
     assert [message["status"] for message in response_starts] == [200]
+
+
+@pytest.fixture
+async def restore_sigterm():
+    """SIGTERM disposition and the loop handler are process-global; put them back.
+
+    Async so that teardown still runs inside the event loop, which
+    ``remove_signal_handler`` needs.
+    """
+    original = signal.getsignal(signal.SIGTERM)
+    try:
+        yield
+    finally:
+        try:
+            asyncio.get_running_loop().remove_signal_handler(signal.SIGTERM)
+        except (NotImplementedError, RuntimeError, ValueError):
+            pass
+        signal.signal(signal.SIGTERM, original)
+
+
+async def test_sigterm_redispatches_to_the_handler_it_displaced(restore_sigterm):
+    """A host ASGI server's SIGTERM handler must survive our installing ours.
+
+    ``loop.add_signal_handler`` installs a dummy handler through
+    ``signal.signal``, so it silently replaces whatever the hosting server
+    registered. uvicorn registers ``Server.handle_exit`` there, and losing it
+    means the server never drains and gets force-killed by its supervisor.
+    """
+    host_handler_calls: List[int] = []
+
+    def host_handler(signum: int, frame: Optional[FrameType]) -> None:
+        host_handler_calls.append(signum)
+
+    # Stand in for uvicorn's Server.capture_signals().
+    signal.signal(signal.SIGTERM, host_handler)
+
+    await drive_health_request(Endpoint().app())
+
+    # Precondition: installing ours displaced theirs.
+    assert signal.getsignal(signal.SIGTERM) is not host_handler
+
+    signal.raise_signal(signal.SIGTERM)
+    for _ in range(50):
+        if host_handler_calls:
+            break
+        await asyncio.sleep(0.02)
+
+    assert host_handler_calls == [signal.SIGTERM], "the displaced host handler was never re-dispatched"
+
+
+async def test_sigterm_without_a_previous_handler_is_harmless(restore_sigterm):
+    """Nothing to re-dispatch to is the common standalone case, not an error."""
+    signal.signal(signal.SIGTERM, signal.SIG_DFL)
+
+    await drive_health_request(Endpoint().app())
+
+    signal.raise_signal(signal.SIGTERM)
+    await asyncio.sleep(0.05)  # a raise that reached SIG_DFL would have killed us