Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 47 additions & 33 deletions nativelink-store/src/redis_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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!(
Expand Down Expand Up @@ -2184,6 +2197,31 @@ impl SchedulerSubscriptionManager for RedisSubscriptionManager {
}
}

impl<C, M> RedisStore<C, M>
where
C: ConnectionLike + Clone + Send + Sync,
M: RedisManager<C> + 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<K: SchedulerIndexProvider>(&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<K: SchedulerIndexProvider>(&self) -> String {
format!(
"{}{}",
self.key_prefix,
get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY)
)
}
}

impl<C, M> SchedulerStore for RedisStore<C, M>
where
C: Clone + ConnectionLike + Sync + Send + 'static,
Expand Down Expand Up @@ -2374,13 +2412,11 @@ where
.err_tip(|| {
format!("In RedisStore::search_by_index_prefix::try_sanitize - {index_value:?}")
})?;
let index_name = self.index_name::<K>();
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 {
Expand Down Expand Up @@ -2410,7 +2446,7 @@ where
});
}
let create_options = FtCreateOptions {
prefixes: vec![K::KEY_PREFIX.into()],
prefixes: vec![self.index_key_prefix::<K>()],
nohl: true,
nofields: true,
nofreqs: true,
Expand All @@ -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
};

Expand All @@ -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(_) => {
Expand Down Expand Up @@ -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})")
})
}
});
Expand All @@ -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
Expand Down Expand Up @@ -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::<K>();
let query = if index_value.is_empty() {
"*".to_string()
} else {
Expand Down
134 changes: 134 additions & 0 deletions nativelink-store/tests/redis_store_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 `<key_prefix><schema prefix>` 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<Value, RedisError>) -> 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<TestSchedulerDataUnversioned> = 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 {
Expand Down Expand Up @@ -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();
Expand Down
Loading