[Bug fix] Adds retry logic to redis store

Adds retry logic and configs for redis store. This will setup redis to
reconnect and retry commands if the connection to redis is lost.

nativelink-config/src/stores.rs

@@ -924,6 +924,71 @@ pub struct RedisStore {
     /// Default: standard,
     #[serde(default)]
     pub mode: RedisMode,
+
+    /// When using pubsub interface, this is the maximum number of items to keep
+    /// queued up before dropping old items.
+    ///
+    /// Default: 4096
+    #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")]
+    pub broadcast_channel_capacity: usize,
+
+    /// The amount of time in milliseconds until the redis store considers the
+    /// command to be timed out. This will trigger a retry of the command and
+    /// potentially a reconnection to the redis server.
+    ///
+    /// Default: 10000 (10 seconds)
+    #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")]
+    pub command_timeout_ms: u64,
+
+    /// The amount of time in milliseconds until the redis store considers the
+    /// connection to unresponsive. This will trigger a reconnection to the
+    /// redis server.
+    ///
+    /// Default: 3000 (3 seconds)
+    #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")]
+    pub connection_timeout_ms: u64,
+
+    /// The amount of data to read from the redis server at a time.
+    /// This is used to limit the amount of memory used when reading
+    /// large objects from the redis server as well as limiting the
+    /// amount of time a single read operation can take.
+    ///
+    /// IMPORTANT: If this value is too high, the `command_timeout_ms`
+    /// might be triggered if the latency or throughput to the redis
+    /// server is too low.
+    ///
+    /// Default: 64KiB
+    #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")]
+    pub read_chunk_size: usize,
+
+    /// The number of connections to keep open to the redis server(s).
+    ///
+    /// Default: 3
+    #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")]
+    pub connection_pool_size: usize,
+
+    /// 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
+    /// think of the data as:
+    /// AVAIL_MEMORY / (read_chunk_size * max_chunk_uploads_per_update) = THORETICAL_MAX_CONCURRENT_UPLOADS
+    /// (note: it is a good idea to divide AVAIL_MAX_MEMORY by ~10 to account for other memory usage)
+    ///
+    /// Default: 10
+    #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")]
+    pub max_chunk_uploads_per_update: usize,
+
+    /// Retry configuration to use when a network request fails.
+    /// See the `Retry` struct for more information.
+    ///
+    /// Default: Retry {
+    ///   max_retries: 0, /* unlimited */
+    ///   delay: 0.1, /* 100ms */
+    ///   jitter: 0.5, /* 50% */
+    ///   retry_on_errors: None, /* not used in redis store */
+    /// }
+    #[serde(default)]
+    pub retry: Retry,
 }
 
 #[derive(Debug, Default, Deserialize, Serialize, Clone, PartialEq, Eq)]

nativelink-metric-collector/Cargo.toml

@@ -6,7 +6,7 @@ rust-version = "1.79.0"
 
 [dependencies]
 nativelink-metric = { path = "../nativelink-metric" }
-opentelemetry = { version = "0.24.0", default-features = false }
+opentelemetry = { version = "0.24.0", features = ["metrics"], default-features = false }
 parking_lot = "0.12.3"
 serde = { version = "1.0.210", default-features = false }
 tracing = { version = "0.1.40", default-features = false }

nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs

@@ -20,11 +20,12 @@ use std::time::{Duration, SystemTime};
 
 use bytes::Bytes;
 use fred::bytes_utils::string::Str;
+use fred::clients::SubscriberClient;
 use fred::error::{RedisError, RedisErrorKind};
 use fred::mocks::{MockCommand, Mocks};
-use fred::prelude::Builder;
-use fred::types::{RedisConfig, RedisValue};
-use mock_instant::thread_local::SystemTime as MockSystemTime;
+use fred::prelude::{Builder, RedisPool};
+use fred::types::{PerformanceConfig, RedisConfig, RedisValue};
+use mock_instant::global::SystemTime as MockSystemTime;
 use nativelink_error::Error;
 use nativelink_macro::nativelink_test;
 use nativelink_scheduler::awaited_action_db::{
@@ -46,6 +47,7 @@ const INSTANCE_NAME: &str = "instance_name";
 const TEMP_UUID: &str = "550e8400-e29b-41d4-a716-446655440000";
 const SCRIPT_VERSION: &str = "3e762c15";
 const VERSION_SCRIPT_HASH: &str = "fdf1152fd21705c8763752809b86b55c5d4511ce";
+const MAX_CHUNK_UPLOADS_PER_UPDATE: usize = 10;
 
 fn mock_uuid_generator() -> String {
     uuid::Uuid::parse_str(TEMP_UUID).unwrap().to_string()
@@ -145,6 +147,20 @@ impl Drop for MockRedisBackend {
     }
 }
 
+fn make_clients(mut builder: Builder) -> (RedisPool, SubscriberClient) {
+    const CONNECTION_POOL_SIZE: usize = 1;
+    let client_pool = builder
+        .set_performance_config(PerformanceConfig {
+            broadcast_channel_capacity: 4096,
+            ..Default::default()
+        })
+        .build_pool(CONNECTION_POOL_SIZE)
+        .unwrap();
+
+    let subscriber_client = builder.build_subscriber_client().unwrap();
+    (client_pool, subscriber_client)
+}
+
 #[nativelink_test]
 async fn add_action_smoke_test() -> Result<(), Error> {
     const CLIENT_OPERATION_ID: &str = "my_client_operation_id";
@@ -389,13 +405,16 @@ async fn add_action_smoke_test() -> Result<(), Error> {
             mocks: Some(Arc::clone(&mocks) as Arc<dyn Mocks>),
             ..Default::default()
         });
-
+        let (client_pool, subscriber_client) = make_clients(builder);
         Arc::new(
             RedisStore::new_from_builder_and_parts(
-                builder,
+                client_pool,
+                subscriber_client,
                 Some(SUB_CHANNEL.into()),
                 mock_uuid_generator,
                 String::new(),
+                4064,
+                MAX_CHUNK_UPLOADS_PER_UPDATE,
             )
             .unwrap(),
         )

nativelink-service/src/worker_api_server.rs

@@ -13,10 +13,10 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::convert::Into;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
-use std::convert::Into;
 
 use futures::stream::unfold;
 use futures::Stream;

nativelink-store/src/default_store_factory.rs

@@ -53,7 +53,7 @@ pub fn store_factory<'a>(
             StoreConfig::experimental_s3_store(config) => {
                 S3Store::new(config, SystemTime::now).await?
             }
-            StoreConfig::redis_store(config) => RedisStore::new(config)?,
+            StoreConfig::redis_store(config) => RedisStore::new(config.clone())?,
             StoreConfig::verify(config) => VerifyStore::new(
                 config,
                 store_factory(&config.backend, store_manager, None).await?,