redis_store: hold subscribed_keys write lock across receiver drop (#2353)

Co-authored-by: Marcus Eagan <marcuseagan@gmail.com>

Author: amankrx

nativelink-store/src/redis_store.rs

@@ -1327,22 +1327,24 @@ impl Drop for RedisSubscription {
             return; // Already dropped, nothing to do.
         };
         let key = receiver.borrow().clone();
-        // IMPORTANT: This must be dropped before receiver_count() is called.
-        drop(receiver);
         let Some(subscribed_keys) = self.weak_subscribed_keys.upgrade() else {
-            return; // Already dropped, nothing to do.
+            return; // Parent dropped — nothing to do.
         };
         let mut subscribed_keys = subscribed_keys.write();
-        let Some(value) = subscribed_keys.get(&key) else {
-            error!(
-                "Key {key} was not found in subscribed keys when checking if it should be removed."
+        let Some(publisher) = subscribed_keys.get(&key) else {
+            warn!(
+                %key,
+                "RedisSubscription::drop: key absent from subscribed_keys under write lock — \
+                 indicates an unexpected removal path",
             );
             return;
         };
-        // If we have no receivers, cleanup the entry from our map.
-        if value.receiver_count() == 0 {
-            subscribed_keys.remove(key);
+        // Count includes our own (still-alive) receiver. If we are the
+        // sole subscriber, remove the publisher entry.
+        if publisher.receiver_count() == 1 {
+            subscribed_keys.remove(&key);
         }
+        drop(receiver);
     }
 }
 

nativelink-store/tests/redis_store_test.rs

@@ -23,8 +23,9 @@ use nativelink_config::stores::{RedisMode, RedisSpec};
 use nativelink_error::{Code, Error, ErrorContext, ResultExt, make_err};
 use nativelink_macro::nativelink_test;
 use nativelink_redis_tester::{
-    ReadOnlyRedis, add_lua_script, add_to_response, fake_redis_sentinel_master_stream,
-    fake_redis_sentinel_stream, fake_redis_stream, make_fake_redis_with_responses,
+    ReadOnlyRedis, SubscriptionManagerNotify, add_lua_script, add_to_response,
+    fake_redis_sentinel_master_stream, fake_redis_sentinel_stream, fake_redis_stream,
+    make_fake_redis_with_responses,
 };
 use nativelink_store::cas_utils::ZERO_BYTE_DIGESTS;
 use nativelink_store::redis_store::{
@@ -36,8 +37,9 @@ use nativelink_util::common::DigestInfo;
 use nativelink_util::health_utils::HealthStatus;
 use nativelink_util::store_trait::{
     FalseValue, SchedulerCurrentVersionProvider, SchedulerIndexProvider, SchedulerStore,
-    SchedulerStoreDataProvider, SchedulerStoreDecodeTo, SchedulerStoreKeyProvider, StoreKey,
-    StoreLike, TrueValue, UploadSizeInfo,
+    SchedulerStoreDataProvider, SchedulerStoreDecodeTo, SchedulerStoreKeyProvider,
+    SchedulerSubscription, SchedulerSubscriptionManager, StoreKey, StoreLike, TrueValue,
+    UploadSizeInfo,
 };
 use pretty_assertions::assert_eq;
 use redis::{PushInfo, RedisError, Value, make_extension_error};
@@ -1734,3 +1736,128 @@ async fn test_update_data_versioned_with_expiry() {
         .await
         .expect("working update");
 }
+
+/// Test key provider that just wraps a string. Reused across the
+/// subscription regression tests below.
+#[derive(Clone)]
+struct TestSubKey(String);
+
+impl SchedulerStoreKeyProvider for TestSubKey {
+    type Versioned = FalseValue;
+    fn get_key(&self) -> StoreKey<'static> {
+        StoreKey::Str(std::borrow::Cow::Owned(self.0.clone()))
+    }
+}
+
+/// Sanity: a single subscriber that drops cleanly produces no warning
+/// and no error.
+#[nativelink_test]
+async fn redis_subscription_single_drop_is_silent() -> Result<(), Error> {
+    let (_tx, rx) = tokio::sync::mpsc::unbounded_channel();
+    let manager = RedisSubscriptionManager::new(rx);
+
+    let sub = manager.subscribe(TestSubKey("solo-key".to_string()))?;
+    drop(sub);
+    sleep(Duration::from_millis(10)).await;
+
+    assert!(
+        !logs_contain("key absent from subscribed_keys under write lock"),
+        "single-subscriber drop unexpectedly logged the absence warning",
+    );
+    assert!(!logs_contain("ERROR"));
+
+    drop(manager);
+    Ok(())
+}
+
+#[nativelink_test]
+async fn redis_subscription_drop_one_of_two_keeps_publisher() -> Result<(), Error> {
+    let (_tx, rx) = tokio::sync::mpsc::unbounded_channel();
+    let manager = RedisSubscriptionManager::new(rx);
+
+    let key = "shared-key";
+    let sub_a = manager.subscribe(TestSubKey(key.to_string()))?;
+    let mut sub_b = manager.subscribe(TestSubKey(key.to_string()))?;
+
+    // Drop the first; the second's subscription must still resolve
+    // when we notify on the same key.
+    drop(sub_a);
+
+    manager.notify_for_test(key.to_string());
+    timeout(Duration::from_secs(2), sub_b.changed())
+        .await
+        .expect("sub_b.changed() did not fire — publisher entry was dropped prematurely")?;
+
+    assert!(
+        !logs_contain("key absent from subscribed_keys under write lock"),
+        "absence warning fired during single drop with another receiver alive",
+    );
+    drop(sub_b);
+    drop(manager);
+    Ok(())
+}
+
+#[nativelink_test]
+async fn redis_subscription_concurrent_drops_no_absence_warn() -> Result<(), Error> {
+    let (_tx, rx) = tokio::sync::mpsc::unbounded_channel();
+    let manager = RedisSubscriptionManager::new(rx);
+
+    const ITERATIONS: usize = 200;
+    for i in 0..ITERATIONS {
+        let key = format!("race-key-{i}");
+        let sub_a = manager.subscribe(TestSubKey(key.clone()))?;
+        let sub_b = manager.subscribe(TestSubKey(key.clone()))?;
+
+        // `spawn_blocking` puts each Drop on its own thread; with the
+        // pre-fix code's "drop receiver, then take lock" sequence,
+        // both threads can race into the lock with already-decremented
+        // counts. With the post-fix "take lock, then evaluate, then
+        // drop receiver" sequence, the lock serialises the decision
+        // and the warning never fires.
+        let h_a = tokio::task::spawn_blocking(move || drop(sub_a));
+        let h_b = tokio::task::spawn_blocking(move || drop(sub_b));
+        h_a.await.unwrap();
+        h_b.await.unwrap();
+    }
+    sleep(Duration::from_millis(50)).await;
+
+    assert!(
+        !logs_contain("key absent from subscribed_keys under write lock"),
+        "concurrent drops produced the absence warning at least once across {ITERATIONS} \
+         iterations — the Drop ordering regressed",
+    );
+    assert!(!logs_contain("ERROR"));
+
+    drop(manager);
+    Ok(())
+}
+
+#[nativelink_test]
+async fn redis_subscription_resubscribe_after_drop_creates_fresh_publisher() -> Result<(), Error> {
+    let (_tx, rx) = tokio::sync::mpsc::unbounded_channel();
+    let manager = RedisSubscriptionManager::new(rx);
+
+    let key = "cycle-key";
+    let sub_a = manager.subscribe(TestSubKey(key.to_string()))?;
+    let sub_b = manager.subscribe(TestSubKey(key.to_string()))?;
+    drop(sub_a);
+    drop(sub_b);
+
+    // Re-subscribe to the same key. If the previous drops left the
+    // map in an inconsistent state (stale publisher kept, or a
+    // partially-deconstructed entry), this either reuses a dead
+    // publisher (changed() never fires) or panics inside the
+    // patricia map.
+    let mut sub_c = manager.subscribe(TestSubKey(key.to_string()))?;
+    manager.notify_for_test(key.to_string());
+    timeout(Duration::from_secs(2), sub_c.changed())
+        .await
+        .expect("re-subscribe after drops produced a dead publisher")?;
+
+    assert!(!logs_contain(
+        "key absent from subscribed_keys under write lock"
+    ));
+    drop(sub_c);
+    drop(manager);
+    Ok(())
+}