Feat/redis read retry backoff

crates/http-service/src/executor/http.rs

@@ -283,6 +283,8 @@ mod tests {
         fn count_kv_read(&self, _value: i32) {}
 
         fn count_kv_byod_read(&self, _value: i32) {}
+
+        fn count_kv_read_cached(&self) {}
     }
 
     impl UserDiagStats for TestStats {
@@ -354,6 +356,7 @@ mod tests {
             &self,
             _request_id: &SmolStr,
             _app: &SmolStr,
+            _caller_ip: std::net::Ipv4Addr,
             _cfg: &App,
         ) -> Arc<dyn StatsVisitor> {
             Arc::new(TestStats)

crates/http-service/src/lib.rs

@@ -359,7 +359,15 @@ where
                 }
             };
 
-            let stats = self.context.new_stats_row(&traceparent, &app_name, &cfg);
+            let caller_ip = request
+                .headers()
+                .get(crate::executor::X_REAL_IP)
+                .and_then(|v| v.to_str().ok())
+                .and_then(|v| v.parse::<std::net::Ipv4Addr>().ok())
+                .unwrap_or(std::net::Ipv4Addr::UNSPECIFIED);
+            let stats = self
+                .context
+                .new_stats_row(&traceparent, &app_name, caller_ip, &cfg);
 
             let response = match executor.execute(request, stats.clone()).await {
                 Ok(mut response) => {

crates/http-service/src/state.rs

@@ -236,6 +236,7 @@ mod tests {
     impl ReadStats for TestStats {
         fn count_kv_read(&self, _value: i32) {}
         fn count_kv_byod_read(&self, _value: i32) {}
+        fn count_kv_read_cached(&self) {}
     }
     impl UserDiagStats for TestStats {
         fn set_user_diag(&self, _diag: &str) {}

crates/key-value-store/Cargo.toml

@@ -8,16 +8,17 @@ description = "key-value store host function"
 
 [features]
 default = []
-redis = ["dep:redis"]
+redis = ["dep:redis", "dep:tokio"]
 
 [dependencies]
 reactor = { path = "../reactor" }
-wasmtime = {workspace = true}
+wasmtime = { workspace = true }
 slab = "0.4"
 async-trait = "0.1"
-smol_str = {workspace = true}
+smol_str = { workspace = true }
 tracing = "0.1"
-redis = { version = "1.2", features = ["aio", "tokio-comp", "connection-manager", "tokio-native-tls-comp"], optional = true}
+tokio = { workspace = true, optional = true }
+redis = { version = "1.2", features = ["aio", "tokio-comp", "connection-manager", "tokio-native-tls-comp"], optional = true }
 
 [dev-dependencies]
 tokio = { version = "1", features = ["macros", "rt-multi-thread"] }

crates/key-value-store/src/lib.rs

@@ -13,10 +13,28 @@ pub use key_value::{Error, Value};
 #[cfg(feature = "redis")]
 pub use redis_impl::RedisStore;
 
+/// Outcome of a [`Store::get_tracked`] call: the value plus where it came from.
+pub struct GetOutcome {
+    pub value: Option<Value>,
+    /// `true` when the value was served from an in-process cache instead of
+    /// the backing store.
+    pub from_cache: bool,
+}
+
 #[async_trait::async_trait]
 pub trait Store: Sync + Send {
     async fn get(&self, key: &str) -> Result<Option<Value>, Error>;
 
+    /// Like [`Store::get`], but also reports whether the value was served from
+    /// an in-process cache. Layers that do not cache inherit this default,
+    /// which delegates to [`Store::get`] and reports a backing-store read.
+    async fn get_tracked(&self, key: &str) -> Result<GetOutcome, Error> {
+        Ok(GetOutcome {
+            value: self.get(key).await?,
+            from_cache: false,
+        })
+    }
+
     async fn zrange_by_score(
         &self,
         key: &str,
@@ -47,6 +65,10 @@ pub trait ReadStats: Sync + Send {
     fn count_kv_read(&self, value: i32);
     /// Increment key-value read count and size for BYOD
     fn count_kv_byod_read(&self, value: i32);
+    /// Increment the count of reads served from the in-process cache. Counted
+    /// in addition to [`ReadStats::count_kv_read`] /
+    /// [`ReadStats::count_kv_byod_read`], so it is a subset of the total reads.
+    fn count_kv_read_cached(&self);
 }
 
 #[derive(Clone)]
@@ -367,6 +389,8 @@ mod tests {
             self.byod_reads
                 .fetch_add(value, std::sync::atomic::Ordering::Relaxed);
         }
+
+        fn count_kv_read_cached(&self) {}
     }
 
     // Mock implementation of StoreManager

crates/key-value-store/src/redis_impl.rs

@@ -16,6 +16,43 @@ const REDIS_CONNECTION_TIMEOUT: Duration = Duration::from_millis(100);
 const REDIS_NUMBER_OF_RETRIES: usize = 2;
 const REDIS_MAX_RECONNECT_DELAY: Duration = Duration::from_millis(100);
 
+/// Retry budget for a single read command (on top of connection-manager
+/// reconnection above). Kept small and fixed for the same reason as the
+/// timeouts: Redis is on the request hot path, so a retry loop must have a
+/// bounded worst case rather than let a stalled dependency pile up requests.
+const REDIS_READ_RETRIES: usize = 2;
+const REDIS_READ_RETRY_BASE_DELAY: Duration = Duration::from_millis(10);
+const REDIS_READ_RETRY_MAX_DELAY: Duration = Duration::from_millis(50);
+
+fn redis_read_retry_delay(attempt: usize) -> Duration {
+    let scale = 1u32 << (attempt.saturating_sub(1) as u32);
+    (REDIS_READ_RETRY_BASE_DELAY * scale).min(REDIS_READ_RETRY_MAX_DELAY)
+}
+
+/// Retry a read command up to `REDIS_READ_RETRIES` times with capped
+/// exponential backoff. `op` is re-invoked from scratch on each attempt so it
+/// can pick a fresh pooled connection (used by `get`/`zrange_by_score`/
+/// `bf_exists`; `scan`/`zscan` retry the same connection since their iterator
+/// borrows it across the loop).
+async fn retry_read<T, F, Fut>(mut op: F) -> Result<T, ::redis::RedisError>
+where
+    F: FnMut() -> Fut,
+    Fut: std::future::Future<Output = Result<T, ::redis::RedisError>>,
+{
+    let mut attempt = 0;
+    loop {
+        match op().await {
+            Ok(value) => return Ok(value),
+            Err(error) if attempt < REDIS_READ_RETRIES => {
+                attempt += 1;
+                tracing::debug!(attempt, cause = ?error, "kv-store: redis read retry");
+                tokio::time::sleep(redis_read_retry_delay(attempt)).await;
+            }
+            Err(error) => return Err(error),
+        }
+    }
+}
+
 /// Build the fail-fast connection-manager config for KV-store Redis connections.
 fn connection_manager_config() -> ConnectionManagerConfig {
     ConnectionManagerConfig::new()
@@ -78,10 +115,12 @@ impl RedisStore {
 #[async_trait::async_trait]
 impl Store for RedisStore {
     async fn get(&self, key: &str) -> Result<Option<Value>, Error> {
-        self.conn().get(key).await.map_err(|error| {
-            tracing::warn!(cause = ?error, key, "kv-store: redis get");
-            Error::InternalError
-        })
+        retry_read(|| async { self.conn().get(key).await })
+            .await
+            .map_err(|error| {
+                tracing::warn!(cause = ?error, key, "kv-store: redis get");
+                Error::InternalError
+            })
     }
 
     async fn zrange_by_score(
@@ -90,8 +129,7 @@ impl Store for RedisStore {
         min: f64,
         max: f64,
     ) -> Result<Vec<(Value, f64)>, Error> {
-        self.conn()
-            .zrangebyscore_withscores(key, min, max)
+        retry_read(|| async { self.conn().zrangebyscore_withscores(key, min, max).await })
             .await
             .map_err(|error| {
                 tracing::warn!(cause = ?error, key, min, max, "kv-store: redis zrangebyscore");
@@ -101,10 +139,21 @@ impl Store for RedisStore {
 
     async fn scan(&self, pattern: &str) -> Result<Vec<String>, Error> {
         let mut conn = self.conn();
-        let mut it = conn.scan_match(pattern).await.map_err(|error| {
-            tracing::warn!(cause = ?error, pattern, "kv-store: redis scan_match");
-            Error::InternalError
-        })?;
+        let mut attempt = 0;
+        let mut it = loop {
+            match conn.scan_match(pattern).await {
+                Ok(it) => break it,
+                Err(error) if attempt < REDIS_READ_RETRIES => {
+                    attempt += 1;
+                    tracing::debug!(attempt, cause = ?error, pattern, "kv-store: redis scan_match retry");
+                    tokio::time::sleep(redis_read_retry_delay(attempt)).await;
+                }
+                Err(error) => {
+                    tracing::warn!(cause = ?error, pattern, "kv-store: redis scan_match");
+                    return Err(Error::InternalError);
+                }
+            }
+        };
         let mut ret = vec![];
         while let Some(element) = it.next_item().await {
             ret.push(element.map_err(|error| {
@@ -117,11 +166,21 @@ impl Store for RedisStore {
 
     async fn zscan(&self, key: &str, pattern: &str) -> Result<Vec<(Value, f64)>, Error> {
         let mut conn = self.conn();
-        let mut it: AsyncIter<(Value, f64)> =
-            conn.zscan_match(key, pattern).await.map_err(|error| {
-                tracing::warn!(cause = ?error, key, pattern, "kv-store: redis zscan_match");
-                Error::InternalError
-            })?;
+        let mut attempt = 0;
+        let mut it: AsyncIter<(Value, f64)> = loop {
+            match conn.zscan_match(key, pattern).await {
+                Ok(it) => break it,
+                Err(error) if attempt < REDIS_READ_RETRIES => {
+                    attempt += 1;
+                    tracing::debug!(attempt, cause = ?error, key, pattern, "kv-store: redis zscan_match retry");
+                    tokio::time::sleep(redis_read_retry_delay(attempt)).await;
+                }
+                Err(error) => {
+                    tracing::warn!(cause = ?error, key, pattern, "kv-store: redis zscan_match");
+                    return Err(Error::InternalError);
+                }
+            }
+        };
         let mut ret = vec![];
         while let Some(element) = it.next_item().await {
             ret.push(element.map_err(|error| {
@@ -133,14 +192,17 @@ impl Store for RedisStore {
     }
 
     async fn bf_exists(&self, key: &str, item: &str) -> Result<bool, Error> {
-        redis::cmd("BF.EXISTS")
-            .arg(key)
-            .arg(item)
-            .query_async(&mut self.conn())
-            .await
-            .map_err(|error| {
-                tracing::warn!(cause = ?error, key, item, "kv-store: redis bf_exists");
-                Error::InternalError
-            })
+        retry_read(|| async {
+            redis::cmd("BF.EXISTS")
+                .arg(key)
+                .arg(item)
+                .query_async(&mut self.conn())
+                .await
+        })
+        .await
+        .map_err(|error| {
+            tracing::warn!(cause = ?error, key, item, "kv-store: redis bf_exists");
+            Error::InternalError
+        })
     }
 }

crates/runtime/src/lib.rs

@@ -1,6 +1,7 @@
 use crate::app::KvStoreOption;
 use crate::store::HasStats;
 use http_backend::stats::ExtStatsTimer;
+use std::net::Ipv4Addr;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::Instant;
@@ -433,6 +434,7 @@ pub trait ContextT {
         &self,
         request_id: &SmolStr,
         app: &SmolStr,
+        caller_ip: Ipv4Addr,
         cfg: &App,
     ) -> Arc<dyn StatsVisitor>;
 }

crates/runtime/src/store.rs

@@ -616,6 +616,7 @@ mod tests {
         impl ReadStats for NoStats {
             fn count_kv_read(&self, _: i32) {}
             fn count_kv_byod_read(&self, _: i32) {}
+            fn count_kv_read_cached(&self) {}
         }
         impl UserDiagStats for NoStats {
             fn set_user_diag(&self, _: &str) {}

crates/runtime/src/util/stats.rs

@@ -144,6 +144,8 @@ mod tests {
         fn count_kv_byod_read(&self, value: i32) {
             self.byod_reads.fetch_add(value, Ordering::Relaxed);
         }
+
+        fn count_kv_read_cached(&self) {}
     }
 
     impl UserDiagStats for MockStatsVisitor {

src/context.rs

@@ -94,6 +94,7 @@ impl ContextT for Context {
         &self,
         _request_id: &SmolStr,
         _app: &SmolStr,
+        _caller_ip: std::net::Ipv4Addr,
         _cfg: &App,
     ) -> Arc<dyn StatsVisitor> {
         Arc::new(StatsStub::default())
@@ -190,6 +191,8 @@ impl ReadStats for StatsStub {
     fn count_kv_read(&self, _value: i32) {}
 
     fn count_kv_byod_read(&self, _value: i32) {}
+
+    fn count_kv_read_cached(&self) {}
 }
 
 impl UserDiagStats for StatsStub {

src/key_value.rs

@@ -32,6 +32,7 @@ mod tests {
     impl ReadStats for NoOpStats {
         fn count_kv_read(&self, _: i32) {}
         fn count_kv_byod_read(&self, _: i32) {}
+        fn count_kv_read_cached(&self) {}
     }
 
     fn make_kv_option(name: &str, param: &str) -> KvStoreOption {