diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index 433e07ede6..c8518f0151 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -2106,6 +2106,19 @@ impl RedisSubscriptionManager { } } }; + // The notice carries the key as Redis holds it, behind + // this store's `key_prefix`; subscriptions are kept under + // the bare store key. A notice from outside the prefix is + // another store's, sharing the channel, and wakes nothing. + let Some(key) = + key.strip_prefix(key_prefix.as_str()).map(str::to_owned) + else { + trace!( + key, + "Published key is not under this store's prefix; ignored" + ); + continue; + }; trace!(key, "New subscription manager key"); let Some(subscribed_keys) = subscribed_keys_weak.upgrade() else { warn!( @@ -2184,6 +2197,31 @@ impl SchedulerSubscriptionManager for RedisSubscriptionManager { } } +impl RedisStore +where + C: ConnectionLike + Clone + Send + Sync, + M: RedisManager + Send + Sync, +{ + /// The keys the search index for `K` covers: the schema's prefix behind + /// this store's `key_prefix`, because that is where every key this store + /// writes lands. An index over the bare schema prefix matches nothing when + /// a `key_prefix` is set, and a scheduler on it never sees a queued action. + fn index_key_prefix(&self) -> String { + format!("{}{}", self.key_prefix, K::KEY_PREFIX) + } + + /// The search index's name, under this store's `key_prefix` for the same + /// reason: two stores with different prefixes on one Redis must not share + /// an index whose PREFIX only one of them wrote. + fn index_name(&self) -> String { + format!( + "{}{}", + self.key_prefix, + get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY) + ) + } +} + impl SchedulerStore for RedisStore where C: Clone + ConnectionLike + Sync + Send + 'static, @@ -2374,13 +2412,11 @@ where .err_tip(|| { format!("In RedisStore::search_by_index_prefix::try_sanitize - {index_value:?}") })?; + let index_name = self.index_name::(); let run_ft_aggregate = |connection_manager: C| async { ft_aggregate( connection_manager, - format!( - "{}", - get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY) - ), + index_name.clone(), if index_value.is_empty() { "*".to_string() } else { @@ -2410,7 +2446,7 @@ where }); } let create_options = FtCreateOptions { - prefixes: vec![K::KEY_PREFIX.into()], + prefixes: vec![self.index_key_prefix::()], nohl: true, nofields: true, nofreqs: true, @@ -2427,10 +2463,7 @@ where // dropped (`FT.DROPINDEX`, without DD); this code recreates it. temporary: None, }; - let index = format!( - "{}", - get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY) - ); + let index = index_name.clone(); ft_create(connection_manager, index, create_options, schema).await }; @@ -2446,10 +2479,7 @@ where let (connection_manager, _connect_id) = self.connection_manager.reconnect(connect_id).await?; run_ft_aggregate(connection_manager).await.err_tip(|| { - format!( - "Error with reconnected ft_aggregate in RedisStore::search_by_index_prefix({})", - get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY), - ) + format!("Error with reconnected ft_aggregate in RedisStore::search_by_index_prefix({index_name})") }) } Err(_) => { @@ -2478,10 +2508,7 @@ where Ok(()) } else { Err(e).err_tip(|| { - format!( - "Error with ft_create in RedisStore::search_by_index_prefix({})", - get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY), - ) + format!("Error with ft_create in RedisStore::search_by_index_prefix({index_name})") }) } }); @@ -2494,21 +2521,11 @@ where // is bounded; past it the aggregate runs anyway and the next // pass re-reads. if create_result.is_ok() { - wait_for_index_ready( - connection_manager.clone(), - &format!( - "{}", - get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY) - ), - ) - .await; + wait_for_index_ready(connection_manager.clone(), &index_name.clone()).await; } let run_result = run_ft_aggregate(connection_manager).await.err_tip(|| { - format!( - "Error with second ft_aggregate in RedisStore::search_by_index_prefix({})", - get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY), - ) + format!("Error with second ft_aggregate in RedisStore::search_by_index_prefix({index_name})") }); // Creating the index will race which is ok. If it fails to create, we only @@ -2616,10 +2633,7 @@ where .err_tip(|| { format!("In RedisStore::count_by_index_prefix::try_sanitize - {index_value:?}") })?; - let index_name = format!( - "{}", - get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY) - ); + let index_name = self.index_name::(); let query = if index_value.is_empty() { "*".to_string() } else { diff --git a/nativelink-store/tests/redis_store_test.rs b/nativelink-store/tests/redis_store_test.rs index 6b3d43a813..288a2b6ba5 100644 --- a/nativelink-store/tests/redis_store_test.rs +++ b/nativelink-store/tests/redis_store_test.rs @@ -1485,6 +1485,100 @@ fn test_search_by_index() -> Result<(), Error> { Ok(()) } +// A store with a `key_prefix` writes every key behind it, so the index it +// creates has to cover `` and be named under the +// prefix too. Indexing the bare schema prefix matched nothing: a scheduler on +// such a store queued every action and dispatched none (railgun, 2026-10-22). +#[nativelink_test] +fn test_search_by_index_covers_the_store_key_prefix() -> Result<(), Error> { + const INDEX: &str = "sched:test:_content_prefix_sort_key_3e762c15"; + fn make_ft_aggregate(result: Result) -> MockCmd { + MockCmd::new( + redis::cmd("FT.AGGREGATE") + .arg(INDEX) + .arg("@content_prefix:{ Searchable }") + .arg("TIMEOUT") + .arg(10000_u64) + .arg("LOAD") + .arg(2) + .arg("data") + .arg("version") + .arg("WITHCURSOR") + .arg("COUNT") + .arg(1500) + .arg("MAXIDLE") + .arg(30000) + .arg("SORTBY") + .arg(2usize) + .arg("@sort_key") + .arg("ASC") + .arg("MAX") + .arg(1_000_000_u64), + result, + ) + } + + let commands = vec![ + // A fresh Redis: the first read finds no index … + make_ft_aggregate(Err(make_extension_error( + "No".to_string(), + Some(format!("such index {INDEX}")), + ))), + // … so the store creates one, over the keys it actually writes. + MockCmd::new( + redis::cmd("FT.CREATE") + .arg(INDEX) + .arg("ON") + .arg("HASH") + .arg("NOHL") + .arg("NOFIELDS") + .arg("NOFREQS") + .arg("NOOFFSETS") + .arg("PREFIX") + .arg(1) + .arg("sched:test:") + .arg("SCHEMA") + .arg("content_prefix") + .arg("TAG") + .arg("sort_key") + .arg("TAG") + .arg("SORTABLE"), + Ok(Value::Okay), + ), + make_ft_info(INDEX, 0), + make_ft_aggregate(Ok(Value::Array(vec![ + Value::Array(vec![ + Value::Int(1), + Value::Array(vec![ + Value::BulkString(b"data".to_vec()), + Value::BulkString(b"1234".to_vec()), + Value::BulkString(b"version".to_vec()), + Value::BulkString(b"1".to_vec()), + ]), + ]), + Value::Int(0), + ]))), + ]; + let store = make_mock_store_with_prefix(commands, "sched:".to_string()).await; + + let search_results: Vec = store + .search_by_index_prefix(SearchByContentPrefix { + prefix: "Searchable".to_string(), + }) + .await + .err_tip(|| "Failed to search by index")? + .try_collect() + .await?; + + assert_eq!( + search_results.len(), + 1, + "the row behind the prefix is found" + ); + assert_eq!(search_results[0].content, "1234"); + Ok(()) +} + /// A provider that reads its index descending, the way the scheduler's /// awaited-action index does. struct SearchByContentPrefixDesc { @@ -2365,6 +2459,46 @@ async fn no_items_from_none_subscription_channel() -> Result<(), Error> { Ok(()) } +// A store with a `key_prefix` publishes the key as Redis holds it, behind +// the prefix, while a subscriber registers the bare store key. The manager +// has to take the prefix off before it looks the subscriber up, or no notice +// ever wakes anyone and every state change waits for the subscriber's own +// timeout re-read (a Redis-backed scheduler ran a build 2.2x slower that way). +#[nativelink_test] +async fn a_notice_behind_the_store_prefix_wakes_the_bare_key_subscriber() -> Result<(), Error> { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let subscription_manager = RedisSubscriptionManager::new( + rx, + Arc::new(async_lock::Mutex::new(vec![])), + "sched:".to_string(), + ); + let mut subscription = subscription_manager.subscribe(TestSchedulerDataUnversioned { + key: "aa_demo".to_string(), + content: String::new(), + version: 0, + })?; + // A fresh subscription reports one change so its owner reads the key + // once; take that, so the wake below can only come from the notice. + subscription.changed().await?; + + tx.send(PushInfo { + kind: redis::PushKind::PMessage, + data: vec![ + Value::BulkString("scheduler_key_change".into()), + Value::BulkString("scheduler_key_change".into()), + Value::BulkString("sched:aa_demo".into()), + ], + }) + .unwrap(); + + timeout(Duration::from_secs(5), subscription.changed()) + .await + .err_tip(|| "the prefixed notice did not wake the subscriber")??; + assert!(!logs_contain("ERROR")); + drop(subscription_manager); + Ok(()) +} + #[nativelink_test] async fn send_messages_to_subscription_channel() -> Result<(), Error> { let (tx, rx) = tokio::sync::mpsc::unbounded_channel();