fix: restore signal handlers on lock release

codecarbon/lock.py

@@ -27,18 +27,20 @@ def __init__(self):
             self.release
         )  # 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()
-        # 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
+        # Reentrant: _handle_exit -> release() can fire on a thread already holding it.
+        self._thread_lock = threading.RLock()
+        # Previous signal handlers, restored on release().
+        self._previous_handlers = {}
 
     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:
+            os.kill(os.getpid(), signum)
 
     def acquire(self):
         """Creates a lock file and ensures it's the only instance running."""
@@ -48,19 +50,34 @@ def acquire(self):
                 with open(LOCKFILE, "x") as _:
                     logger.debug(f"Lock file created. Path: {LOCKFILE}")
                     self._has_created_lock = True
+                # Only now that we own the lock file: a failed acquire() must not
+                # leave the host application's handlers hijacked for good, since
+                # release() is never reached on that path.
+                if threading.current_thread() is threading.main_thread():
+                    for sig in (signal.SIGINT, signal.SIGTERM):
+                        self._previous_handlers[sig] = signal.signal(
+                            sig, self._handle_exit
+                        )
             except FileExistsError:
                 logger.debug(
                     f"Lock file {LOCKFILE} already exists. This usually means another instance of codecarbon is running. You can safely delete it if you want or use allow_multiple_runs parameter to always bypass it."
                 )
                 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.release)
             try:
                 # Remove the lock file only if it was created by this instance
                 if self._has_created_lock:
+                    self._has_created_lock = False
                     os.remove(LOCKFILE)
             except OSError as e:
                 logger.debug(f"Error: {e}")

tests/test_lock.py

@@ -1,12 +1,29 @@
+import atexit
+import os
+import signal
 import threading
 import unittest
 from unittest.mock import mock_open, patch
 
 from codecarbon.lock import LOCKFILE, Lock
 
 
-class TestLock(unittest.TestCase):
+class SignalSafeTestCase(unittest.TestCase):
+    """acquire() installs process-wide handlers: put the originals back."""
+
     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)
+
+
+class TestLock(SignalSafeTestCase):
+    def setUp(self):
+        super().setUp()
         self.lock = Lock()
 
     @patch("codecarbon.lock.os.remove")
@@ -30,6 +47,17 @@ def test_release_removes_lock_file(self, mock_file, mock_remove):
         self.lock.release()
         mock_remove.assert_called_once_with(LOCKFILE)
 
+    @patch("codecarbon.lock.os.remove")
+    @patch("codecarbon.lock.open", new_callable=mock_open)
+    def test_release_is_idempotent(self, mock_file, mock_remove):
+        # A second release() must not delete the lock file again: by then it may
+        # have been re-created by another instance of codecarbon.
+        self.lock.acquire()
+        self.lock.release()
+        self.lock.release()
+        mock_remove.assert_called_once_with(LOCKFILE)
+        self.assertFalse(self.lock._has_created_lock)
+
     @patch("codecarbon.lock.os.remove")
     @patch("codecarbon.lock.open", new_callable=mock_open)
     def test_release_does_not_release_when_not_created_by_this_instance(
@@ -77,5 +105,123 @@ def thread_target():
         self.assertTrue(mock_remove.called)
 
 
+class TestLockSignalHandlers(SignalSafeTestCase):
+    """The lock must not permanently steal the host application's handlers."""
+
+    @patch("codecarbon.lock.os.remove")
+    @patch("codecarbon.lock.open", new_callable=mock_open)
+    def test_release_restores_previous_handlers(self, mock_file, mock_remove):
+        def sentinel(signum, frame):
+            pass
+
+        signal.signal(signal.SIGTERM, sentinel)
+        lock = Lock()
+        lock.acquire()
+        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")
+    @patch("codecarbon.lock.open", new_callable=mock_open)
+    def test_signal_is_forwarded_to_previous_handler(self, mock_file, mock_remove):
+        received = []
+
+        signal.signal(signal.SIGTERM, lambda signum, frame: received.append(signum))
+        lock = Lock()
+        lock.acquire()
+        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")
+    @patch("codecarbon.lock.open", new_callable=mock_open)
+    def test_default_disposition_is_reproduced(self, mock_file, 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.acquire()
+
+        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")
+    @patch("codecarbon.lock.open", new_callable=mock_open)
+    def test_ignored_signal_stays_ignored(self, mock_file, mock_remove, mock_kill):
+        signal.signal(signal.SIGTERM, signal.SIG_IGN)
+        lock = Lock()
+        lock.acquire()
+
+        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
+            # Without this, a reverted (non-reentrant) lock.py wedges this thread
+            # while still holding the lock, and the atexit hook then blocks the
+            # interpreter forever on exit instead of letting the test fail.
+            atexit.unregister(lock.release)
+            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"
+
+    @patch("codecarbon.lock.open", new_callable=mock_open)
+    def test_failed_acquire_leaves_the_handlers_alone(self, mock_file):
+        # Another instance already holds the lock: stop() returns early and never
+        # calls release(), so acquire() must not have taken the handlers at all.
+        def sentinel(signum, frame):
+            pass
+
+        signal.signal(signal.SIGTERM, sentinel)
+        mock_file.side_effect = FileExistsError
+        lock = Lock()
+        with self.assertRaises(FileExistsError):
+            lock.acquire()
+
+        self.assertIs(signal.getsignal(signal.SIGTERM), sentinel)
+        self.assertIs(
+            signal.getsignal(signal.SIGINT), self.original_handlers[signal.SIGINT]
+        )
+
+
 if __name__ == "__main__":
     unittest.main()