Handle correctly subscription messages (TraceMachina#2201)

Author: palfrey

nativelink-store/src/redis_store.rs

@@ -337,9 +337,7 @@ where
     /// A manager for subscriptions to keys in Redis.
     subscription_manager: tokio::sync::OnceCell<Arc<RedisSubscriptionManager>>,
 
-    /// Channel for getting subscription messages. Only used by cluster mode where
-    /// the sender is connected at construction time. For standard mode, this is
-    /// None and created on demand in `subscription_manager()`.
+    /// Channel for getting subscription messages
     subscriber_channel: Mutex<Option<UnboundedReceiver<PushInfo>>>,
 
     /// Permits to limit inflight Redis requests. Technically only
@@ -1366,11 +1364,22 @@ impl RedisSubscriptionManager {
                                 },
                                 maybe_push_info = local_subscriber_channel.next() => {
                                     if let Some(push_info) = maybe_push_info {
-                                        if push_info.data.len() != 1 {
-                                            error!(?push_info, "Expected exactly one message on subscriber_channel");
+                                        match push_info.kind {
+                                            redis::PushKind::PMessage => {},
+                                            redis::PushKind::PSubscribe => {
+                                                trace!(?push_info, "PSubscribe, ignore");
+                                                continue;
+                                            }
+                                            _ => {
+                                                warn!(?push_info, "Other push_info message, discarded");
+                                                continue;
+                                            },
+                                        }
+                                        if push_info.data.len() != 3 {
+                                            error!(?push_info, "Expected exactly 3 values on subscriber channel (pattern, channel, value)");
                                             continue;
                                         }
-                                        match push_info.data.first().unwrap() {
+                                        match push_info.data.last().unwrap() {
                                             Value::SimpleString(s) => {
                                                 s.clone()
                                             }
@@ -1388,6 +1397,7 @@ impl RedisSubscriptionManager {
                                     }
                                 }
                             };
+                            trace!(key, "New subscription manager key");
                             let Some(subscribed_keys) = subscribed_keys_weak.upgrade() else {
                                 warn!(
                                     "It appears our parent has been dropped, exiting RedisSubscriptionManager spawn"

nativelink-store/tests/redis_store_test.rs

@@ -40,9 +40,9 @@ use nativelink_util::store_trait::{
     StoreLike, TrueValue, UploadSizeInfo,
 };
 use pretty_assertions::assert_eq;
-use redis::{RedisError, Value};
+use redis::{PushInfo, RedisError, Value};
 use redis_test::{MockCmd, MockRedisConnection};
-use tokio::time::sleep;
+use tokio::time::{sleep, timeout};
 use tracing::{Instrument, info, info_span};
 
 const VALID_HASH1: &str = "3031323334353637383961626364656630303030303030303030303030303030";
@@ -1338,3 +1338,49 @@ async fn no_items_from_none_subscription_channel() -> Result<(), Error> {
 
     Ok(())
 }
+
+#[nativelink_test]
+async fn send_messages_to_subscription_channel() -> Result<(), Error> {
+    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+    let subscription_manager = RedisSubscriptionManager::new(rx);
+
+    tx.send(PushInfo {
+        kind: redis::PushKind::PSubscribe,
+        data: vec![
+            // Pattern
+            Value::BulkString("scheduler_key_change".into()),
+            // Subscribe count
+            Value::Int(1),
+        ],
+    })
+    .unwrap();
+    tx.send(PushInfo {
+        kind: redis::PushKind::PMessage,
+        data: vec![
+            // First is the pattern
+            Value::BulkString("scheduler_key_change".into()),
+            // Second is the matching channel. Which in this case is the same as the pattern.
+            Value::BulkString("scheduler_key_change".into()),
+            // And then the actual message
+            Value::BulkString("demo-key".into()),
+        ],
+    })
+    .unwrap();
+
+    timeout(Duration::from_secs(5), async {
+        loop {
+            assert!(!logs_contain("ERROR"));
+            if logs_contain("New subscription manager key key=\"demo-key\"") {
+                break;
+            }
+            sleep(Duration::from_millis(100)).await;
+        }
+    })
+    .await
+    .unwrap();
+
+    // Because otherwise it gets dropped immediately, and we need it to live to do things
+    drop(subscription_manager);
+
+    Ok(())
+}