From 942770cd6a4f60ba3b60b0ea15c36f27a7456a8e Mon Sep 17 00:00:00 2001 From: Shaun Cutts Date: Wed, 3 Jun 2026 19:28:05 -0400 Subject: [PATCH] Fix CallStack.register orphan lock if exception fires before try/finally CallStack.register currently calls self._lock.acquire() and self._stack.append(...) BEFORE entering the try/finally that releases the lock. Any exception fired between the acquire and the try-block -- for example raised by CallFrame.__init__, an OOM, or an asynchronously-injected interrupt -- orphans the lock; subsequent threads queued on it block forever. Restructure with `with self._lock:` so the release is unconditional. No behavior change on the happy path; strictly stronger under failure. Adds a regression test in numba_cuda/numba/cuda/tests/nocuda that stubs CallFrame to raise on construction and asserts the lock is still free afterwards (verified to fail on the previous code). --- .../test_callstack_register_orphan_lock.py | 114 ++++++++++++++++++ numba_cuda/numba/cuda/typing/context.py | 16 +-- 2 files changed, 123 insertions(+), 7 deletions(-) create mode 100644 numba_cuda/numba/cuda/tests/nocuda/test_callstack_register_orphan_lock.py diff --git a/numba_cuda/numba/cuda/tests/nocuda/test_callstack_register_orphan_lock.py b/numba_cuda/numba/cuda/tests/nocuda/test_callstack_register_orphan_lock.py new file mode 100644 index 000000000..e2cace37c --- /dev/null +++ b/numba_cuda/numba/cuda/tests/nocuda/test_callstack_register_orphan_lock.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-2-Clause + +"""Regression test for CallStack.register orphan-lock defect. + +If an exception fires after `CallStack._lock` is acquired but before +the surrounding `try/finally` is entered (e.g. raised by +`CallFrame.__init__`, an OOM, or an asynchronously-injected interrupt), +the lock must still be released. A bug in the previous implementation +called `self._lock.acquire()` outside the try-block, so any such +exception would orphan the lock and cause subsequent registrations on +the same `CallStack` to block forever. +""" + +import threading +import unittest +from unittest import mock + +from numba.cuda.typing import context as typing_context + + +class _BoomError(RuntimeError): + pass + + +class _FakeFuncId: + """Minimal stand-in for FunctionIdentity used by CallStack.register.""" + + def __init__(self): + # `func` only needs to be something `self.match` can compare by + # identity. A bare object suffices: an empty CallStack will not + # match anything regardless. + self.func = object() + + +class TestCallStackRegisterOrphanLock(unittest.TestCase): + def test_lock_released_when_callframe_construction_raises(self): + """register() must release its lock if CallFrame() raises.""" + stack = typing_context.CallStack() + func_id = _FakeFuncId() + + with mock.patch.object( + typing_context, + "CallFrame", + side_effect=_BoomError("simulated failure"), + ): + with self.assertRaises(_BoomError): + with stack.register( + target=None, + typeinfer=None, + func_id=func_id, + args=(), + ): + # Should never enter the body: CallFrame() raised first. + self.fail("register body should not have executed") + + # The lock must be free. Use a non-blocking acquire with a + # background thread to prove no thread holds it (RLock would + # let the same thread re-enter, masking the bug). + acquired = [False] + + def _try_acquire(): + acquired[0] = stack._lock.acquire(blocking=False) + if acquired[0]: + stack._lock.release() + + thread = threading.Thread(target=_try_acquire) + thread.start() + thread.join(timeout=5.0) + self.assertFalse( + thread.is_alive(), "background thread blocked on orphaned lock" + ) + self.assertTrue( + acquired[0], + "CallStack._lock was orphaned after CallFrame() raised", + ) + # Stack must remain empty since the append never happened. + self.assertEqual(len(stack), 0) + + def test_lock_released_on_normal_exit(self): + """Sanity check: happy path still releases the lock and pops.""" + stack = typing_context.CallStack() + func_id = _FakeFuncId() + + with mock.patch.object( + typing_context, "CallFrame", autospec=False + ) as fake_frame: + fake_frame.return_value = mock.MagicMock(func_id=func_id, args=()) + with stack.register( + target=None, + typeinfer=None, + func_id=func_id, + args=(), + ): + self.assertEqual(len(stack), 1) + + self.assertEqual(len(stack), 0) + # Lock free in a fresh thread. + acquired = [False] + + def _try_acquire(): + acquired[0] = stack._lock.acquire(blocking=False) + if acquired[0]: + stack._lock.release() + + thread = threading.Thread(target=_try_acquire) + thread.start() + thread.join(timeout=5.0) + self.assertFalse(thread.is_alive()) + self.assertTrue(acquired[0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/numba_cuda/numba/cuda/typing/context.py b/numba_cuda/numba/cuda/typing/context.py index c9d4c6832..58f998fb0 100644 --- a/numba_cuda/numba/cuda/typing/context.py +++ b/numba_cuda/numba/cuda/typing/context.py @@ -67,13 +67,15 @@ def register(self, target, typeinfer, func_id, args): if self.match(func_id.func, args): msg = "compiler re-entrant to the same function signature" raise errors.NumbaRuntimeError(msg) - self._lock.acquire() - self._stack.append(CallFrame(target, typeinfer, func_id, args)) - try: - yield - finally: - self._stack.pop() - self._lock.release() + # Use `with self._lock:` so an exception between acquire and the + # try-block (e.g. raised by CallFrame.__init__, an OOM, or an + # asynchronously-injected interrupt) cannot orphan the lock. + with self._lock: + self._stack.append(CallFrame(target, typeinfer, func_id, args)) + try: + yield + finally: + self._stack.pop() def finditer(self, py_func): """