[Ansh-699]

Fixes #1144.

clean_locks counted a key's lock references under the map mutex, but each operation dropped its own reference only after returning. Two operations on the same key finishing together could each count the other, both keep the entry, and leave an unused lock cached.

clean_locks now takes the reference by value and drops it before releasing the mutex.

Tests:

• reclaims_lock_after_overlapping_cleanups fails on main (with the old &Arc signature) and passes with this change.
• new_writer_racing_cleanup_shares_live_lock checks that a new writer racing cleanup keeps sharing the key's only lock.

PostgresStore and VssStore use the same pattern. Happy to follow up there if you'd like.

AI-assisted: written with help from Claude.

[ldk-reviews-bot]

I've assigned @tnull as a reviewer!
I'll wait for their review and will help manage the review process.
Once they submit their review, I'll check if a second reviewer would be helpful.

[shubhamkmr04]

This drop looks redundant, but the fix depends on it: if it's removed, the code still compiles, inner_lock_ref is dropped after outer_lock (parameters are dropped after locals), and the race comes back. Both new tests still pass in that case, since they run the cleanups one after another.

A test with two threads catches it. With the drop removed (or on main), this fails on every run for me, and with the fix it passes:

#[test]
fn concurrently_finishing_operations_remove_their_lock_entries() {
    let store = SqliteStore::new(random_storage_path(), None, None).unwrap();
    let locking_keys: Vec<String> =
        (0..100).map(|i| store.build_locking_key("test", "", &format!("key{}", i))).collect();

    // Two threads each run an operation on every key, so that on each key both are in flight and finish at about
    // the same time. An operation that counted the other's reference before releasing its own would leave the
    // key's lock entry behind.
    for _ in 0..200 {
        let barrier = Arc::new(std::sync::Barrier::new(2));
        let threads: Vec<_> = (0..2)
            .map(|_| {
                let operations: Vec<_> = locking_keys
                    .iter()
                    .map(|locking_key| {
                        let (lock_ref, version) =
                            store.get_new_version_and_lock_ref(locking_key.clone());
                        (lock_ref, locking_key.clone(), version)
                    })
                    .collect();
                let inner = Arc::clone(&store.inner);
                let barrier = Arc::clone(&barrier);
                std::thread::spawn(move || {
                    barrier.wait();
                    for (lock_ref, locking_key, version) in operations {
                        inner
                            .execute_locked_write(lock_ref, locking_key, version, || Ok(()))
                            .unwrap();
                    }
                })
            })
            .collect();
        for thread in threads {
            thread.join().unwrap();
        }

        assert!(
            store.inner.write_version_locks.lock().unwrap().is_empty(),
            "lock entries outlived the operations on their keys"
        );
    }
}

Another option is to drop first and then remove the entry if the map's Arc is the only one left (strong_count == 1), so the order is part of the logic rather than a trailing drop.

[Ansh-699]

Good catch, thanks. Went with your second option: clean_locks now drops our arc first and removes the entry only if the map's Arc is the only one left. Removing the drop now makes both existing tests fail every time.

I left the threaded test out. On main it failed for me in about half the runs (11 of 21), so the two sequential tests are the more reliable check now. Happy to add it if you'd still like it.

[shubhamkmr04]

Heads-up: #1186 adds self.clean_locks(&inner_lock_ref, locking_key) calls in write_internal/remove_internal. The two PRs merge without a conflict, but together they don't compile (expected Arc<Mutex<u64>>, found &Arc<Mutex<u64>>), so whichever lands second needs to pass the Arc by value there.

[Ansh-699]

Thanks, noted. #1186 is still open, so whichever lands second will pass the Arc by value there. Happy to rebase on top of it if that's easier.