diff --git a/python/restate/server.py b/python/restate/server.py index c5e233f9..43e7bfa4 100644 --- a/python/restate/server.py +++ b/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 diff --git a/tests/server.py b/tests/server.py index 83bd3143..b93615d8 100644 --- a/tests/server.py +++ b/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()