fix: restore signal handlers on lock release

codecarbon/lock.py

@@ -23,22 +23,37 @@ class Lock:
     def __init__(self):
         self._has_created_lock = False
         self.lockfile_path = LOCKFILE
+        # Keep one reference to the bound method: atexit.unregister() matches on
+        # identity, so registering `self.release` twice would not be undoable.
+        self._atexit_hook = self.release
         atexit.register(
-            self.release
+            self._atexit_hook
         )  # Ensure release() is called on unexpected exit of the user's python code
         # If there is more than one thread add a lock
-        self._thread_lock = threading.Lock()
+        # Reentrant: a signal delivered while this thread is inside acquire() or
+        # release() runs _handle_exit -> release() on that same thread, which a
+        # plain Lock would deadlock on.
+        self._thread_lock = threading.RLock()
+        # Previous signal handlers, restored on release() so we do not take
+        # permanent ownership of the host application's signal disposition.
+        self._previous_handlers = {}
         # If the current thread is the main thread, register signal handlers
         if threading.current_thread() is threading.main_thread():
             # Register signal handlers to ensure lock release on interruption
-            signal.signal(signal.SIGINT, self._handle_exit)  # Ctrl+C
-            signal.signal(signal.SIGTERM, self._handle_exit)  # Termination signal
+            for sig in (signal.SIGINT, signal.SIGTERM):  # Ctrl+C, termination signal
+                self._previous_handlers[sig] = signal.signal(sig, self._handle_exit)
 
     def _handle_exit(self, signum, frame):
-        """Ensures the lock file is removed when the script is interrupted."""
-        logger.debug(f"Signal {signum} received. Releasing lock and exiting.")
-        self.release()
-        raise SystemExit(1)  # Exit gracefully to prevent further execution
+        """Releases the lock, then delegates to the handler we replaced."""
+        logger.debug(f"Signal {signum} received. Releasing lock.")
+        previous = self._previous_handlers.get(signum, signal.SIG_DFL)
+        self.release()  # also restores the previous handlers
+        if callable(previous):
+            return previous(signum, frame)
+        if previous == signal.SIG_DFL:
+            # Reproduce the default disposition (usually terminate).
+            os.kill(os.getpid(), signum)
+        # signal.SIG_IGN: nothing to do
 
     def acquire(self):
         """Creates a lock file and ensures it's the only instance running."""
@@ -55,9 +70,15 @@ def acquire(self):
                 raise
 
     def release(self):
-        """Removes the lock file on exit."""
+        """Removes the lock file and restores the signal handlers we replaced."""
         with self._thread_lock:
             logger.debug("Removing the lock")
+            while self._previous_handlers:
+                sig, handler = self._previous_handlers.popitem()
+                # Only restore if nobody installed another handler after us.
+                if signal.getsignal(sig) == self._handle_exit:
+                    signal.signal(sig, handler)
+            atexit.unregister(self._atexit_hook)
             try:
                 # Remove the lock file only if it was created by this instance
                 if self._has_created_lock:

tests/test_lock.py

@@ -1,3 +1,5 @@
+import os
+import signal
 import threading
 import unittest
 from unittest.mock import mock_open, patch
@@ -77,5 +79,103 @@ def thread_target():
         self.assertTrue(mock_remove.called)
 
 
+class TestLockSignalHandlers(unittest.TestCase):
+    """The lock must not permanently steal the host application's handlers."""
+
+    def setUp(self):
+        self.original_handlers = {
+            sig: signal.getsignal(sig) for sig in (signal.SIGINT, signal.SIGTERM)
+        }
+
+    def tearDown(self):
+        for sig, handler in self.original_handlers.items():
+            signal.signal(sig, handler)
+
+    def test_release_restores_previous_handlers(self):
+        def sentinel(signum, frame):
+            pass
+
+        signal.signal(signal.SIGTERM, sentinel)
+        lock = Lock()
+        self.assertEqual(signal.getsignal(signal.SIGTERM), lock._handle_exit)
+        lock.release()
+        self.assertIs(signal.getsignal(signal.SIGTERM), sentinel)
+        self.assertIs(
+            signal.getsignal(signal.SIGINT), self.original_handlers[signal.SIGINT]
+        )
+
+    @unittest.skipIf(
+        not hasattr(signal, "raise_signal"), "requires signal.raise_signal (3.8+)"
+    )
+    @patch("codecarbon.lock.os.remove")
+    def test_signal_is_forwarded_to_previous_handler(self, mock_remove):
+        received = []
+
+        signal.signal(signal.SIGTERM, lambda signum, frame: received.append(signum))
+        lock = Lock()
+        signal.raise_signal(signal.SIGTERM)
+        self.assertEqual(received, [signal.SIGTERM])
+        self.assertFalse(lock._previous_handlers)
+
+    @unittest.skipIf(
+        not hasattr(signal, "raise_signal"), "requires signal.raise_signal (3.8+)"
+    )
+    @patch("codecarbon.lock.os.kill")
+    @patch("codecarbon.lock.os.remove")
+    def test_default_disposition_is_reproduced(self, mock_remove, mock_kill):
+        # No handler installed by the host application: the default disposition
+        # of SIGTERM is to terminate, which the lock must reproduce after having
+        # released the lock instead of silently swallowing the signal.
+        signal.signal(signal.SIGTERM, signal.SIG_DFL)
+        lock = Lock()
+        lock._has_created_lock = True
+
+        signal.raise_signal(signal.SIGTERM)
+
+        mock_kill.assert_called_once_with(os.getpid(), signal.SIGTERM)
+        # The lock was released before re-raising, and the default disposition
+        # was put back so the re-raised signal is not caught again.
+        self.assertTrue(mock_remove.called)
+        self.assertIs(signal.getsignal(signal.SIGTERM), signal.SIG_DFL)
+
+    @unittest.skipIf(
+        not hasattr(signal, "raise_signal"), "requires signal.raise_signal (3.8+)"
+    )
+    @patch("codecarbon.lock.os.kill")
+    @patch("codecarbon.lock.os.remove")
+    def test_ignored_signal_stays_ignored(self, mock_remove, mock_kill):
+        signal.signal(signal.SIGTERM, signal.SIG_IGN)
+        lock = Lock()
+        lock._has_created_lock = True
+
+        signal.raise_signal(signal.SIGTERM)
+
+        # The host application asked to ignore SIGTERM: release the lock, but do
+        # not terminate on its behalf.
+        self.assertTrue(mock_remove.called)
+        mock_kill.assert_not_called()
+        self.assertIs(signal.getsignal(signal.SIGTERM), signal.SIG_IGN)
+
+    @patch("codecarbon.lock.os.remove")
+    def test_release_from_within_the_critical_section_does_not_deadlock(
+        self, mock_remove
+    ):
+        """A signal handler fires on the thread that may already hold the lock."""
+        done = threading.Event()
+
+        def hold_then_release():
+            # Built off the main thread : no signal handlers to restore, so this
+            # exercises the thread lock only (signal.signal is main-thread only).
+            lock = Lock()
+            lock._has_created_lock = True
+            with lock._thread_lock:
+                lock.release()
+            done.set()
+
+        worker = threading.Thread(target=hold_then_release, daemon=True)
+        worker.start()
+        assert done.wait(timeout=5), "release() deadlocked on its own thread lock"
+
+
 if __name__ == "__main__":
     unittest.main()