From da08151cb1bcf9da59a37052c6526485ee6eb8e7 Mon Sep 17 00:00:00 2001 From: Ansh-699 Date: Thu, 22 Oct 2026 00:07:44 +0530 Subject: [PATCH 1/2] Release SQLite lock refs atomically with their cleanup `clean_locks` counted a key's lock references while holding the map mutex, but the finishing operation only dropped its own reference after returning. Two operations on the same key finishing at the same time could each count the other's reference and both keep the entry, leaving an unused lock cached until the key was used again. Take the reference by value, drop it under the map mutex, and remove the entry only if the map's reference is the last one left, so counting and releasing are atomic with respect to other cleanups. AI-assisted: written with help from Claude. --- src/io/sqlite_store/mod.rs | 90 +++++++++++++++++++++++++++++++++----- 1 file changed, 79 insertions(+), 11 deletions(-) diff --git a/src/io/sqlite_store/mod.rs b/src/io/sqlite_store/mod.rs index 8a2c6a0690..8848137b62 100644 --- a/src/io/sqlite_store/mod.rs +++ b/src/io/sqlite_store/mod.rs @@ -425,7 +425,7 @@ impl SqliteStoreInner { if let Err(e) = check_namespace_key_validity(primary_namespace, secondary_namespace, Some(key), "write") { - self.clean_locks(&inner_lock_ref, locking_key); + self.clean_locks(inner_lock_ref, locking_key); return Err(e.into()); } @@ -477,7 +477,7 @@ impl SqliteStoreInner { Some(key), "remove", ) { - self.clean_locks(&inner_lock_ref, locking_key); + self.clean_locks(inner_lock_ref, locking_key); return Err(e.into()); } @@ -716,22 +716,24 @@ impl SqliteStoreInner { } }; - self.clean_locks(&inner_lock_ref, locking_key); + self.clean_locks(inner_lock_ref, locking_key); res } - fn clean_locks(&self, inner_lock_ref: &Arc>, locking_key: String) { + fn clean_locks(&self, inner_lock_ref: Arc>, locking_key: String) { // If there no arcs in use elsewhere, this means that there are no in-flight writes. We can remove the map entry - // to prevent leaking memory. The two arcs that are expected are the one in the map and the one held here in - // inner_lock_ref. The outer lock is obtained first, to avoid a new arc being cloned after we've already - // counted. + // to prevent leaking memory. The outer lock is obtained first, to avoid a new arc being cloned after we've already + // counted. Our own arc is dropped before counting, so two operations finishing at the same time can't each count + // the other's arc and both keep the entry. let mut outer_lock = self.write_version_locks.lock().expect("lock"); + debug_assert!( + Arc::strong_count(&inner_lock_ref) >= 2, + "Unexpected SqliteStore strong count" + ); + drop(inner_lock_ref); - let strong_count = Arc::strong_count(&inner_lock_ref); - debug_assert!(strong_count >= 2, "Unexpected SqliteStore strong count"); - - if strong_count == 2 { + if outer_lock.get(&locking_key).is_some_and(|lock| Arc::strong_count(lock) == 1) { outer_lock.remove(&locking_key); } } @@ -946,6 +948,72 @@ mod tests { } } + #[test] + fn reclaims_lock_after_overlapping_cleanups() { + let mut temp_path = random_storage_path(); + temp_path.push("reclaims_lock_after_overlapping_cleanups"); + let store = SqliteStore::new( + temp_path, + Some("test_db".to_string()), + Some("test_table".to_string()), + ) + .unwrap(); + let locking_key = store.build_locking_key("namespace", "", "key"); + + // Two overlapping operations on the same key reach their final cleanup while both still hold + // a reference. + let (first_lock_ref, _) = store.get_new_version_and_lock_ref(locking_key.clone()); + let (second_lock_ref, _) = store.get_new_version_and_lock_ref(locking_key.clone()); + store.inner.clean_locks(first_lock_ref, locking_key.clone()); + store.inner.clean_locks(second_lock_ref, locking_key); + + assert!(store.inner.write_version_locks.lock().unwrap().is_empty()); + } + + #[test] + fn new_writer_racing_cleanup_shares_live_lock() { + let mut temp_path = random_storage_path(); + temp_path.push("new_writer_racing_cleanup_shares_live_lock"); + let store = SqliteStore::new( + temp_path, + Some("test_db".to_string()), + Some("test_table".to_string()), + ) + .unwrap(); + let locking_key = store.build_locking_key("namespace", "", "key"); + + // A new writer takes its lock reference just before another write on the same key cleans up. + let (finishing_lock_ref, finishing_version) = + store.get_new_version_and_lock_ref(locking_key.clone()); + let (new_lock_ref, new_version) = store.get_new_version_and_lock_ref(locking_key.clone()); + store + .inner + .execute_locked_write( + finishing_lock_ref, + locking_key.clone(), + finishing_version, + || Ok(()), + ) + .unwrap(); + + // The cleanup must keep the entry, so the new writer still holds the only lock for the key and + // sees the version written before it. + assert!(store + .inner + .write_version_locks + .lock() + .unwrap() + .get(&locking_key) + .is_some_and(|lock| Arc::ptr_eq(lock, &new_lock_ref))); + assert_eq!(*new_lock_ref.lock().unwrap(), finishing_version); + + store + .inner + .execute_locked_write(new_lock_ref, locking_key, new_version, || Ok(())) + .unwrap(); + assert!(store.inner.write_version_locks.lock().unwrap().is_empty()); + } + #[tokio::test] async fn read_write_remove_list_persist() { let mut temp_path = random_storage_path(); From bd0f9939e69d270d66b9f1b412269ce726c1bb27 Mon Sep 17 00:00:00 2001 From: Ansh-699 Date: Sat, 24 Oct 2026 14:08:50 +0530 Subject: [PATCH 2/2] Release Postgres and VSS lock refs atomically with their cleanup `PostgresStore` and `VssStore` share `SqliteStore`'s lock cleanup and have the same race: two operations on the same key finishing at the same time could each count the other's reference and both keep the entry. Apply the same fix: take the reference by value, drop it under the map mutex, and remove the entry only if the map's reference is the last one left. AI-assisted: written with help from Claude. --- src/io/postgres_store/mod.rs | 41 ++++++++++++++++++++++++++---------- src/io/vss_store.rs | 32 ++++++++++++++++++++-------- 2 files changed, 53 insertions(+), 20 deletions(-) diff --git a/src/io/postgres_store/mod.rs b/src/io/postgres_store/mod.rs index a80fec2952..c3d28f8aac 100644 --- a/src/io/postgres_store/mod.rs +++ b/src/io/postgres_store/mod.rs @@ -857,7 +857,7 @@ impl PostgresStoreInner { if let Err(e) = check_namespace_key_validity(primary_namespace, secondary_namespace, Some(key), "write") { - self.clean_locks(&inner_lock_ref, locking_key); + self.clean_locks(inner_lock_ref, locking_key); return Err(e.into()); } @@ -925,7 +925,7 @@ impl PostgresStoreInner { Some(key), "remove", ) { - self.clean_locks(&inner_lock_ref, locking_key); + self.clean_locks(inner_lock_ref, locking_key); return Err(e.into()); } @@ -1131,22 +1131,24 @@ impl PostgresStoreInner { } }; - self.clean_locks(&inner_lock_ref, locking_key); + self.clean_locks(inner_lock_ref, locking_key); res } - fn clean_locks(&self, inner_lock_ref: &Arc>, locking_key: String) { + fn clean_locks(&self, inner_lock_ref: Arc>, locking_key: String) { // If there are no arcs in use elsewhere, this means that there are no in-flight writes. We can remove the map - // entry to prevent leaking memory. The two arcs that are expected are the one in the map and the one held here - // in inner_lock_ref. The outer lock is obtained first, to avoid a new arc being cloned after we've already - // counted. + // entry to prevent leaking memory. The outer lock is obtained first, to avoid a new arc being cloned after we've + // already counted. Our own arc is dropped before counting, so two operations finishing at the same time can't + // each count the other's arc and both keep the entry. let mut outer_lock = self.write_version_locks.lock().unwrap(); + debug_assert!( + Arc::strong_count(&inner_lock_ref) >= 2, + "Unexpected PostgresStore strong count" + ); + drop(inner_lock_ref); - let strong_count = Arc::strong_count(inner_lock_ref); - debug_assert!(strong_count >= 2, "Unexpected PostgresStore strong count"); - - if strong_count == 2 { + if outer_lock.get(&locking_key).is_some_and(|lock| Arc::strong_count(lock) == 1) { outer_lock.remove(&locking_key); } } @@ -1806,4 +1808,21 @@ mod tests { let result = PostgresStore::build_tls_connector(Some("not-a-valid-pem".to_string())); assert!(result.is_err()); } + + #[tokio::test(flavor = "multi_thread")] + async fn test_postgres_reclaims_lock_after_overlapping_cleanups() { + let store = create_test_store("test_pg_reclaims_lock").await; + let locking_key = store.build_locking_key("namespace", "", "key"); + + // Two overlapping operations on the same key reach their final cleanup while both still hold + // a reference. The first cleanup must keep the entry for the second, which must remove it. + let (first_lock_ref, _) = store.get_new_version_and_lock_ref(locking_key.clone()); + let (second_lock_ref, _) = store.get_new_version_and_lock_ref(locking_key.clone()); + store.inner.clean_locks(first_lock_ref, locking_key.clone()); + assert!(store.inner.write_version_locks.lock().unwrap().contains_key(&locking_key)); + store.inner.clean_locks(second_lock_ref, locking_key); + assert!(store.inner.write_version_locks.lock().unwrap().is_empty()); + + cleanup_store(&store).await; + } } diff --git a/src/io/vss_store.rs b/src/io/vss_store.rs index 6a44be02d9..df8ba60762 100644 --- a/src/io/vss_store.rs +++ b/src/io/vss_store.rs @@ -757,22 +757,21 @@ impl VssStoreInner { } }; - self.clean_locks(&inner_lock_ref, locking_key); + self.clean_locks(inner_lock_ref, locking_key); res } - fn clean_locks(&self, inner_lock_ref: &Arc>, locking_key: String) { + fn clean_locks(&self, inner_lock_ref: Arc>, locking_key: String) { // If there no arcs in use elsewhere, this means that there are no in-flight writes. We can remove the map entry - // to prevent leaking memory. The two arcs that are expected are the one in the map and the one held here in - // inner_lock_ref. The outer lock is obtained first, to avoid a new arc being cloned after we've already - // counted. + // to prevent leaking memory. The outer lock is obtained first, to avoid a new arc being cloned after we've already + // counted. Our own arc is dropped before counting, so two operations finishing at the same time can't each count + // the other's arc and both keep the entry. let mut outer_lock = self.locks.lock().expect("lock"); + debug_assert!(Arc::strong_count(&inner_lock_ref) >= 2, "Unexpected VssStore strong count"); + drop(inner_lock_ref); - let strong_count = Arc::strong_count(&inner_lock_ref); - debug_assert!(strong_count >= 2, "Unexpected VssStore strong count"); - - if strong_count == 2 { + if outer_lock.get(&locking_key).is_some_and(|lock| Arc::strong_count(lock) == 1) { outer_lock.remove(&locking_key); } } @@ -1359,4 +1358,19 @@ mod tests { assert_eq!(response.keys.len(), 1); assert!(response.keys.contains(&"key_3".to_string())); } + + #[tokio::test] + async fn vss_reclaims_lock_after_overlapping_cleanups() { + let store = build_vss_store(); + let locking_key = store.build_locking_key("namespace", "", "key"); + + // Two overlapping operations on the same key reach their final cleanup while both still hold + // a reference. The first cleanup must keep the entry for the second, which must remove it. + let (first_lock_ref, _) = store.get_new_version_and_lock_ref(locking_key.clone()); + let (second_lock_ref, _) = store.get_new_version_and_lock_ref(locking_key.clone()); + store.inner.clean_locks(first_lock_ref, locking_key.clone()); + assert!(store.inner.locks.lock().unwrap().contains_key(&locking_key)); + store.inner.clean_locks(second_lock_ref, locking_key); + assert!(store.inner.locks.lock().unwrap().is_empty()); + } }