[Bug fix] Adds retry logic to redis store

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-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?,