Skip to content

Commit 0820c90

Browse files
memoize s3 client (#5377)
* cache S3Client in factory * remove unused constructor for s3 storage
1 parent aac8b49 commit 0820c90

File tree

2 files changed

+40
-22
lines changed

2 files changed

+40
-22
lines changed

quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs

Lines changed: 19 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ fn get_region(s3_storage_config: &S3StorageConfig) -> Option<Region> {
127127
})
128128
}
129129

130-
async fn create_s3_client(s3_storage_config: &S3StorageConfig) -> S3Client {
130+
pub async fn create_s3_client(s3_storage_config: &S3StorageConfig) -> S3Client {
131131
let aws_config = get_aws_config().await;
132132
let credentials_provider =
133133
get_credentials_provider(s3_storage_config).or(aws_config.credentials_provider());
@@ -155,41 +155,40 @@ async fn create_s3_client(s3_storage_config: &S3StorageConfig) -> S3Client {
155155
}
156156

157157
impl S3CompatibleObjectStorage {
158-
/// Creates an object storage given a region and a bucket name.
159-
pub async fn new(
158+
/// Creates an object storage given a region and an uri.
159+
pub async fn from_uri(
160160
s3_storage_config: &S3StorageConfig,
161-
uri: Uri,
162-
bucket: String,
161+
uri: &Uri,
163162
) -> Result<Self, StorageResolverError> {
164163
let s3_client = create_s3_client(s3_storage_config).await;
164+
Self::from_uri_and_client(s3_storage_config, uri, s3_client).await
165+
}
166+
167+
/// Creates an object storage given a region, an uri and an S3 client.
168+
pub async fn from_uri_and_client(
169+
s3_storage_config: &S3StorageConfig,
170+
uri: &Uri,
171+
s3_client: S3Client,
172+
) -> Result<Self, StorageResolverError> {
173+
let (bucket, prefix) = parse_s3_uri(uri).ok_or_else(|| {
174+
let message = format!("failed to extract bucket name from S3 URI: {uri}");
175+
StorageResolverError::InvalidUri(message)
176+
})?;
165177
let retry_params = RetryParams::aggressive();
166178
let disable_multi_object_delete = s3_storage_config.disable_multi_object_delete;
167179
let disable_multipart_upload = s3_storage_config.disable_multipart_upload;
168180
Ok(Self {
169181
s3_client,
170-
uri,
182+
uri: uri.clone(),
171183
bucket,
172-
prefix: PathBuf::new(),
184+
prefix,
173185
multipart_policy: MultiPartPolicy::default(),
174186
retry_params,
175187
disable_multi_object_delete,
176188
disable_multipart_upload,
177189
})
178190
}
179191

180-
/// Creates an object storage given a region and an uri.
181-
pub async fn from_uri(
182-
s3_storage_config: &S3StorageConfig,
183-
uri: &Uri,
184-
) -> Result<Self, StorageResolverError> {
185-
let (bucket, prefix) = parse_s3_uri(uri).ok_or_else(|| {
186-
let message = format!("failed to extract bucket name from S3 URI: {uri}");
187-
StorageResolverError::InvalidUri(message)
188-
})?;
189-
let storage = Self::new(s3_storage_config, uri.clone(), bucket).await?;
190-
Ok(storage.with_prefix(prefix))
191-
}
192-
193192
/// Sets a specific for all buckets.
194193
///
195194
/// This method overrides any existing prefix. (It does NOT

quickwit/quickwit-storage/src/object_storage/s3_compatible_storage_resolver.rs

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,22 +20,34 @@
2020
use std::sync::Arc;
2121

2222
use async_trait::async_trait;
23+
use aws_sdk_s3::Client as S3Client;
2324
use quickwit_common::uri::Uri;
2425
use quickwit_config::{S3StorageConfig, StorageBackend};
26+
use tokio::sync::OnceCell;
2527

28+
use super::s3_compatible_storage::create_s3_client;
2629
use crate::{
2730
DebouncedStorage, S3CompatibleObjectStorage, Storage, StorageFactory, StorageResolverError,
2831
};
2932

3033
/// S3 compatible object storage resolver.
3134
pub struct S3CompatibleObjectStorageFactory {
3235
storage_config: S3StorageConfig,
36+
// we cache the S3Client so we don't rebuild one every time we build a new Storage (for
37+
// every search query).
38+
// We don't build it in advance because we don't know if this factory is one that will
39+
// end up being used, or if something like azure, gcs, or even local files, will be used
40+
// instead.
41+
s3_client: OnceCell<S3Client>,
3342
}
3443

3544
impl S3CompatibleObjectStorageFactory {
3645
/// Creates a new S3-compatible storage factory.
3746
pub fn new(storage_config: S3StorageConfig) -> Self {
38-
Self { storage_config }
47+
Self {
48+
storage_config,
49+
s3_client: OnceCell::new(),
50+
}
3951
}
4052
}
4153

@@ -46,7 +58,14 @@ impl StorageFactory for S3CompatibleObjectStorageFactory {
4658
}
4759

4860
async fn resolve(&self, uri: &Uri) -> Result<Arc<dyn Storage>, StorageResolverError> {
49-
let storage = S3CompatibleObjectStorage::from_uri(&self.storage_config, uri).await?;
61+
let s3_client = self
62+
.s3_client
63+
.get_or_init(|| create_s3_client(&self.storage_config))
64+
.await
65+
.clone();
66+
let storage =
67+
S3CompatibleObjectStorage::from_uri_and_client(&self.storage_config, uri, s3_client)
68+
.await?;
5069
Ok(Arc::new(DebouncedStorage::new(storage)))
5170
}
5271
}

0 commit comments

Comments
 (0)