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/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(); 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()); + } }