Fix ASGI SIGTERM handler to chain to host server and skip lifespan scope

python/restate/server.py

@@ -251,26 +251,33 @@ async def get_codec() -> Optional[JournalValueCodec]:
             return codec_config
         return await _codec_task()
 
-    def _on_sigterm() -> None:
-        """Notify all active receive channels of graceful shutdown."""
-        for ch in active_channels:
-            ch.notify_shutdown()
-
     async def app(scope: Scope, receive: Receive, send: Send):
         nonlocal sigterm_installed
+        if scope["type"] == "lifespan":
+            raise LifeSpanNotImplemented()
+        if scope["type"] != "http":
+            raise NotImplementedError(f"Unknown scope type {scope['type']}")
+
         if not sigterm_installed:
             loop = asyncio.get_running_loop()
+            prev_handle = getattr(loop, "_signal_handlers", {}).get(signal.SIGTERM)
+            prev_signal = signal.getsignal(signal.SIGTERM)
+
+            def _on_sigterm() -> None:
+                for ch in active_channels:
+                    ch.notify_shutdown()
+                # ASGI hosts use either asyncio or signal.signal for shutdown.
+                if prev_handle is not None:
+                    prev_handle._run()
+                elif callable(prev_signal):
+                    prev_signal(signal.SIGTERM, None)
+
             try:
                 loop.add_signal_handler(signal.SIGTERM, _on_sigterm)
             except (NotImplementedError, RuntimeError, ValueError):
                 pass  # Windows or non-main thread
             sigterm_installed = True
         try:
-            if scope["type"] == "lifespan":
-                raise LifeSpanNotImplemented()
-            if scope["type"] != "http":
-                raise NotImplementedError(f"Unknown scope type {scope['type']}")
-
             request_path = scope["path"]
             assert isinstance(request_path, str)
             request: ParsedPath = parse_path(request_path)
@@ -330,8 +337,6 @@ async def app(scope: Scope, receive: Receive, send: Send):
             finally:
                 active_channels.discard(receive_channel)
                 await receive_channel.close()
-        except LifeSpanNotImplemented as e:
-            raise e
         except Exception as e:
             logger.exception("Exception in ASGI app")
             raise e

tests/server.py

@@ -9,11 +9,13 @@
 #  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
 #
 import asyncio
-from unittest.mock import Mock
+import signal
+from unittest.mock import AsyncMock, Mock
 
 import pytest
-
+from restate import server
 from restate.endpoint import Endpoint
+from restate.server import LifeSpanNotImplemented
 
 
 @pytest.fixture(scope="session")
@@ -61,3 +63,62 @@ 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.mark.parametrize("registration", ["loop", "signal"])
+async def test_sigterm_chains_host_after_notifying_channels(monkeypatch: pytest.MonkeyPatch, registration):
+    loop = asyncio.get_running_loop()
+    previous_signal = signal.getsignal(signal.SIGTERM)
+    previous_handle = getattr(loop, "_signal_handlers", {}).get(signal.SIGTERM)
+    calls = []
+    host = Mock(side_effect=lambda *args: calls.append("host"))
+    try:
+        if registration == "loop":
+            loop.add_signal_handler(signal.SIGTERM, host, "shutdown")
+        else:
+            loop.remove_signal_handler(signal.SIGTERM)
+            signal.signal(signal.SIGTERM, host)
+
+        endpoint = Endpoint()
+        endpoint.services["test"] = Mock()
+        monkeypatch.setattr(server, "VMWrapper", Mock())
+
+        async def process(vm, handler, headers, receive_channel, send, **kwargs):
+            original_notify = receive_channel.notify_shutdown
+
+            def notify():
+                calls.append("channel")
+                original_notify()
+
+            monkeypatch.setattr(receive_channel, "notify_shutdown", notify)
+            getattr(loop, "_signal_handlers", {})[signal.SIGTERM]._run()
+            assert calls == ["channel", "host"]
+            await asyncio.wait_for(receive_channel.block_until_http_input_closed(), 1)
+
+        monkeypatch.setattr(server, "process_invocation_to_completion", process)
+        receive = AsyncMock(return_value={"type": "http.disconnect"})
+        await endpoint.app()(
+            {"type": "http", "path": "/invoke/test/handler", "headers": []},
+            receive,
+            AsyncMock(),
+        )
+        if registration == "loop":
+            host.assert_called_once_with("shutdown")
+        else:
+            host.assert_called_once_with(signal.SIGTERM, None)
+    finally:
+        loop.remove_signal_handler(signal.SIGTERM)
+        if previous_handle is not None:
+            loop.add_signal_handler(signal.SIGTERM, previous_handle._callback, *previous_handle._args)
+        else:
+            signal.signal(signal.SIGTERM, previous_signal)
+
+
+@pytest.mark.parametrize("scope_type", ["lifespan", "websocket"])
+async def test_non_http_scope_does_not_install_signal_handler(monkeypatch: pytest.MonkeyPatch, scope_type):
+    install = Mock()
+    monkeypatch.setattr(asyncio.get_running_loop(), "add_signal_handler", install)
+    expected = LifeSpanNotImplemented if scope_type == "lifespan" else NotImplementedError
+    with pytest.raises(expected):
+        await Endpoint().app()({"type": scope_type}, AsyncMock(), AsyncMock())
+    install.assert_not_called()