[ENH] Timeout for s3 storage

chroma-core · Jun 18, 2024 · 9c5254f · 9c5254f
1 parent fd12d4a
commit 9c5254f
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 0 deletions.
diff --git a/rust/worker/chroma_config.yaml b/rust/worker/chroma_config.yaml
@@ -24,6 +24,8 @@ query_service:
         S3:
             bucket: "chroma-storage"
             credentials: "Minio"
+            connect_timeout_ms: 5000
+            request_timeout_ms: 1000
     log:
         Grpc:
             host: "logservice.chroma"
@@ -56,6 +58,8 @@ compaction_service:
         S3:
             bucket: "chroma-storage"
             credentials: "Minio"
+            connect_timeout_ms: 5000
+            request_timeout_ms: 1000
     log:
         Grpc:
             host: "logservice.chroma"

diff --git a/rust/worker/src/config.rs b/rust/worker/src/config.rs
@@ -174,6 +174,8 @@ mod tests {
                         S3:
                             bucket: "chroma"
                             credentials: Minio
+                            connect_timeout_ms: 5000
+                            request_timeout_ms: 1000
                     log:
                         Grpc:
                             host: "localhost"
@@ -206,6 +208,8 @@ mod tests {
                         S3:
                             bucket: "chroma"
                             credentials: Minio
+                            connect_timeout_ms: 5000
+                            request_timeout_ms: 1000
                     log:
                         Grpc:
                             host: "localhost"
@@ -262,6 +266,8 @@ mod tests {
                         S3:
                             bucket: "chroma"
                             credentials: Minio
+                            connect_timeout_ms: 5000
+                            request_timeout_ms: 1000
                     log:
                         Grpc:
                             host: "localhost"
@@ -294,6 +300,8 @@ mod tests {
                         S3:
                             bucket: "chroma"
                             credentials: Minio
+                            connect_timeout_ms: 5000
+                            request_timeout_ms: 1000
                     log:
                         Grpc:
                             host: "localhost"
@@ -368,6 +376,8 @@ mod tests {
                         S3:
                             bucket: "chroma"
                             credentials: Minio
+                            connect_timeout_ms: 5000
+                            request_timeout_ms: 1000
                     log:
                         Grpc:
                             host: "localhost"
@@ -400,6 +410,8 @@ mod tests {
                         S3:
                             bucket: "chroma"
                             credentials: Minio
+                            connect_timeout_ms: 5000
+                            request_timeout_ms: 1000
                     log:
                         Grpc:
                             host: "localhost"
@@ -438,6 +450,14 @@ mod tests {
             let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__MY_PORT", 50051);
             let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__STORAGE__S3__BUCKET", "buckets!");
             let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__STORAGE__S3__CREDENTIALS", "AWS");
+            let _ = jail.set_env(
+                "CHROMA_COMPACTION_SERVICE__STORAGE__S3__CONNECT_TIMEOUT_MS",
+                5000,
+            );
+            let _ = jail.set_env(
+                "CHROMA_COMPACTION_SERVICE__STORAGE__S3__REQUEST_TIMEOUT_MS",
+                1000,
+            );
             let _ = jail.create_file(
                 "chroma_config.yaml",
                 r#"
@@ -460,6 +480,8 @@ mod tests {
                         S3:
                             bucket: "chroma"
                             credentials: Minio
+                            connect_timeout_ms: 5000
+                            request_timeout_ms: 1000
                     log:
                         Grpc:
                             host: "localhost"
@@ -517,6 +539,8 @@ mod tests {
                         s.credentials,
                         crate::storage::config::S3CredentialsConfig::AWS
                     );
+                    assert_eq!(s.connect_timeout_ms, 5000);
+                    assert_eq!(s.request_timeout_ms, 1000);
                 }
                 _ => panic!("Invalid storage config"),
             }

diff --git a/rust/worker/src/storage/config.rs b/rust/worker/src/storage/config.rs
@@ -29,6 +29,8 @@ pub(crate) enum S3CredentialsConfig {
 pub(crate) struct S3StorageConfig {
     pub(crate) bucket: String,
     pub(crate) credentials: S3CredentialsConfig,
+    pub(crate) connect_timeout_ms: u64,
+    pub(crate) request_timeout_ms: u64,
 }
 
 #[derive(Deserialize, Debug)]

diff --git a/rust/worker/src/storage/s3.rs b/rust/worker/src/storage/s3.rs
@@ -12,12 +12,14 @@ use super::config::StorageConfig;
 use crate::config::Configurable;
 use crate::errors::ChromaError;
 use async_trait::async_trait;
+use aws_config::timeout::TimeoutConfigBuilder;
 use aws_sdk_s3;
 use aws_sdk_s3::error::ProvideErrorMetadata;
 use aws_sdk_s3::error::SdkError;
 use aws_sdk_s3::operation::create_bucket::CreateBucketError;
 use aws_smithy_types::byte_stream::ByteStream;
 use std::clone::Clone;
+use std::time::Duration;
 use thiserror::Error;
 use tokio::io::AsyncBufRead;
 
@@ -227,13 +229,18 @@ impl Configurable<StorageConfig> for S3Storage {
                             "loaded-from-env",
                         );
 
+                        let timeout_config_builder = TimeoutConfigBuilder::default()
+                            .connect_timeout(Duration::from_millis(s3_config.connect_timeout_ms))
+                            .read_timeout(Duration::from_millis(s3_config.request_timeout_ms));
+
                         // Set up s3 client
                         let config = aws_sdk_s3::config::Builder::new()
                             .endpoint_url("http://minio.chroma:9000".to_string())
                             .credentials_provider(cred)
                             .behavior_version_latest()
                             .region(aws_sdk_s3::config::Region::new("us-east-1"))
                             .force_path_style(true)
+                            .timeout_config(timeout_config_builder.build())
                             .build();
                         aws_sdk_s3::Client::from_conf(config)
                     }