diff --git a/.dockerignore b/.dockerignore index 637ab31..c3d2503 100644 --- a/.dockerignore +++ b/.dockerignore @@ -187,4 +187,9 @@ fabric.properties # Test results junit/ -.github/ \ No newline at end of file +.github/ + +# Not needed to build the image +.git/ +.claude/ +tests/ diff --git a/src/main.py b/src/main.py index e2d90d9..6b5fe42 100644 --- a/src/main.py +++ b/src/main.py @@ -9,6 +9,7 @@ from configuration.parser import process_command_line from mqtt_gateway import MqttGateway +from shutdown import run_until_terminated if __name__ == "__main__": # Keep this at the top! @@ -28,4 +29,4 @@ mqtt_gateway = MqttGateway(configuration) - asyncio.run(mqtt_gateway.run(), debug=debug_log_enabled()) + asyncio.run(run_until_terminated(mqtt_gateway.run()), debug=debug_log_enabled()) diff --git a/src/shutdown.py b/src/shutdown.py new file mode 100644 index 0000000..61157b9 --- /dev/null +++ b/src/shutdown.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import asyncio +import logging +import signal +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Coroutine + +LOG = logging.getLogger(__name__) + + +async def run_until_terminated(main: Coroutine[Any, Any, None]) -> None: + """Run ``main``, cancelling it cleanly when the process receives SIGTERM. + + Python does not handle SIGTERM by default, and as PID 1 in a container it + would simply ignore it. Cancelling the task lets its ``finally`` blocks run. + """ + loop = asyncio.get_running_loop() + task = loop.create_task(main) + terminated = False + + def request_shutdown() -> None: + nonlocal terminated + terminated = True + LOG.info("Received SIGTERM, shutting down") + task.cancel() + + handler_installed = False + try: + loop.add_signal_handler(signal.SIGTERM, request_shutdown) + handler_installed = True + except (NotImplementedError, RuntimeError, ValueError): + LOG.debug("Cannot install a SIGTERM handler on this platform") + + try: + await task + except asyncio.CancelledError: + if not terminated: + raise + finally: + if handler_installed: + loop.remove_signal_handler(signal.SIGTERM) diff --git a/tests/test_shutdown.py b/tests/test_shutdown.py new file mode 100644 index 0000000..bc90f3b --- /dev/null +++ b/tests/test_shutdown.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import asyncio +import os +import signal +import sys +import unittest + +import pytest + +from shutdown import run_until_terminated + + +@pytest.mark.skipif(sys.platform == "win32", reason="needs POSIX signals") +class TestRunUntilTerminated(unittest.IsolatedAsyncioTestCase): + async def test_sigterm_cancels_main_and_runs_cleanup(self) -> None: + cleaned_up = False + + async def main() -> None: + nonlocal cleaned_up + try: + await asyncio.sleep(30) + finally: + cleaned_up = True + + loop = asyncio.get_running_loop() + loop.call_later(0.05, os.kill, os.getpid(), signal.SIGTERM) + + await run_until_terminated(main()) + + assert cleaned_up + + async def test_normal_completion_returns(self) -> None: + done = False + + async def main() -> None: + nonlocal done + done = True + + await run_until_terminated(main()) + + assert done + + async def test_exceptions_from_main_propagate(self) -> None: + async def main() -> None: + msg = "boom" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError): + await run_until_terminated(main()) + + async def test_external_cancellation_is_not_swallowed(self) -> None: + async def main() -> None: + await asyncio.sleep(30) + + runner = asyncio.ensure_future(run_until_terminated(main())) + await asyncio.sleep(0.05) + runner.cancel() + + with pytest.raises(asyncio.CancelledError): + await runner + + async def test_sigterm_handler_removed_afterwards(self) -> None: + async def main() -> None: + return None + + await run_until_terminated(main()) + + assert signal.getsignal(signal.SIGTERM) == signal.SIG_DFL