Add an optional TTL for keys a redis store writes

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

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<C>
 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<Duration>,
+
     /// 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<Duration>,
         subscriber_channel: UnboundedReceiver<PushInfo>,
         connection_manager: M,
     ) -> Result<Self, Error> {
@@ -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<ClusterConnection, ClusterRedisManager<ClusterConnection>> {
             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<ConnectionManager, StandardRedisManager<ConnectionManager>> {
             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::<u32, Error>(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;
                                             };
 

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<MockCmd>,
     key_prefix: String,
+    key_ttl: Option<Duration>,
     subscriber_channel: UnboundedReceiver<PushInfo>,
 ) -> RedisStore<MockRedisConnection, ClusterRedisManager<MockRedisConnection>> {
     commands.insert(
@@ -111,19 +112,29 @@ 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,
     )
     .await
     .unwrap()
 }
 
+async fn make_mock_store_with_ttl(
+    commands: Vec<MockCmd>,
+    key_ttl: Duration,
+) -> RedisStore<MockRedisConnection, ClusterRedisManager<MockRedisConnection>> {
+    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<MockCmd>,
     key_prefix: String,
 ) -> RedisStore<MockRedisConnection, ClusterRedisManager<MockRedisConnection>> {
     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<F>(
     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,18 +2590,136 @@ where
     .unwrap();
 
     assert!(logs_contain(&format!(
-        "Eviction key eviction_key=\"keyprefix:{prefix}3031323334353637383961626364656630303030303030303030303030303030-2\""
+        "Removed key removed_key=\"keyprefix:{prefix}3031323334353637383961626364656630303030303030303030303030303030-2\""
     )));
 
     Ok(())
 }
 
 #[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(())
 }