JAX PR #40556: Why an Explain Exception Left a Cache Entry Stuck In Progress
JAX PR #40556 fixes a cache miss that could publish an in-progress entry before the code responsible for cleaning it up existed.
The failure began in WeakrefLRUCacheBase::Call().
On a miss, the cache inserted a new CacheEntry whose completed notification was still unset. Other calls used that entry to learn whether they should compute the value, reject a recursive lookup, or wait for another thread.
The function intended to guarantee two things on every exit:
an inserted entry eventually signals completion
and
local shared owners are released while the cache lock is heldThose guarantees were implemented with absl::Cleanup, but the cleanup objects were constructed after the initial critical section.
If the optional explain factory or key collection raised after the entry had been inserted but before the critical section ended, control never reached the cleanup declarations. The cache retained an entry that still looked actively computed.
A later call on the same thread could receive a false RecursionError. A call on another thread could wait for a notification that would never arrive.
The patch moves both cleanups above the critical section so they cover the publication of the entry itself.
A cache entry represented work in progress
This weak-reference LRU cache does not store only completed values.
Its CacheEntry also contains coordination state:
completed
a one-shot notification indicating that computation ended
result
valid only after completion
thread_id
the thread responsible for computing the entry
LRU links
intrusive links into the cache eviction list
weak and strong keys
identities used by the nested maps and reverse indexThe cache uses a two-level structure. A weak-key map points to an inner strong-key cache, and the inner cache owns shared pointers to entries.
Publishing a newly allocated entry lets concurrent callers converge on one computation instead of invoking the cached function repeatedly.
That publication creates an obligation. Once another caller can observe the entry, every exit path must convert it from “being computed” into a terminal state or remove it safely.
The normal miss path had several stages
A simplified successful miss looked like this:
acquire the cache object lock
↓
look up the weak and strong keys
↓
insert a new CacheEntry
↓
mark the current call as the inserting call
↓
prepare optional explanation data
↓
release the lock
↓
call the cached Python function
↓
publish the result
↓
notify completionThe entry has to become visible before the expensive user function runs. Otherwise another thread cannot discover that the same key is already being computed.
The dangerous interval begins at insertion and ends when exception cleanup is active.
Before this PR, that interval included the remaining work inside the first locked block.
The cleanup was correct only after execution reached it
The old code declared destroy_entry_under_lock after the first critical section.
It then declared a separate notification cleanup inside the later inserted branch, immediately before invoking the optional explainer and cached function.
That protected exceptions occurring during the main computation.
It did not protect exceptions thrown earlier while the miss was still being prepared.
The commit identifies two examples:
calling the optional explain factory
collecting the current cache keys through GetKeys()Both could occur after the entry had been inserted.
If either operation raised, C++ stack unwinding left the critical section and returned the Python exception. Since execution had not reached the later cleanup declarations, no cleanup object existed to notify the entry.
RAII cannot protect code that runs before the guard is constructed.
The stale entry looked like recursion on the same thread
Each new entry records the ID of the thread computing it.
When a caller finds an incomplete entry, the cache distinguishes two cases.
If the recorded thread ID matches the current thread, waiting would deadlock: the thread would be waiting for itself to finish the same computation. JAX therefore raises a RecursionError for a genuine recursive invocation.
If the first call failed during miss preparation, however, its computation was no longer running. The stale entry still contained:
completed = false
thread_id = the original threadA retry using the same key on that thread matched the recursion detector even though no recursion was taking place.
The detector behaved consistently with the state it received. The exception path had failed to update that state.
Another thread could wait forever
A different thread follows the other incomplete-entry path.
It releases the Python GIL and calls:
entry->completed.WaitForNotification();That is the correct behavior while another thread is genuinely computing the value.
It becomes a hang when the inserting call has already exited through an exception and nobody retains the responsibility to call Notify().
The first failure is visible and understandable: the explain callback raised.
The second failure can appear later, in a different caller, as an unrelated blocked cache lookup.
This is why the notification is part of the exception-safety contract rather than merely an optimization for duplicate work.
Entry destruction also had a lock requirement
The incomplete notification was only one problem.
CacheEntry participates directly in shared cache structures. Its destructor can unlink the entry from the intrusive LRU list and remove it from the reverse weak-reference index.
Those structures are protected by the cache object lock.
The function holds local shared_ptr owners named entry and cache_ptr. Usually the cache maps also retain ownership. A concurrent Clear() can remove those map references while the current call is still unwinding.
In that race, the local pointer may become the final owner. Resetting it can immediately run ~CacheEntry().
The old destroy_entry_under_lock cleanup was designed for exactly this case, but an early exception could occur before that cleanup was declared. Normal destruction of the local shared_ptr would then happen after the nested lock guard had already been released.
The bug therefore crossed two contracts:
completion must always be signaled
and
cache-structure destruction must occur under the object lockThe patch registers both obligations before publication
The merged function now declares destroy_entry_under_lock before entering the first critical section.
Its cleanup body first checks whether either local pointer was acquired. If so, it takes the object lock and resets both pointers there.
The notification cleanup is declared next:
absl::Cleanup notify = [&inserted, &entry]() {
if (inserted) {
entry->completed.Notify();
}
};Only the call that inserted the entry sends the signal. Cache hits and callers waiting on an existing entry do not claim ownership of another thread’s completion event.
The function then enters the critical section and may publish an entry.
Every subsequent return or exception now passes through both cleanup objects.
Cleanup order keeps the entry alive for notification
C++ destroys automatic objects in reverse declaration order.
The pointer-reset cleanup is declared first. The notification cleanup is declared second.
At function exit, the order is therefore:
notify the inserted entry
↓
acquire the cache object lock
↓
reset entry and cache_ptrThe entry remains alive while Notify() uses it.
Only afterward are the local shared owners released under the required lock.
This ordering applies to normal returns and exception unwinding. The patch does not need separate catch blocks for each Python or C++ operation that might fail.
The regression test reproduces the stale state
The new test creates an explain factory controlled by a boolean.
On the first call, the factory raises:
RuntimeError("explain factory error")The test verifies that the original exception still reaches the caller.
It then disables the exception and invokes the cache again with the same weak key.
The second call must compute and return 42.
This directly protects the same-thread symptom. Under the old ordering, the first call could leave the incomplete entry carrying the current thread ID, and the retry could be classified as recursive.
The test also confirms that the patch does not swallow the original error or cache a fabricated result.
The test does not exercise every stated failure mode
The test comment says subsequent calls on the same thread or another thread must avoid RecursionError or deadlock.
The code directly performs only the same-thread retry.
It does not start a second thread, place a timeout around a waiter, or force an exception while a peer is already blocked on the notification.
It also does not combine the explain exception with a concurrent cache_clear() to make the local shared_ptr the final owner and observe the destructor-lock boundary.
The test file contains broader multithreaded and cache-clear stress coverage, but the new regression does not join those conditions into this precise failure.
The notification and locked-reset behavior are visible in the implementation. Their specific concurrent failure modes remain less directly tested than the same-thread RecursionError case.
The public review record is limited
This pull request was imported into the public JAX repository through Copybara.
The public GitHub thread contains the final description, two-file patch, test, and merged commit, but no review comments or submitted reviews.
That is enough to inspect the lifetime mechanism and regression.
It does not expose the internal review discussion, the original production trigger, or a multithreaded failure log.
No performance claim is needed. The patch changes exception cleanup and coordination correctness, not the successful cache-hit algorithm.
The contract restored by the change
Exception safety begins when shared state becomes externally observable, not when the main computation starts.
The corrected sequence is:
register destruction cleanup
↓
register completion cleanup
↓
enter the critical section
↓
publish an in-progress entry
↓
perform every remaining fallible operation
↓
notify on every exit
↓
release shared owners under the lockA cleanup guard created after publication cannot protect the publication step.
JAX PR #40556 closes that interval. An exception during miss preparation can still fail the current call, but it no longer leaves future callers trapped in work that has already ended.
Sources
https://github.com/jax-ml/jax/pull/40556
https://github.com/jax-ml/jax/commit/637326002d07205e49b44a31c5b5bec8f47a3e6e