From 004a7e31be42b300380a79117570f911389b4db6 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Thu, 20 Aug 2026 08:13:10 +0200 Subject: [PATCH] fix: restore signal handlers on lock release `Lock` installed SIGINT/SIGTERM handlers and threw away the previous ones, so the host application's handlers were destroyed and Ctrl-C stopped raising KeyboardInterrupt. Save the previous handlers, chain to them from `_handle_exit`, and restore them in `release()`, unregistering the atexit hook so a released lock is not pinned. Handlers are installed in `acquire()` after `open(LOCKFILE, "x")` succeeds, not in `__init__`. On the "another instance is already running" path `acquire()` raises, the tracker sets `_another_instance_already_running`, and `stop()` returns at its early guard without ever reaching `release()` -- so handlers installed in the constructor stayed hijacked for the life of the process. The thread lock is reentrant: `_handle_exit` calls `release()`, which takes `_thread_lock`, so a signal delivered while the same thread was inside `acquire()`/`release()` deadlocked on a plain `Lock`. `release()` is also idempotent now (moved here from #1336, since it edits the same few lines this branch already rewrites), and the `_atexit_hook` indirection is dropped: `atexit.unregister()` compares with `==`, not identity, so a bound method unregisters fine. Tests cover the default and ignored signal dispositions, and the deadlock test unregisters its atexit hook so a reverted lock.py fails the suite instead of wedging the interpreter at exit. Co-Authored-By: Claude Opus 5 (1M context) --- codecarbon/lock.py | 39 ++++++++---- tests/test_lock.py | 148 ++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 175 insertions(+), 12 deletions(-) diff --git a/codecarbon/lock.py b/codecarbon/lock.py index 38d112324..68ba547f9 100644 --- a/codecarbon/lock.py +++ b/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,6 +50,14 @@ 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." @@ -55,12 +65,19 @@ 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.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}") diff --git a/tests/test_lock.py b/tests/test_lock.py index aafb46a1b..46fba39bd 100644 --- a/tests/test_lock.py +++ b/tests/test_lock.py @@ -1,3 +1,6 @@ +import atexit +import os +import signal import threading import unittest from unittest.mock import mock_open, patch @@ -5,8 +8,22 @@ 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()