diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index c074b0dbf..b8e827492 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1785,6 +1785,30 @@ pub struct RedisSpec { #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] pub connection_pool_size: usize, + /// Expire keys written by this store after this many seconds. + /// + /// Redis does not expire these on its own, so a store whose consumer + /// stops keeps every key it ever wrote. A BEP store is the case this + /// exists for: if the ETL stops consuming, the backlog grows until the + /// Redis node runs out of memory. + /// + /// Set this per store, not globally. The same store type backs the CAS + /// fast tier and the scheduler, and expiring scheduler state would drop + /// in-flight actions. + /// + /// This trades data for a bound. Anything not consumed within the window + /// is deleted, so the value has to exceed the longest consumer outage you + /// intend to survive, and it must also exceed how long a single upload can + /// take: the temp key an upload builds carries this same TTL, so a value + /// below the upload duration would expire the write in flight. + /// + /// Zero disables expiry. One second is the smallest value that enables it, + /// and any value that small is almost certainly a mistake. + /// + /// Default: 0 (keys never expire) + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub key_ttl_s: u64, + /// The maximum number of upload chunks to allow per update. /// This is used to limit the amount of memory used when uploading /// large objects to the redis server. A good rule of thumb is to diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index 4097f638d..a58cf3cad 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -125,6 +125,16 @@ pub const DEFAULT_MAX_COUNT_PER_CURSOR: u64 = 1_500; const DEFAULT_CLIENT_PERMITS: usize = 500; +/// Converts a TTL to the seconds `EXPIRE` takes. +/// +/// Clamped to at least one second. Redis treats `EXPIRE` with a value of zero +/// or less as "delete now", so a sub-second duration reaching here would +/// destroy the value it was meant to protect. Config already rejects zero, so +/// this only guards a future caller. +fn ttl_seconds(key_ttl: Duration) -> i64 { + i64::try_from(key_ttl.as_secs()).unwrap_or(i64::MAX).max(1) +} + /// A wrapper around Redis to allow it to be reconnected. pub trait RedisManager where @@ -401,6 +411,11 @@ where /// Per-call ceiling for `check_health` PING. health_check_timeout: Duration, + /// Expire keys this store writes after this long. `None` keeps them + /// forever, which is the default and what every store other than a BEP + /// store wants. + key_ttl: Option, + /// Have we done a subscribe for messages for `remove_callback` subscribes? has_remove_callback_subscribe: OnceCell<()>, @@ -499,6 +514,7 @@ where max_client_permits: usize, max_count_per_cursor: u64, health_check_timeout: Duration, + key_ttl: Option, subscriber_channel: UnboundedReceiver, connection_manager: M, ) -> Result { @@ -526,6 +542,7 @@ where client_permits: Arc::new(Semaphore::new(max_client_permits)), max_count_per_cursor, health_check_timeout, + key_ttl, has_remove_callback_subscribe: OnceCell::const_new(), remove_callbacks, }) @@ -684,6 +701,7 @@ impl RedisStore> { spec.max_client_permits, spec.max_count_per_cursor, Duration::from_millis(spec.health_check_timeout_ms), + (spec.key_ttl_s > 0).then(|| Duration::from_secs(spec.key_ttl_s)), subscriber_channel, ClusterRedisManager::new(client.get_async_connection().await?).await?, ) @@ -811,6 +829,7 @@ impl RedisStore> { spec.max_client_permits, spec.max_count_per_cursor, Duration::from_millis(spec.health_check_timeout_ms), + (spec.key_ttl_s > 0).then(|| Duration::from_secs(spec.key_ttl_s)), subscriber_channel, StandardRedisManager::new(Box::new(move || { Box::pin(Self::connect(spec.clone(), tx.clone())) @@ -1088,6 +1107,29 @@ where return Err(error); } } + // Guard the temp key as soon as it exists. An update that + // dies anywhere after this — mid-upload, during the length + // check, during the rename — would otherwise leave the + // temp key behind forever, and nothing cleans those up. + // EXPIRE on a key that does not exist yet is a no-op, so + // this has to happen after the first chunk lands rather + // than before the loop. + if offset == 0 && let Some(key_ttl) = self.key_ttl { + let ttl_secs = ttl_seconds(key_ttl); + match connection_manager.expire::<_, ()>(temp_key_ref, ttl_secs).await { + Ok(()) => {} + Err(err) if is_retryable_redis_error(&err) => { + let (mut connection_manager, _connect_id) = self.connection_manager.reconnect(connect_id).await?; + connection_manager + .expire::<_, ()>(temp_key_ref, ttl_secs) + .await + .err_tip(|| format!("(after reconnect) while setting TTL on temp key ({temp_key_ref}) in RedisStore::update"))?; + } + Err(err) => { + return Err(Error::from(err).append(format!("While setting TTL on temp key ({temp_key_ref}) in RedisStore::update"))); + } + } + } Ok::(end_pos) }) }) @@ -1176,6 +1218,21 @@ where } } + // Expire the final key, not the temp one: RENAME would carry a TTL + // over, but setting it here keeps the window where the key exists + // without one down to a single round trip and keeps the intent + // obvious. A failure here is a real error rather than something to + // swallow, since a key that silently never expires is the bug this + // option exists to prevent. + if let Some(key_ttl) = self.key_ttl { + let ttl_secs = i64::try_from(key_ttl.as_secs()).unwrap_or(i64::MAX); + client + .connection_manager + .expire::<_, ()>(final_key.as_ref(), ttl_secs) + .await + .err_tip(|| format!("While setting TTL on {final_key} in RedisStore::update()"))?; + } + // If we have a publish channel configured, send a notice that the key has been set. if let Some(pub_sub_channel) = &self.pub_sub_channel { client @@ -1774,22 +1831,29 @@ impl RedisSubscriptionManager { continue; } }; - if value == "evicted" { - trace!(?push_info, "Eviction event"); - let eviction_key = if let Some(key) = push_info.data.get(1) { + // Redis reports maxmemory eviction as + // "evicted" and TTL expiry as "expired". + // Both mean the key is gone, so both have + // to invalidate anything caching its + // existence; treating only one as a removal + // leaves an ExistenceCacheStore claiming to + // hold a key that Redis has already dropped. + if value == "evicted" || value == "expired" { + trace!(?push_info, %value, "Key removal event"); + let removed_key = if let Some(key) = push_info.data.get(1) { if let Value::BulkString(s) = key { String::from_utf8(s.clone()).expect("String message") } else { - error!(?push_info, "Eviction key wasn't bulk-string"); + error!(?push_info, "Removed key wasn't bulk-string"); continue; } } else { - error!(?push_info, "No key in eviction event"); + error!(?push_info, "No key in removal event"); continue; }; - trace!(?eviction_key, "Eviction key"); - let Some((_prefix, internal_key)) = eviction_key.split_once(':') else { - error!(?eviction_key, "Eviction key doesn't contain a colon"); + trace!(?removed_key, "Removed key"); + let Some((_prefix, internal_key)) = removed_key.split_once(':') else { + error!(?removed_key, "Removed key doesn't contain a colon"); continue; }; diff --git a/nativelink-store/tests/redis_store_test.rs b/nativelink-store/tests/redis_store_test.rs index 783ad772f..140087cba 100644 --- a/nativelink-store/tests/redis_store_test.rs +++ b/nativelink-store/tests/redis_store_test.rs @@ -90,6 +90,7 @@ async fn fake_redis_sentinel_master_stream_with_script() -> u16 { async fn make_mock_store_with_prefix_and_subscriber_channel( mut commands: Vec, key_prefix: String, + key_ttl: Option, subscriber_channel: UnboundedReceiver, ) -> RedisStore> { commands.insert( @@ -111,6 +112,7 @@ async fn make_mock_store_with_prefix_and_subscriber_channel( DEFAULT_MAX_PERMITS, DEFAULT_MAX_COUNT_PER_CURSOR, Duration::from_secs(4), + key_ttl, subscriber_channel, manager, ) @@ -118,12 +120,21 @@ async fn make_mock_store_with_prefix_and_subscriber_channel( .unwrap() } +async fn make_mock_store_with_ttl( + commands: Vec, + key_ttl: Duration, +) -> RedisStore> { + let (_tx, rx) = tokio::sync::mpsc::unbounded_channel(); + make_mock_store_with_prefix_and_subscriber_channel(commands, String::new(), Some(key_ttl), rx) + .await +} + async fn make_mock_store_with_prefix( commands: Vec, key_prefix: String, ) -> RedisStore> { let (_tx, rx) = tokio::sync::mpsc::unbounded_channel(); - make_mock_store_with_prefix_and_subscriber_channel(commands, key_prefix, rx).await + make_mock_store_with_prefix_and_subscriber_channel(commands, key_prefix, None, rx).await } #[nativelink_test] @@ -2443,7 +2454,7 @@ async fn send_eviction_to_subscription_channel() -> Result<(), Error> { drop(subscription_manager); assert!(logs_contain( - "Eviction key eviction_key=\"keyprefix:test-eviction\"" + "Removed key removed_key=\"keyprefix:test-eviction\"" )); Ok(()) @@ -2474,6 +2485,7 @@ async fn store_key_coding_round_trip_with_prefix() -> Result<(), Error> { async fn evict_keys_for_existence_cache_core( prefix: String, + removal_event: &str, logs_contain: F, ) -> Result<(), Error> where @@ -2532,7 +2544,8 @@ where ]; let redis_store = Arc::new( - make_mock_store_with_prefix_and_subscriber_channel(commands, prefix.clone(), rx).await, + make_mock_store_with_prefix_and_subscriber_channel(commands, prefix.clone(), None, rx) + .await, ); let existence_store = ExistenceCacheStore::new(&spec, Store::new(redis_store.clone())); @@ -2559,7 +2572,7 @@ where data: vec![ Value::BulkString("demo_pattern".into()), Value::BulkString(format!("keyprefix:{real_key}").into()), - Value::BulkString("evicted".into()), + Value::BulkString(removal_event.into()), ], }) .unwrap(); @@ -2577,7 +2590,7 @@ where .unwrap(); assert!(logs_contain(&format!( - "Eviction key eviction_key=\"keyprefix:{prefix}3031323334353637383961626364656630303030303030303030303030303030-2\"" + "Removed key removed_key=\"keyprefix:{prefix}3031323334353637383961626364656630303030303030303030303030303030-2\"" ))); Ok(()) @@ -2585,10 +2598,128 @@ where #[nativelink_test] async fn evict_keys_for_existence_cache_no_prefix() -> Result<(), Error> { - evict_keys_for_existence_cache_core(String::new(), logs_contain).await + evict_keys_for_existence_cache_core(String::new(), "evicted", logs_contain).await } #[nativelink_test] async fn evict_keys_for_existence_cache_prefix() -> Result<(), Error> { - evict_keys_for_existence_cache_core(String::from("demo_prefix:"), logs_contain).await + evict_keys_for_existence_cache_core(String::from("demo_prefix:"), "evicted", logs_contain).await +} + +/// A key that reaches its TTL has to invalidate the existence cache exactly +/// as a maxmemory eviction does. Redis reports the two differently: `evicted` +/// for maxmemory pressure, `expired` for a TTL. Handling only the former left +/// the cache claiming to hold a key Redis had already dropped, which is +/// newly reachable now that a store can be given a TTL. +#[nativelink_test] +async fn expired_keys_invalidate_the_existence_cache() -> Result<(), Error> { + evict_keys_for_existence_cache_core(String::new(), "expired", logs_contain).await +} + +#[nativelink_test] +async fn expired_keys_invalidate_the_existence_cache_with_prefix() -> Result<(), Error> { + evict_keys_for_existence_cache_core(String::from("demo_prefix:"), "expired", logs_contain).await +} + +/// A store configured with a TTL must expire what it writes. Without this a +/// BEP store keeps every key it ever wrote once its consumer stops, which is +/// how one deployment reached 1.79M keys with `expires=0`. +#[nativelink_test] +async fn update_sets_a_ttl_when_configured() -> Result<(), Error> { + const TTL_SECONDS: u64 = 7 * 24 * 60 * 60; + let data = Bytes::from_static(b"14"); + let digest = DigestInfo::try_new(VALID_HASH1, 2)?; + let packed_hash_hex = format!("{digest}"); + let temp_key = make_temp_key(&packed_hash_hex); + let real_key = packed_hash_hex; + + let store = make_mock_store_with_ttl( + vec![ + MockCmd::new( + redis::cmd("SETRANGE") + .arg(&temp_key) + .arg(0) + .arg(data.to_vec()), + Ok(Value::Int(0)), + ), + // The temp key is guarded as soon as it exists, so an update that + // dies before the rename cannot orphan it. + MockCmd::new( + redis::cmd("EXPIRE") + .arg(&temp_key) + .arg(i64::try_from(TTL_SECONDS).unwrap()), + Ok(Value::Int(1)), + ), + MockCmd::new( + redis::cmd("STRLEN").arg(&temp_key), + Ok(Value::Int(data.len().try_into().unwrap_or(i64::MAX))), + ), + MockCmd::new( + redis::cmd("RENAME").arg(&temp_key).arg(&real_key), + Ok(Value::Nil), + ), + // Re-set after the rename so the retention window starts at + // completion rather than at the first byte. + MockCmd::new( + redis::cmd("EXPIRE") + .arg(&real_key) + .arg(i64::try_from(TTL_SECONDS).unwrap()), + Ok(Value::Int(1)), + ), + // A read afterwards, purely so the expectations that follow EXPIRE + // get consumed. The mock ignores expectations nobody reaches, so + // without something after it a skipped EXPIRE would pass silently. + MockCmd::with_values( + redis::pipe() + .cmd("STRLEN") + .arg(&real_key) + .cmd("EXISTS") + .arg(&real_key), + Ok(vec![Value::Int(2), Value::Boolean(true)]), + ), + ], + Duration::from_secs(TTL_SECONDS), + ) + .await; + + store.update_oneshot(digest, data).await?; + let size = store.has(digest).await?; + assert_eq!(size, Some(2), "the value must survive the expiry being set"); + Ok(()) +} + +/// The default must stay off. The same store type backs the CAS fast tier and +/// the scheduler, and expiring scheduler state would drop in-flight actions. +/// +/// The mock rejects any command it was not given, so an EXPIRE issued here +/// would fail this test. +#[nativelink_test] +async fn update_sets_no_ttl_by_default() -> Result<(), Error> { + let data = Bytes::from_static(b"14"); + let digest = DigestInfo::try_new(VALID_HASH1, 2)?; + let packed_hash_hex = format!("{digest}"); + let temp_key = make_temp_key(&packed_hash_hex); + let real_key = packed_hash_hex; + + let store = make_mock_store(vec![ + MockCmd::new( + redis::cmd("SETRANGE") + .arg(&temp_key) + .arg(0) + .arg(data.to_vec()), + Ok(Value::Int(0)), + ), + MockCmd::new( + redis::cmd("STRLEN").arg(&temp_key), + Ok(Value::Int(data.len().try_into().unwrap_or(i64::MAX))), + ), + MockCmd::new( + redis::cmd("RENAME").arg(&temp_key).arg(&real_key), + Ok(Value::Nil), + ), + ]) + .await; + + store.update_oneshot(digest, data).await?; + Ok(()) }