From d4f89f61596c228c83291d401ee1c307455347b5 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Wed, 12 Aug 2026 16:55:33 +0200 Subject: [PATCH 1/3] 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(). Also unregister the atexit hook so a released lock is not pinned. Co-Authored-By: Claude Opus 5 (1M context) --- codecarbon/lock.py | 34 ++++++++++++++++++++++++++-------- tests/test_lock.py | 40 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 8 deletions(-) diff --git a/codecarbon/lock.py b/codecarbon/lock.py index 38d112324..a25f1aa1d 100644 --- a/codecarbon/lock.py +++ b/codecarbon/lock.py @@ -23,22 +23,34 @@ 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() + # 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 +67,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: diff --git a/tests/test_lock.py b/tests/test_lock.py index aafb46a1b..aaae34ce8 100644 --- a/tests/test_lock.py +++ b/tests/test_lock.py @@ -1,3 +1,4 @@ +import signal import threading import unittest from unittest.mock import mock_open, patch @@ -77,5 +78,44 @@ 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) + + if __name__ == "__main__": unittest.main() From a2e1eab0e1f8be930d5a28e0dda4955f372011fe Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Wed, 12 Aug 2026 17:43:24 +0200 Subject: [PATCH 2/3] test: cover default and ignored signal dispositions Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_lock.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/test_lock.py b/tests/test_lock.py index aaae34ce8..08c46f66d 100644 --- a/tests/test_lock.py +++ b/tests/test_lock.py @@ -1,3 +1,4 @@ +import os import signal import threading import unittest @@ -116,6 +117,45 @@ def test_signal_is_forwarded_to_previous_handler(self, mock_remove): 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) + if __name__ == "__main__": unittest.main() From a91020f05339473c074ec2c6315ea9631c5330f4 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Wed, 12 Aug 2026 19:59:18 +0200 Subject: [PATCH 3/3] fix(lock): make the thread lock reentrant for the signal path _handle_exit calls release(), which takes _thread_lock; a signal delivered while the same thread is inside acquire()/release() deadlocked on a plain Lock. Co-Authored-By: Claude Opus 5 (1M context) --- codecarbon/lock.py | 5 ++++- tests/test_lock.py | 20 ++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/codecarbon/lock.py b/codecarbon/lock.py index a25f1aa1d..4c60119df 100644 --- a/codecarbon/lock.py +++ b/codecarbon/lock.py @@ -30,7 +30,10 @@ def __init__(self): 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 = {} diff --git a/tests/test_lock.py b/tests/test_lock.py index 08c46f66d..e211ed8ec 100644 --- a/tests/test_lock.py +++ b/tests/test_lock.py @@ -156,6 +156,26 @@ def test_ignored_signal_stays_ignored(self, mock_remove, mock_kill): 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()