From f0dc8d94dee45eeb8bf4ec2d53cc0acc5f65a8a6 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 15 Oct 2026 12:43:47 -0400 Subject: [PATCH 1/2] redis: retry the eviction subscription properly, and keep reconnect off the lock Follow-ups from review of the subscription-lifecycle fix. reconnect: subscribing the new connection to patterns recorded mid-reconnect happened while holding the connection write lock, so every get_connection (every Redis operation) waited on those round trips. Subscribing now happens without the lock; the lock is taken only to check nothing is still missing and to publish, repeating if a pattern arrived in between. A pattern recorded after the final check still waits for the lock in get_connection and then subscribes on the new connection, so the original race stays closed. Eviction subscription retry: - The task holds a weak reference, so a dropped store ends it instead of being kept alive and retried forever. - It logs the first failure as an error, then only at powers of two, so a permanent failure (an ACL forbidding PSUBSCRIBE) doesn't log every 5 s. - On a connection-type failure it re-resolves the connection, as the store's other retry paths do, instead of retrying a dead handle. - Reading notify-keyspace-events stays advisory, now logged at warn with its own message so it can't be mistaken for the subscription failing. Tests: - The existing init test only covered the advisory CONFIG GET and passed without any retry. A new test fails the first PSUBSCRIBE and checks the subscription is still established; it fails if the retry loop gives up. - A new test checks a dropped store ends the retry task; it fails if the task holds a strong reference. - A new unit test checks get_connection is not blocked while reconnect catches up; it fails against the previous locking. - The reconnect race test uses #[nativelink_test] instead of a hand-built runtime, so the clippy and lint exemptions are gone, and the test names now describe the fixed behaviour. - The fake Redis gains make_fake_redis_with_first_responses, which answers a command differently the first time only. Co-Authored-By: Claude Opus 5.5 --- nativelink-redis-tester/src/fake_redis.rs | 50 ++ nativelink-redis-tester/src/lib.rs | 3 +- nativelink-store/BUILD.bazel | 1 + nativelink-store/src/redis_store.rs | 566 ++++++++++++--------- nativelink-store/tests/redis_store_test.rs | 187 +++++-- 5 files changed, 528 insertions(+), 279 deletions(-) diff --git a/nativelink-redis-tester/src/fake_redis.rs b/nativelink-redis-tester/src/fake_redis.rs index 95a49685b87..4f38592314f 100644 --- a/nativelink-redis-tester/src/fake_redis.rs +++ b/nativelink-redis-tester/src/fake_redis.rs @@ -15,6 +15,8 @@ use core::fmt::Write; use core::hash::BuildHasher; use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use nativelink_util::background_spawn; use redis::Value; @@ -292,3 +294,51 @@ pub async fn make_fake_redis_with_responses u16 { make_fake_redis_with_multiple_responses(vec![responses]).await } + +/// Like [`make_fake_redis_with_responses`], except that the first request +/// matching a key of `first_responses` gets that key's response instead, once. +/// Every later request gets its normal response. Use this to test that a +/// caller retries after a transient error. +pub async fn make_fake_redis_with_first_responses< + B: BuildHasher + Clone + Send + 'static + Sync, +>( + responses: HashMap, + first_responses: HashMap, +) -> u16 { + let first_responses: Vec<(String, String, Arc)> = first_responses + .into_iter() + .map(|(key, value)| (key, value, Arc::new(AtomicBool::new(false)))) + .collect(); + let handler = move |buf: &[u8]| -> String { + let str_buf = String::from_utf8_lossy(buf).into_owned(); + for (key, value, used) in &first_responses { + if str_buf.starts_with(key) && !used.swap(true, Ordering::SeqCst) { + info!("Responding once to {}", str_buf.replace("\r\n", "\\r\\n")); + return value.clone(); + } + } + for (key, value) in &responses { + if str_buf.starts_with(key) { + info!("Responding to {}", str_buf.replace("\r\n", "\\r\\n")); + return value.clone(); + } + } + warn!( + "Unknown command: {}", + str_buf.chars().take(1000).collect::() + ); + String::new() + }; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let (listener_ready_tx, listener_ready_rx) = oneshot::channel::<()>(); + background_spawn!("listener", async move { + fake_redis_internal(listener, listener_ready_tx, vec![handler]).await; + }); + listener_ready_rx + .await + .expect("Expected successful listener boot"); + info!(port, "Fake redis booted"); + port +} diff --git a/nativelink-redis-tester/src/lib.rs b/nativelink-redis-tester/src/lib.rs index 896d384df96..65ec4874674 100644 --- a/nativelink-redis-tester/src/lib.rs +++ b/nativelink-redis-tester/src/lib.rs @@ -20,7 +20,8 @@ mod read_only_redis; pub use dynamic_fake_redis::{FakeRedisBackend, SubscriptionManagerNotify}; pub use fake_redis::{ add_lua_script, add_to_response, add_to_response_raw, fake_redis_sentinel_master_stream, - fake_redis_sentinel_stream, fake_redis_stream, make_fake_redis_with_responses, + fake_redis_sentinel_stream, fake_redis_stream, make_fake_redis_with_first_responses, + make_fake_redis_with_responses, }; pub use pubsub::MockPubSub; pub use read_only_redis::ReadOnlyRedis; diff --git a/nativelink-store/BUILD.bazel b/nativelink-store/BUILD.bazel index 5742556338f..34a3dd9f296 100644 --- a/nativelink-store/BUILD.bazel +++ b/nativelink-store/BUILD.bazel @@ -267,6 +267,7 @@ rust_test( "@crates//:redis", "@crates//:serde_json", "@crates//:sha2", + "@crates//:tracing-test", ], ) diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index 433e07ede6c..54a3893407f 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -90,6 +90,14 @@ const DEFAULT_RETRY_DELAY: f32 = 0.1; /// paths so a transient error doesn't fail an otherwise healthy request. const MAX_REDIS_RETRY_ATTEMPTS: u32 = 5; +/// First delay between attempts to subscribe to eviction events after a +/// failure. It doubles after each failure, up to +/// [`EVICTION_SUBSCRIBE_MAX_DELAY`]. +const EVICTION_SUBSCRIBE_INITIAL_DELAY: Duration = Duration::from_millis(50); + +/// Longest delay between attempts to subscribe to eviction events. +const EVICTION_SUBSCRIBE_MAX_DELAY: Duration = Duration::from_secs(5); + /// Whether a Redis error is worth re-resolving the master and retrying: a /// dropped/refused connection, an IO error, or a write that hit a freshly /// demoted replica. Anything else (e.g. a real protocol/logic error) is @@ -270,6 +278,16 @@ where Ok(()) } + /// The recorded subscriptions that are not in `subscribed`. + fn subscriptions_not_in(&self, subscribed: &HashSet) -> Vec { + self.subscriptions + .lock() + .iter() + .filter(|pattern| !subscribed.contains(*pattern)) + .cloned() + .collect() + } + async fn new(connect_func: Box>) -> Result { let connection_manager = connect_func().await?; let update_if_version_matches_script = Script::new(LUA_VERSION_SET_SCRIPT); @@ -320,36 +338,27 @@ impl RedisManager for StandardRedisManager let mut connection_manager = (self.connect_func)().await?; let new_uuid = Uuid::new_v4(); self.configure(&mut connection_manager).await?; - let subscriptions = { - let guard = self.subscriptions.lock(); - guard.iter().map(Clone::clone).collect::>() - }; - for subscription in &subscriptions { - connection_manager.psubscribe(subscription).await?; - } - let mut replayed = subscriptions.into_iter().collect::>(); - // Publish the new handle under a brief exclusive lock. Before - // publishing, replay any patterns registered since the snapshot above: - // `psubscribe` records its pattern in `self.subscriptions` *before* it - // fetches a connection, so a pattern missing from this final delta was - // inserted while we hold the write lock, and its `get_connection` will - // block here and observe the new handle — it subscribes on the live - // connection itself. Either way there is no window where a pattern is - // recorded as active but subscribed only on the old connection. - { - let mut guard = self.connection_manager.write().await; - let missing = { - let subs = self.subscriptions.lock(); - subs.iter() - .filter(|s| !replayed.contains(*s)) - .cloned() - .collect::>() - }; - for subscription in missing { + // Subscribe the new connection to every recorded pattern before it is + // published, so a pattern is never recorded as active but missing from + // the live connection. `psubscribe` records its pattern before it calls + // `get_connection`, so a pattern recorded while this runs is either + // picked up by a later pass of this loop, or is recorded after the + // final check, in which case its `get_connection` waits for the write + // lock below and then subscribes on the new connection itself. The + // subscribing happens without the write lock, so `get_connection` + // readers keep using the old handle meanwhile; the lock is held only to + // check that nothing is missing and to publish. + let mut subscribed = HashSet::new(); + loop { + for subscription in self.subscriptions_not_in(&subscribed) { connection_manager.psubscribe(&subscription).await?; - replayed.insert(subscription); + subscribed.insert(subscription); + } + let mut guard = self.connection_manager.write().await; + if self.subscriptions_not_in(&subscribed).is_empty() { + *guard = (connection_manager.clone(), new_uuid); + break; } - *guard = (connection_manager.clone(), new_uuid); } info!(old = %uuid, new = %new_uuid, "StandardRedisManager re-resolved the Redis master"); Ok((connection_manager, new_uuid)) @@ -609,6 +618,86 @@ where } } + /// Subscribes to Redis keyspace events, so evictions and expiries reach the + /// remove callbacks. Succeeds at most once; later calls return at once. + async fn subscribe_to_eviction_events(&self) -> Result<(), Error> { + self.has_remove_callback_subscribe + .get_or_try_init(|| async { + self.warn_if_eviction_events_are_disabled().await; + let (_, uuid) = self.connection_manager.get_connection().await?; + if let Err(err) = self.connection_manager.psubscribe("__key*__:*").await { + // A dropped connection or a failover in progress: re-resolve + // the connection, as the other retry paths in this store do, + // so the next attempt does not reuse a dead one. + if matches!(err.code, Code::Unavailable | Code::DeadlineExceeded) + && let Err(reconnect_err) = self.connection_manager.reconnect(uuid).await + { + debug!(?reconnect_err, "Reconnect after a failed PSUBSCRIBE failed"); + } + return Err(err); + } + Ok(()) + }) + .await?; + Ok(()) + } + + /// Logs if the server's `notify-keyspace-events` setting will not deliver + /// the eviction events the remove callbacks need. Advisory only: the + /// subscription is attempted either way, including when the setting + /// cannot be read (some managed Redis services block `CONFIG`). + async fn warn_if_eviction_events_are_disabled(&self) { + let cfg = async { + let mut client = self.get_client().await?; + redis::cmd("CONFIG") + .arg("GET") + .arg("notify-keyspace-events") + .query_async::>(&mut client.connection_manager) + .await + .map_err(|e| Error::from(e).append("Parsing notify-keyspace-events")) + } + .await; + let cfg = match cfg { + Ok(cfg) => cfg, + Err(err) => { + warn!( + ?err, + "Could not read notify-keyspace-events; subscribing to eviction events anyway" + ); + return; + } + }; + if cfg.len() != 1 { + warn!(?cfg, "Got multiple items for CONFIG GET, expected one"); + return; + } + let Some((_, events_cfg)) = cfg.first() else { + return; + }; + if events_cfg.is_empty() { + error!( + "notify-keyspace-events not enabled for Redis, will fail to get remove callbacks" + ); + } else if !events_cfg.contains('K') { + error!( + notify_keyspace_events = events_cfg, + "notify-keyspace-events does not contain 'K' so won't get keyspace events we need for eviction events" + ); + } else if !events_cfg.contains('A') { + error!( + notify_keyspace_events = events_cfg, + "notify-keyspace-events does not contain 'A' so we won't get eviction events" + ); + } + // FIXME: Redis events spec appears unreliable, so we subscribe anyways + // It should just need Ke as per https://redis.io/docs/latest/develop/pubsub/keyspace-notifications/ + // but I'm yet to get reliable eviction events out of that + info!( + notify_keyspace_events = events_cfg, + "Attempting to subscribe to eviction events" + ); + } + async fn get_client(&self) -> Result, Error> { let local_client_permits = self.client_permits.clone(); let remaining = local_client_permits.available_permits(); @@ -1550,67 +1639,55 @@ where fn register_remove_callback(self: Arc, callback: RemoveCallback) -> Result<(), Error> { debug!(?callback, "New callback"); + // The task holds only a weak reference, so a store that is dropped + // while the subscription is still being retried ends the task instead + // of being kept alive by it. + let weak_self = Arc::downgrade(&self); background_spawn!("remove_callback_subscribe", async move { - self.remove_callbacks.lock().await.push(callback); - // The subscription init is retried until it succeeds: this - // function already told its caller Ok(()) (e.g. - // ExistenceCacheStore, which relies on eviction callbacks for - // invalidation), so a transient Redis error here (e.g. -LOADING - // during a restart/failover) must not permanently disable remove - // callbacks. Previously a one-shot `get_or_try_init` failure was - // only logged, the OnceCell stayed empty, and nothing ever - // retried, so the `__key*__:*` subscription was silently never - // established. - let mut retry_delay = Duration::from_millis(50); + let Some(store) = weak_self.upgrade() else { + return; + }; + store.remove_callbacks.lock().await.push(callback); + drop(store); + // This function has already returned Ok(()) to its caller (an + // `ExistenceCacheStore` relies on these callbacks to invalidate its + // entries), so the subscription is retried until it succeeds. A + // single transient error, such as -LOADING during a restart or + // failover, used to leave remove callbacks off for good. + let mut retry_delay = EVICTION_SUBSCRIBE_INITIAL_DELAY; + let mut failures: u32 = 0; loop { - let local_self = self.clone(); - let init_result = self - .has_remove_callback_subscribe - .get_or_try_init(|| async move { - // The notify-keyspace-events check is advisory: we log - // misconfiguration loudly, but a failure to *read* the - // config must not block the subscription itself. - let cfg_result = async { - let mut client = local_self.get_client().await?; - redis::cmd("CONFIG").arg("GET").arg("notify-keyspace-events").to_owned().query_async::>(&mut client.connection_manager).await.map_err(|e| Error::from(e).append("Parsing notify-keyspace-events")) - }.await; - match &cfg_result { - Ok(cfg) => { - if cfg.len() != 1 { - warn!(?cfg, "Got multiple items for CONFIG GET, expected one"); - } else if let Some((_, events_cfg)) = cfg.first() { - if events_cfg.is_empty() { - error!("notify-keyspace-events not enabled for Redis, will fail to get remove callbacks"); - } else if !events_cfg.contains('K') { - error!(notify_keyspace_events=events_cfg, "notify-keyspace-events does not contain 'K' so won't get keyspace events we need for eviction events"); - } else if !events_cfg.contains('A') { - error!(notify_keyspace_events=events_cfg, "notify-keyspace-events does not contain 'A' so we won't get eviction events"); - } - // FIXME: Redis events spec appears unreliable, so we subscribe anyways - // It should just need Ke as per https://redis.io/docs/latest/develop/pubsub/keyspace-notifications/ - // but I'm yet to get reliable eviction events out of that - info!(notify_keyspace_events=events_cfg, "Attempting to subscribe to eviction events"); - } - } - Err(err) => { - error!(?err, "Error while trying to initialise remove_callback_subscribe: could not verify notify-keyspace-events; subscribing anyway"); - } - } - local_self.connection_manager.psubscribe("__key*__:*").await?; - Ok::<(), Error>(()) - }) - .await; - match init_result { - Ok(()) => break, - Err(err) => { - error!( - ?err, - "Error while trying to initialise remove_callback_subscribe; will retry" + let Some(store) = weak_self.upgrade() else { + debug!("RedisStore dropped before it subscribed to eviction events"); + return; + }; + let Err(err) = store.subscribe_to_eviction_events().await else { + if failures > 0 { + info!( + failures, + "Subscribed to Redis eviction events after retrying" ); - sleep(retry_delay).await; - retry_delay = (retry_delay * 2).min(Duration::from_secs(5)); } + return; + }; + drop(store); + failures += 1; + // Log the first failure as an error, then only at powers of + // two, so a permanent failure (for example an ACL that forbids + // PSUBSCRIBE) does not log every few seconds forever. + if failures == 1 { + error!( + ?err, + "Could not subscribe to Redis eviction events; retrying until it succeeds" + ); + } else if failures.is_power_of_two() { + warn!( + ?err, + failures, "Still unable to subscribe to Redis eviction events" + ); } + sleep(retry_delay).await; + retry_delay = (retry_delay * 2).min(EVICTION_SUBSCRIBE_MAX_DELAY); } }); Ok(()) @@ -2719,33 +2796,21 @@ where } #[cfg(test)] -// This module hand-rolls a minimal RESP3 server plus a dedicated -// current-thread runtime so the psubscribe/reconnect interleaving is exact and -// timing-independent; that legitimately needs the otherwise-disallowed tokio -// runtime/spawn primitives. -#[expect( - clippy::disallowed_methods, - reason = "deterministic in-test RESP server needs a hand-built runtime and raw spawns" -)] -#[expect( - clippy::doc_markdown, - reason = "choreography docs reference bare Redis/RESP identifiers (PSUBSCRIBE, Sha1) by name" -)] -#[expect( - trivial_casts, - reason = "explicit Pin> coercion for the connect closure's return type" -)] -mod psubscribe_reconnect_race_tests { +mod psubscribe_reconnect_tests { use core::time::Duration; - use std::collections::HashSet; + use std::collections::{HashMap, HashSet}; use std::sync::Arc; + use nativelink_macro::nativelink_test; + use nativelink_util::{background_spawn, spawn}; use redis::aio::{ConnectionManager, ConnectionManagerConfig}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::Notify; + use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; + use tokio::time::timeout; - use super::{RedisManager, StandardRedisManager}; + use super::{RedisConnectFn, RedisManager, StandardRedisManager}; /// Reads one CRLF-terminated line from the stream. async fn read_line(stream: &mut TcpStream) -> Option { @@ -2765,7 +2830,7 @@ mod psubscribe_reconnect_race_tests { Some(String::from_utf8_lossy(&line).into_owned()) } - /// Parses one RESP command (array of bulk strings) from the client. + /// Parses one RESP command (an array of bulk strings) from the client. async fn parse_command(stream: &mut TcpStream) -> Option> { let header = read_line(stream).await?; if !header.starts_with('*') { @@ -2783,21 +2848,51 @@ mod psubscribe_reconnect_race_tests { Some(parts) } + /// A `PSUBSCRIBE` reply the fake server holds back on the reconnect + /// connection until the test releases it. + #[derive(Default)] + struct HeldReply { + /// Notified when the `PSUBSCRIBE` arrives. + arrived: Notify, + /// The reply is sent once this is notified. + release: Notify, + } + + #[derive(Default)] struct FakeServerState { - /// Sha1 to return for SCRIPT LOAD (must match the client's own hash). + /// SHA1 to return for `SCRIPT LOAD` (must match the client's own hash). script_hash: String, - /// Patterns each connection (by accept order) was PSUBSCRIBEd to. + /// Patterns each connection, by accept order, was subscribed to. patterns_per_conn: parking_lot::Mutex>>, - /// Signalled when conn #1 (the reconnect connection) receives its - /// first PSUBSCRIBE (the replay of the pre-existing pattern). - replay_psubscribe_arrived: Notify, - /// The server holds conn #1's first PSUBSCRIBE reply until this is - /// notified, keeping reconnect() parked inside its replay loop. - release_replay_reply: Notify, + /// Replies to hold back on connection 1, the one `reconnect` opens, + /// keyed by pattern. + held_on_reconnect: HashMap>, + } + + impl FakeServerState { + fn new(held_patterns: &[&str]) -> Arc { + Arc::new(Self { + script_hash: redis::Script::new(super::LUA_VERSION_SET_SCRIPT) + .get_hash() + .to_string(), + held_on_reconnect: held_patterns + .iter() + .map(|pattern| ((*pattern).to_string(), Arc::default())) + .collect(), + ..Self::default() + }) + } + + fn held(&self, pattern: &str) -> &HeldReply { + &self.held_on_reconnect[pattern] + } + + fn patterns_on(&self, conn_idx: usize) -> Vec { + self.patterns_per_conn.lock()[conn_idx].clone() + } } async fn handle_conn(mut stream: TcpStream, conn_idx: usize, state: Arc) { - let mut first_psubscribe_on_this_conn = true; while let Some(cmd) = parse_command(&mut stream).await { let name = cmd .first() @@ -2820,13 +2915,11 @@ mod psubscribe_reconnect_race_tests { "PSUBSCRIBE" => { let pattern = cmd.get(1).cloned().unwrap_or_default(); state.patterns_per_conn.lock()[conn_idx].push(pattern.clone()); - if conn_idx == 1 && first_psubscribe_on_this_conn { - first_psubscribe_on_this_conn = false; - // Tell the test that reconnect() is now parked awaiting - // this reply inside its replay loop, then hold the - // reply until the test releases it. - state.replay_psubscribe_arrived.notify_one(); - state.release_replay_reply.notified().await; + if conn_idx == 1 + && let Some(held) = state.held_on_reconnect.get(&pattern) + { + held.arrived.notify_one(); + held.release.notified().await; } let reply = format!( ">3\r\n$10\r\npsubscribe\r\n${}\r\n{}\r\n:1\r\n", @@ -2842,121 +2935,134 @@ mod psubscribe_reconnect_race_tests { } } - /// Deterministic repro: a `psubscribe` that interleaves with a `reconnect` - /// whose subscription snapshot was taken before the insert lands only on - /// the old (about-to-be-replaced) connection. The new live connection is - /// never subscribed, yet the pattern is recorded as active in - /// `self.subscriptions`, so notifications are silently lost until an - /// unrelated future reconnect replays the set. - /// - /// Choreography (all real await points, no timing races): - /// 1. Manager connects (conn #0) and psubscribes "pre". - /// 2. Task B calls reconnect(): connects conn #1, snapshots - /// subscriptions = {"pre"}, and parks in the replay loop because the - /// fake server withholds conn #1's PSUBSCRIBE "pre" reply. - /// 3. The test (task A) calls psubscribe("pat"): get_connection still - /// returns conn #0 (reconnect has not published yet), the pattern is - /// inserted into `subscriptions`, and PSUBSCRIBE succeeds on conn #0. - /// 4. The server releases the reply; reconnect publishes conn #1. - /// 5. Conn #1 — now the live connection — was never subscribed to "pat". - #[test] - fn psubscribe_lost_across_reconnect() { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - runtime.block_on(async { - let state = Arc::new(FakeServerState { - script_hash: redis::Script::new(super::LUA_VERSION_SET_SCRIPT) - .get_hash() - .to_string(), - patterns_per_conn: parking_lot::Mutex::new(Vec::new()), - replay_psubscribe_arrived: Notify::new(), - release_replay_reply: Notify::new(), - }); - - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - { - let state = state.clone(); - tokio::spawn(async move { - let mut conn_idx = 0usize; - loop { - let (stream, _) = listener.accept().await.unwrap(); - state.patterns_per_conn.lock().push(Vec::new()); - tokio::spawn(handle_conn(stream, conn_idx, state.clone())); - conn_idx += 1; - } - }); + /// Starts the fake server and returns a manager connected to it, the + /// generation id of its first connection, and the receiver for server + /// pushes, which must stay alive for the length of the test. + async fn connect( + state: &Arc, + ) -> ( + Arc>, + uuid::Uuid, + UnboundedReceiver, + ) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let server_state = state.clone(); + background_spawn!("fake_redis_accept", async move { + let mut conn_idx = 0usize; + loop { + let (stream, _) = listener.accept().await.unwrap(); + server_state.patterns_per_conn.lock().push(Vec::new()); + background_spawn!( + "fake_redis_conn", + handle_conn(stream, conn_idx, server_state.clone()) + ); + conn_idx += 1; } + }); - // Keep the push receiver alive for the whole test. - let (push_tx, _push_rx) = tokio::sync::mpsc::unbounded_channel(); - let url = format!("redis://127.0.0.1:{port}/?protocol=resp3"); - let connect_func = Box::new(move || { - let url = url.clone(); - let push_tx = push_tx.clone(); - Box::pin(async move { - let client = redis::Client::open(url.as_str()).unwrap(); - let config = ConnectionManagerConfig::new() - .set_number_of_retries(0) - .set_connection_timeout(Some(Duration::from_secs(5))) - .set_response_timeout(Some(Duration::from_secs(5))) - .set_push_sender(push_tx); - ConnectionManager::new_with_config(client, config) - .await - .map_err(Into::into) - }) as core::pin::Pin> - }); - - let manager = Arc::new(StandardRedisManager::new(connect_func).await.unwrap()); - - // Step 1: a pre-existing subscription so reconnect's replay loop - // has an await point we can park it on. - manager.psubscribe("pre").await.unwrap(); - let (_conn, uuid0) = manager.get_connection().await.unwrap(); - - // Step 2: start reconnect; it will connect conn #1, snapshot - // subscriptions (= {"pre"}), and park replaying "pre". - let reconnect_task = { - let manager = manager.clone(); - tokio::spawn(async move { manager.reconnect(uuid0).await }) - }; - tokio::time::timeout( - Duration::from_secs(5), - state.replay_psubscribe_arrived.notified(), - ) + let (push_tx, push_rx) = unbounded_channel(); + let url = format!("redis://127.0.0.1:{port}/?protocol=resp3"); + let connect_func: Box> = Box::new(move || { + let url = url.clone(); + let push_tx = push_tx.clone(); + Box::pin(async move { + let client = redis::Client::open(url.as_str()).unwrap(); + let config = ConnectionManagerConfig::new() + .set_number_of_retries(0) + .set_connection_timeout(Some(Duration::from_secs(5))) + .set_response_timeout(Some(Duration::from_secs(5))) + .set_push_sender(push_tx); + ConnectionManager::new_with_config(client, config) + .await + .map_err(Into::into) + }) + }); + let manager = Arc::new(StandardRedisManager::new(connect_func).await.unwrap()); + let (_, uuid) = manager.get_connection().await.unwrap(); + (manager, uuid, push_rx) + } + + /// A `psubscribe` that runs while a `reconnect` is subscribing its new + /// connection must end up subscribed on that new connection. + /// + /// The interleaving, with every step at a real await point: + /// 1. The manager connects (connection 0) and subscribes to `pre`. + /// 2. `reconnect` opens connection 1 and parks while subscribing it to + /// `pre`, because the fake server holds that reply. + /// 3. The test subscribes to `pat`. `get_connection` still returns + /// connection 0, so `pat` lands there. + /// 4. The server releases the reply and `reconnect` publishes connection 1. + /// + /// Before the fix, `reconnect` only replayed a snapshot taken before step 3, + /// so the live connection never subscribed to `pat` and its events were lost + /// until some later reconnect. + #[nativelink_test] + async fn psubscribe_during_reconnect_lands_on_new_connection() { + let state = FakeServerState::new(&["pre"]); + let (manager, uuid0, _push_rx) = connect(&state).await; + manager.psubscribe("pre").await.unwrap(); + + let reconnect = spawn!("reconnect", { + let manager = manager.clone(); + async move { manager.reconnect(uuid0).await } + }); + timeout(Duration::from_secs(5), state.held("pre").arrived.notified()) .await - .expect("reconnect never reached its subscription replay loop"); - - // Step 3: concurrent psubscribe while reconnect is parked - // post-snapshot, pre-publish. It succeeds — on the old connection. - manager.psubscribe("pat").await.unwrap(); - - // Step 4: let reconnect finish and publish conn #1. - state.release_replay_reply.notify_one(); - let (_new_conn, new_uuid) = reconnect_task.await.unwrap().unwrap(); - assert_ne!( - uuid0, new_uuid, - "reconnect must have published a new connection" - ); + .expect("reconnect never started subscribing its new connection"); - // The manager believes "pat" is an active subscription... - let recorded: HashSet = manager.subscriptions.lock().clone(); - assert!( - recorded.contains("pat"), - "sanity: pattern recorded as active" - ); + manager.psubscribe("pat").await.unwrap(); - // ...but the live connection (conn #1) was never subscribed to it. - let conn1_patterns = state.patterns_per_conn.lock()[1].clone(); - assert!( - conn1_patterns.contains(&"pat".to_string()), - "RACE REPRODUCED: pattern \"pat\" is recorded as an active \ - subscription but the live post-reconnect connection was only \ - subscribed to {conn1_patterns:?}; keyspace notifications for \ - \"pat\" are silently lost until the next reconnect", - ); + state.held("pre").release.notify_one(); + let (_new_conn, new_uuid) = reconnect.await.unwrap().unwrap(); + assert_ne!(uuid0, new_uuid, "reconnect must publish a new connection"); + + let recorded: HashSet = manager.subscriptions.lock().clone(); + assert!( + recorded.contains("pat"), + "pattern must be recorded as active" + ); + let live = state.patterns_on(1); + assert!( + live.contains(&"pat".to_string()), + "pattern \"pat\" is recorded as active but the live connection is only \ + subscribed to {live:?}" + ); + } + + /// While `reconnect` subscribes patterns recorded after its first pass, it + /// must not hold the connection write lock, or every `get_connection` + /// (every Redis operation) waits on those network round trips. + #[nativelink_test] + async fn reconnect_catch_up_does_not_block_get_connection() { + let state = FakeServerState::new(&["pre", "pat"]); + let (manager, uuid0, _push_rx) = connect(&state).await; + manager.psubscribe("pre").await.unwrap(); + + let reconnect = spawn!("reconnect", { + let manager = manager.clone(); + async move { manager.reconnect(uuid0).await } }); + timeout(Duration::from_secs(5), state.held("pre").arrived.notified()) + .await + .expect("reconnect never started subscribing its new connection"); + manager.psubscribe("pat").await.unwrap(); + state.held("pre").release.notify_one(); + + // `reconnect` now subscribes connection 1 to `pat`, and the server holds + // that reply. Readers must still get the old connection meanwhile. + timeout(Duration::from_secs(5), state.held("pat").arrived.notified()) + .await + .expect("reconnect never caught up on the pattern recorded mid-reconnect"); + let (_conn, uuid) = timeout(Duration::from_secs(2), manager.get_connection()) + .await + .expect("get_connection blocked while reconnect was subscribing") + .unwrap(); + assert_eq!(uuid, uuid0, "the new connection must not be published yet"); + + state.held("pat").release.notify_one(); + let (_new_conn, new_uuid) = reconnect.await.unwrap().unwrap(); + assert_ne!(uuid0, new_uuid, "reconnect must publish a new connection"); + assert!(state.patterns_on(1).contains(&"pat".to_string())); } } diff --git a/nativelink-store/tests/redis_store_test.rs b/nativelink-store/tests/redis_store_test.rs index 6b3d43a8135..29e12f4a762 100644 --- a/nativelink-store/tests/redis_store_test.rs +++ b/nativelink-store/tests/redis_store_test.rs @@ -26,13 +26,13 @@ use nativelink_macro::nativelink_test; use nativelink_redis_tester::{ ReadOnlyRedis, SubscriptionManagerNotify, add_lua_script, add_to_response, add_to_response_raw, fake_redis_sentinel_master_stream, fake_redis_sentinel_stream, fake_redis_stream, - make_fake_redis_with_responses, + make_fake_redis_with_first_responses, make_fake_redis_with_responses, }; use nativelink_store::cas_utils::ZERO_BYTE_DIGESTS; use nativelink_store::existence_cache_store::ExistenceCacheStore; use nativelink_store::redis_store::{ ClusterRedisManager, DEFAULT_MAX_CHUNK_UPLOADS_PER_UPDATE, DEFAULT_MAX_COUNT_PER_CURSOR, - LUA_VERSION_SET_SCRIPT, RedisStore, RedisSubscriptionManager, decode_key, + LUA_VERSION_SET_SCRIPT, RedisStore, RedisSubscriptionManager, StandardRedisManager, decode_key, }; use nativelink_util::buf_channel::make_buf_channel_pair; use nativelink_util::common::DigestInfo; @@ -44,6 +44,7 @@ use nativelink_util::store_trait::{ TrueValue, UploadSizeInfo, }; use pretty_assertions::assert_eq; +use redis::aio::ConnectionManager; use redis::{PushInfo, RedisError, Value, make_extension_error}; use redis_test::{MockCmd, MockRedisConnection}; use tokio::sync::mpsc::UnboundedReceiver; @@ -2694,7 +2695,10 @@ where Arc::into_inner(RedisStore::new_standard(spec).await.expect("Working spec")).unwrap(); raw_store.replace_temp_name_generator(mock_uuid_generator); let store = Arc::new(raw_store); - store.register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + // Keep `store` alive: the background subscription holds only a weak reference. + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; timeout(Duration::from_secs(3), async move { loop { @@ -2743,21 +2747,16 @@ async fn callback_for_eviction_good() -> Result<(), Error> { Ok(()) } -// Repro for a fail-silent one-shot init bug: `register_remove_callback` -// returns Ok(()) immediately and performs the CONFIG GET + PSUBSCRIBE lazily -// inside a `background_spawn` via `OnceCell::get_or_try_init`. If that init -// hits a single transient Redis error (here: `-LOADING`), the error is only -// logged, the OnceCell stays uninitialized, and nothing ever retries — so the -// `__key*__:*` subscription is never established and evicted/expired keys -// never fire remove callbacks, silently and permanently. +// `register_remove_callback` returns Ok(()) at once and sets up the +// `__key*__:*` subscription in the background. Reading `notify-keyspace-events` +// first is only advisory, so a server that refuses `CONFIG GET` (here with a +// transient `-LOADING`) must still end up subscribed. This used to leave remove +// callbacks off for the life of the process. #[nativelink_test] -async fn transient_init_error_permanently_disables_remove_callbacks() -> Result<(), Error> { +async fn remove_callback_subscribes_even_if_config_cannot_be_read() -> Result<(), Error> { let redis_span = info_span!("redis"); let mut responses = add_lua_version_script(fake_redis_stream()); - // The one-time init query fails with a transient, retryable server error - // (Redis restarting/failing over). The server is otherwise healthy and - // would happily accept a PSUBSCRIBE. add_to_response_raw( &mut responses, redis::cmd("CONFIG") @@ -2773,6 +2772,125 @@ async fn transient_init_error_permanently_disables_remove_callbacks() -> Result< let redis_port = make_fake_redis_with_responses(responses) .instrument(redis_span) .await; + let store = remove_callback_test_store(redis_port).await; + // Keep `store` alive: the background subscription holds only a weak reference. + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + + wait_for_log( + logs_contain, + "Could not read notify-keyspace-events; subscribing to eviction events anyway", + ) + .await; + wait_for_log( + logs_contain, + "new psubscribe complete pattern=\"__key*__:*\" new_subscription=true", + ) + .await; + Ok(()) +} + +// A `PSUBSCRIBE` that fails once with a transient error must be retried until +// the subscription is established. +#[nativelink_test] +async fn remove_callback_subscription_is_retried_after_a_failed_psubscribe() -> Result<(), Error> { + let redis_span = info_span!("redis"); + + let mut responses = add_lua_version_script(fake_redis_stream()); + add_to_response( + &mut responses, + redis::cmd("CONFIG") + .arg("GET") + .arg("notify-keyspace-events"), + vec![Value::Map(vec![( + Value::BulkString(b"notify-keyspace-events".into()), + Value::BulkString(b"KA".into()), + )])], + ); + add_to_response( + &mut responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + vec![Value::Nil], + ); + let mut first_responses = HashMap::new(); + add_to_response_raw( + &mut first_responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + "-LOADING Redis is loading the dataset in memory\r\n".to_string(), + ); + let redis_port = make_fake_redis_with_first_responses(responses, first_responses) + .instrument(redis_span) + .await; + let store = remove_callback_test_store(redis_port).await; + // Keep `store` alive: the background subscription holds only a weak reference. + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + + wait_for_log( + logs_contain, + "Could not subscribe to Redis eviction events; retrying until it succeeds", + ) + .await; + wait_for_log( + logs_contain, + "new psubscribe complete pattern=\"__key*__:*\" new_subscription=true", + ) + .await; + wait_for_log( + logs_contain, + "Subscribed to Redis eviction events after retrying failures=1", + ) + .await; + Ok(()) +} + +// The subscription retry task holds only a weak reference to the store, so a +// store dropped while its subscription keeps failing (here every PSUBSCRIBE is +// refused) ends the task instead of being kept alive and retried forever. +#[nativelink_test] +async fn remove_callback_retry_stops_when_the_store_is_dropped() -> Result<(), Error> { + let redis_span = info_span!("redis"); + + let mut responses = add_lua_version_script(fake_redis_stream()); + add_to_response_raw( + &mut responses, + redis::cmd("CONFIG") + .arg("GET") + .arg("notify-keyspace-events"), + "-NOPERM this user has no permissions to run the 'config|get' command\r\n".to_string(), + ); + add_to_response_raw( + &mut responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + "-NOPERM this user has no permissions to access one of the channels\r\n".to_string(), + ); + let redis_port = make_fake_redis_with_responses(responses) + .instrument(redis_span) + .await; + let store = remove_callback_test_store(redis_port).await; + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + + wait_for_log( + logs_contain, + "Could not subscribe to Redis eviction events; retrying until it succeeds", + ) + .await; + drop(store); + wait_for_log( + logs_contain, + "RedisStore dropped before it subscribed to eviction events", + ) + .await; + Ok(()) +} + +async fn remove_callback_test_store( + redis_port: u16, +) -> Arc>> { let spec = RedisSpec { addresses: vec![format!("redis://127.0.0.1:{redis_port}/")], mode: RedisMode::Standard, @@ -2781,46 +2899,19 @@ async fn transient_init_error_permanently_disables_remove_callbacks() -> Result< let mut raw_store = Arc::into_inner(RedisStore::new_standard(spec).await.expect("Working spec")).unwrap(); raw_store.replace_temp_name_generator(mock_uuid_generator); - let store = Arc::new(raw_store); - // The caller (e.g. ExistenceCacheStore at construction, its only call - // site) is told Ok(()) and believes eviction tracking is live. - store.register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + Arc::new(raw_store) +} - // Deterministic sync point: wait until the background init task has hit - // the transient error and logged it. +/// Waits up to 5 s for `needle` to appear in the calling test's logs. +/// `logs_contain` is the function `#[nativelink_test]` defines in each test. +async fn wait_for_log(logs_contain: impl Fn(&str) -> bool + Send + Sync, needle: &str) { timeout(Duration::from_secs(5), async { - loop { - if logs_contain("Error while trying to initialise remove_callback_subscribe") { - break; - } + while !logs_contain(needle) { sleep(Duration::from_millis(20)).await; } }) .await - .expect("background init never ran/logged its error"); - - // Since register_remove_callback reported success, the keyspace-event - // subscription must (eventually) be established — otherwise remove - // callbacks are silently dead forever. On current code there is no retry - // path at all (the OnceCell is left empty and only another - // register_remove_callback call could re-drive it), so this never happens. - let subscribed = timeout(Duration::from_secs(2), async { - loop { - if logs_contain("new psubscribe complete pattern=\"__key*__:*\"") { - break; - } - sleep(Duration::from_millis(50)).await; - } - }) - .await - .is_ok(); - assert!( - subscribed, - "register_remove_callback returned Ok(()) but a single transient CONFIG GET \ - error permanently disabled the __key*__:* subscription; evicted/expired keys \ - will never invoke remove callbacks (stale ExistenceCacheStore entries)" - ); - Ok(()) + .unwrap_or_else(|_| panic!("timed out waiting for log line: {needle}")); } #[nativelink_test] From 53cfe6f7507a8a73ac10ff682292fa1a9651af5b Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 22 Oct 2026 14:10:25 +0000 Subject: [PATCH 2/2] redis: require Send of the connection in RedisStore's inherent impl `warn_if_eviction_events_are_disabled` holds a `ClientWithPermit` across an await, which the nightly clippy lane (`future_not_send`) rejects while `C` is only `Sync` there. Every connection type used is `Send`, and the `StoreDriver` impl already requires it. Co-Authored-By: Claude Fable 5.1 --- nativelink-store/src/redis_store.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index 54a3893407f..15cfad3cc04 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -541,8 +541,8 @@ pub fn decode_key<'a>( impl RedisStore where - C: ConnectionLike + Clone + Sync, - M: RedisManager + Sync, + C: ConnectionLike + Clone + Send + Sync, + M: RedisManager + Send + Sync, { /// Used for testing when determinism is required. #[expect(clippy::too_many_arguments)]