refactor: graceful_shutdown.py

Author: vgorkavenko

src/modules/common/graceful_shutdown.py

@@ -1,6 +1,7 @@
 import signal
-from collections.abc import Callable, Iterator
-from contextlib import contextmanager
+import threading
+from collections.abc import Iterator
+from contextlib import ExitStack, contextmanager
 from types import FrameType
 
 from src.metrics.logging import logging
@@ -9,35 +10,28 @@
 logger = logging.getLogger(__name__)
 
 
-type SignalHandler = Callable[[int, FrameType | None], object] | int | None
-
-
-def _shutdown_signal_handler(signum: int, _frame: FrameType | None) -> None:
-    try:
-        signal_name = signal.Signals(signum).name
-    except ValueError:
-        signal_name = str(signum)
-
-    logger.info({'msg': 'Received shutdown signal. Requesting graceful exit.', 'signal': signal_name})
-    raise SystemExit(0)
-
-
 @contextmanager
 def graceful_shutdown_signal_handlers() -> Iterator[None]:
-    previous_handlers: dict[signal.Signals, SignalHandler] = {}
-
-    for stop_signal in (signal.SIGINT, signal.SIGTERM):
-        try:
-            previous_handlers[stop_signal] = signal.getsignal(stop_signal)
-            signal.signal(stop_signal, _shutdown_signal_handler)
-        except ValueError:
-            logger.info({'msg': 'Cannot register signal handler outside main thread', 'signal': stop_signal})
-
-    try:
+    """
+    Temporarily convert `SIGINT` and `SIGTERM` to `SystemExit(0)`.
+
+    Use this around a top-level module run loop so regular `finally` cleanup
+    still runs on process shutdown. The previous handlers are restored when the
+    context exits.
+    """
+    if threading.current_thread() is not threading.main_thread():
+        logger.info({'msg': 'Cannot manage signal handlers outside main thread'})
+        yield
+        return
+
+    def shutdown_signal_handler(signum: int, _frame: FrameType | None) -> None:
+        logger.info(
+            {'msg': 'Received shutdown signal. Requesting graceful exit.', 'signal': signal.Signals(signum).name}
+        )
+        raise SystemExit(0)
+
+    with ExitStack() as stack:
+        for stop_signal in (signal.SIGINT, signal.SIGTERM):
+            previous_handler = signal.signal(stop_signal, shutdown_signal_handler)
+            stack.callback(signal.signal, stop_signal, previous_handler)
         yield
-    finally:
-        for stop_signal, previous_handler in previous_handlers.items():
-            try:
-                signal.signal(stop_signal, previous_handler)
-            except ValueError:
-                logger.info({'msg': 'Cannot restore signal handler outside main thread', 'signal': stop_signal})

src/modules/oracles/common/runtime.py

@@ -125,13 +125,15 @@ def build_staking_module_web3(module_name: str) -> Web3StakingModule:
 
 def run_oracle_module(module: OracleModule):
     module.check_contract_configs()
+
     try:
         module.w3.telemetry_data_bus.send_telemetry(TelemetryEventId.ORACLE_STARTUP)
     except Exception:
         logger.warning({'msg': 'Failed to send startup telemetry to DataBus.'}, exc_info=True)
+
     try:
-        # Docker sends SIGTERM on `docker stop`; convert it to SystemExit so `finally`
-        # runs module.shutdown() and cleanup is not skipped.
+        # Convert termination signals to SystemExit so regular cleanup in
+        # `finally` still runs during process shutdown.
         with graceful_shutdown_signal_handlers():
             if variables.DAEMON:
                 module.run_as_daemon()

tests/modules/submodules/test_oracle_module.py

@@ -1,4 +1,3 @@
-import signal
 from http import HTTPStatus
 from unittest.mock import MagicMock, Mock, patch
 
@@ -143,7 +142,6 @@ def test_run_oracle_module__single_cycle__calls_shutdown(oracle: OracleModule, m
     oracle.shutdown = Mock()
 
     with (
-        patch("src.modules.common.graceful_shutdown.signal.getsignal", return_value=signal.SIG_DFL),
         patch("src.modules.common.graceful_shutdown.signal.signal"),
     ):
         run_oracle_module(oracle)
@@ -162,7 +160,6 @@ def test_run_oracle_module__daemon_exit__calls_shutdown(oracle: OracleModule, mo
     oracle.shutdown = Mock()
 
     with (
-        patch("src.modules.common.graceful_shutdown.signal.getsignal", return_value=signal.SIG_DFL),
         patch("src.modules.common.graceful_shutdown.signal.signal"),
         pytest.raises(SystemExit, match="0"),
     ):