From 004a7e31be42b300380a79117570f911389b4db6 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Thu, 20 Aug 2026 08:13:10 +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()`, 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() From f13217857e7ff3eef24baa5d5179d7bdc4ee6525 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Wed, 23 Sep 2026 16:50:55 +0900 Subject: [PATCH 2/3] fix: keep lock release safe off the main thread and on SIGTERM - release() removes the lock file first and only restores signal handlers on the main thread, so stop() from a worker thread still writes its CSV. - With a default previous handler, raise SystemExit(128 + signum) (or KeyboardInterrupt for SIGINT) instead of re-sending the signal, so finally / __exit__ still write the final emissions. - Treat a None previous handler (installed from C) as SIG_DFL. Co-Authored-By: Claude Opus 5.5 (1M context) --- codecarbon/lock.py | 31 +++++++++++++++++----- tests/test_lock.py | 65 +++++++++++++++++++++++++++++++++++----------- 2 files changed, 74 insertions(+), 22 deletions(-) diff --git a/codecarbon/lock.py b/codecarbon/lock.py index 68ba547f9..8258fca18 100644 --- a/codecarbon/lock.py +++ b/codecarbon/lock.py @@ -39,8 +39,13 @@ def _handle_exit(self, signum, frame): self.release() # also restores the previous handlers if callable(previous): return previous(signum, frame) - if previous == signal.SIG_DFL: - os.kill(os.getpid(), signum) + if previous == signal.SIG_IGN: + return None + # Default disposition: raise instead of re-sending the signal, so that + # `finally` / `__exit__` still run and the tracker writes its last row. + if signum == signal.SIGINT: + raise KeyboardInterrupt + raise SystemExit(128 + signum) def acquire(self): """Creates a lock file and ensures it's the only instance running.""" @@ -68,11 +73,6 @@ def release(self): """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 @@ -83,3 +83,20 @@ def release(self): logger.debug(f"Error: {e}") if e.errno != errno.ENOENT: raise + finally: + self._restore_signal_handlers() + + def _restore_signal_handlers(self): + # signal.signal() only works on the main thread. Elsewhere, keep our + # handlers: a later signal still releases (a no-op) and forwards. + if threading.current_thread() is not threading.main_thread(): + return + while self._previous_handlers: + sig, handler = self._previous_handlers.popitem() + try: + # Only restore if nobody installed another handler after us. + if signal.getsignal(sig) == self._handle_exit: + # None: the handler was installed from C, use the default. + signal.signal(sig, signal.SIG_DFL if handler is None else handler) + except (ValueError, TypeError, OSError) as e: + logger.debug(f"Could not restore handler for signal {sig}: {e}") diff --git a/tests/test_lock.py b/tests/test_lock.py index 46fba39bd..0b6891e44 100644 --- a/tests/test_lock.py +++ b/tests/test_lock.py @@ -1,5 +1,4 @@ import atexit -import os import signal import threading import unittest @@ -142,44 +141,80 @@ def test_signal_is_forwarded_to_previous_handler(self, mock_file, mock_remove): @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. + def test_default_disposition_exits_with_cleanup(self, mock_file, mock_remove): + # No handler installed by the host application: exit, but through + # SystemExit so `finally` / `__exit__` still write the final emissions. signal.signal(signal.SIGTERM, signal.SIG_DFL) lock = Lock() lock.acquire() - signal.raise_signal(signal.SIGTERM) + with self.assertRaises(SystemExit) as ctx: + 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.assertEqual(ctx.exception.code, 128 + signal.SIGTERM) + # The lock was released and the default disposition put back. 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): + def test_ignored_signal_stays_ignored(self, mock_file, mock_remove): signal.signal(signal.SIGTERM, signal.SIG_IGN) lock = Lock() lock.acquire() + # The host application asked to ignore SIGTERM: release the lock, but do + # not terminate on its behalf (no SystemExit). 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") + @patch("codecarbon.lock.open", new_callable=mock_open) + def test_release_from_worker_thread_removes_lock_file(self, mock_file, mock_remove): + # stop() is often called from a callback thread: signal.signal() would + # raise there, so release() must still remove the lock file. + lock = Lock() + lock.acquire() + errors = [] + + def release(): + try: + lock.release() + except Exception as e: # pragma: no cover - the failure being tested + errors.append(e) + + worker = threading.Thread(target=release) + worker.start() + worker.join() + + self.assertEqual(errors, []) + mock_remove.assert_called_once_with(LOCKFILE) + # Handlers could not be restored off the main thread: still ours. + self.assertEqual(signal.getsignal(signal.SIGTERM), lock._handle_exit) + lock.release() + self.assertIs( + signal.getsignal(signal.SIGTERM), self.original_handlers[signal.SIGTERM] + ) + + @patch("codecarbon.lock.os.remove") + @patch("codecarbon.lock.open", new_callable=mock_open) + def test_handler_installed_from_c_is_restored_as_default( + self, mock_file, mock_remove + ): + # signal.getsignal() returns None for handlers not installed from Python. + lock = Lock() + lock.acquire() + lock._previous_handlers[signal.SIGTERM] = None + lock.release() + self.assertIs(signal.getsignal(signal.SIGTERM), signal.SIG_DFL) + @patch("codecarbon.lock.os.remove") def test_release_from_within_the_critical_section_does_not_deadlock( self, mock_remove From def9badfee9570de58a62a11faf94b9c3a05fed2 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Sun, 27 Sep 2026 15:49:24 +0900 Subject: [PATCH 3/3] test: cover SIGINT default, non-ENOENT release error and failed handler restore --- tests/test_lock.py | 49 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/tests/test_lock.py b/tests/test_lock.py index 0b6891e44..bdfffd500 100644 --- a/tests/test_lock.py +++ b/tests/test_lock.py @@ -1,4 +1,5 @@ import atexit +import errno import signal import threading import unittest @@ -239,6 +240,54 @@ def hold_then_release(): worker.start() assert done.wait(timeout=5), "release() deadlocked on its own thread lock" + @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_default_disposition_sigint_raises_keyboard_interrupt( + self, mock_file, mock_remove + ): + # SIGINT's default disposition raises KeyboardInterrupt rather than + # SystemExit, matching what an uncaught Ctrl+C would raise. + signal.signal(signal.SIGINT, signal.SIG_DFL) + lock = Lock() + lock.acquire() + + with self.assertRaises(KeyboardInterrupt): + signal.raise_signal(signal.SIGINT) + + self.assertTrue(mock_remove.called) + self.assertIs(signal.getsignal(signal.SIGINT), signal.SIG_DFL) + + @patch("codecarbon.lock.os.remove") + @patch("codecarbon.lock.open", new_callable=mock_open) + def test_release_reraises_non_enoent_oserror(self, mock_file, mock_remove): + # Only a missing lock file (ENOENT) is swallowed; any other OSError + # (e.g. a permissions issue) must still surface to the caller. + mock_remove.side_effect = OSError(errno.EACCES, "Permission denied") + lock = Lock() + lock.acquire() + with self.assertRaises(OSError): + lock.release() + + @patch("codecarbon.lock.os.remove") + @patch("codecarbon.lock.open", new_callable=mock_open) + def test_restore_signal_handlers_swallows_restore_errors( + self, mock_file, mock_remove + ): + # If putting a previous handler back fails (e.g. a signal invalid on + # this platform), release() must not raise: the lock is still removed. + lock = Lock() + lock.acquire() # installs lock._handle_exit for real via signal.signal + + with patch( + "codecarbon.lock.signal.signal", side_effect=ValueError("bad signal") + ): + lock.release() # must not raise despite the restore failing + + mock_remove.assert_called_once_with(LOCKFILE) + @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