Add region back in file_store::Settings and add tracking of new clien… (#1073)

* Add region back in file_store::Settings and add tracking of new clients to reduce aws credential caches

* Update aws_local and clippy fixes

* Re-add comment for force_path_style

* Add debug statements

* Add BucketClient to file_store and update mobile-verifier to use it

* only set region if Option is Some

* Updated aws_local::AwsLocal::new to take region as a parameter

Author: bbalser

aws_local/Cargo.toml

@@ -16,4 +16,4 @@ prost = { workspace = true }
 anyhow = { workspace = true }
 uuid = { workspace = true }
 tempfile = { workspace = true }
-file-store = { path = "../file_store", features = ["local"] }
+file-store = { path = "../file_store"}

aws_local/src/lib.rs

@@ -53,8 +53,9 @@ impl AwsLocal {
         Client::from_conf(s3_config.build())
     }
 
-    pub async fn new(endpoint: &str, bucket: &str) -> AwsLocal {
+    pub async fn new(region: &str, endpoint: &str, bucket: &str) -> AwsLocal {
         let settings = Settings {
+            region: Some(region.into()),
             endpoint: Some(endpoint.into()),
             access_key_id: Some("random".into()),
             secret_access_key: Some("random2".into()),

file_store/Cargo.toml

@@ -52,5 +52,4 @@ tempfile = { workspace = true }
 
 [features]
 default = ["sqlx-postgres"]
-local = []
 sqlx-postgres = ["sqlx/postgres"]

file_store/src/bucket_client.rs

@@ -0,0 +1,74 @@
+use std::path::Path;
+
+use aws_sdk_s3::primitives::ByteStream;
+use chrono::{DateTime, Utc};
+
+use crate::{error::Result, BytesMutStream, FileInfo, FileInfoStream};
+
+#[derive(Clone, Debug)]
+pub struct BucketClient {
+    pub client: aws_sdk_s3::Client,
+    pub bucket: String,
+}
+
+impl BucketClient {
+    pub async fn new(
+        bucket: String,
+        region: Option<String>,
+        endpoint: Option<String>,
+        access_key_id: Option<String>,
+        secret_access_key: Option<String>,
+    ) -> Self {
+        let client = crate::new_client(region, endpoint, access_key_id, secret_access_key).await;
+        Self { client, bucket }
+    }
+
+    pub fn list_files<A, B>(&self, prefix: impl Into<String>, after: A, before: B) -> FileInfoStream
+    where
+        A: Into<Option<DateTime<Utc>>> + Copy,
+        B: Into<Option<DateTime<Utc>>> + Copy,
+    {
+        crate::list_files(&self.client, self.bucket.clone(), prefix, after, before)
+    }
+
+    pub async fn list_all_files<A, B>(
+        &self,
+        prefix: impl Into<String>,
+        after: A,
+        before: B,
+    ) -> Result<Vec<FileInfo>>
+    where
+        A: Into<Option<DateTime<Utc>>> + Copy,
+        B: Into<Option<DateTime<Utc>>> + Copy,
+    {
+        crate::list_all_files(&self.client, self.bucket.clone(), prefix, after, before).await
+    }
+
+    pub async fn put_file(&self, file: &Path) -> Result {
+        crate::put_file(&self.client, self.bucket.clone(), file).await
+    }
+
+    pub async fn remove_file(&self, key: impl Into<String>) -> Result {
+        crate::remove_file(&self.client, self.bucket.clone(), key).await
+    }
+
+    pub async fn get_raw_file(&self, key: impl Into<String>) -> Result<ByteStream> {
+        crate::get_raw_file(&self.client, self.bucket.clone(), key).await
+    }
+
+    pub async fn get_file(&self, key: impl Into<String>) -> Result<BytesMutStream> {
+        crate::get_file(&self.client, self.bucket.clone(), key).await
+    }
+
+    pub fn source_files(&self, infos: FileInfoStream) -> BytesMutStream {
+        crate::source_files(&self.client, self.bucket.clone(), infos)
+    }
+
+    pub fn source_files_unordered(&self, workers: usize, infos: FileInfoStream) -> BytesMutStream {
+        crate::source_files_unordered(&self.client, self.bucket.clone(), workers, infos)
+    }
+
+    pub async fn stream_single_file(&self, file_info: FileInfo) -> Result<BytesMutStream> {
+        crate::stream_single_file(&self.client, self.bucket.clone(), file_info).await
+    }
+}

file_store/src/file_info_poller.rs

@@ -1,4 +1,4 @@
-use crate::{traits::MsgDecode, Error, FileInfo, Result};
+use crate::{traits::MsgDecode, BucketClient, Error, FileInfo, Result};
 use aws_sdk_s3::primitives::ByteStream;
 use chrono::{DateTime, Utc};
 use derive_builder::Builder;
@@ -173,6 +173,13 @@ impl<Message, State, Parser>
     pub fn file_store(self, client: crate::Client, bucket: impl Into<String>) -> Self {
         self.store(FileStoreInfoPollerStore::new(client, bucket))
     }
+
+    pub fn bucket_client(self, bucket_client: BucketClient) -> Self {
+        self.store(FileStoreInfoPollerStore::new(
+            bucket_client.client,
+            bucket_client.bucket,
+        ))
+    }
 }
 
 #[derive(Clone)]

file_store/src/file_source.rs

@@ -100,6 +100,7 @@ mod test {
         //
 
         let settings = Settings {
+            region: None,
             endpoint: None,
             access_key_id: None,
             secret_access_key: None,

file_store/src/file_upload.rs

@@ -1,4 +1,4 @@
-use crate::{Error, Result};
+use crate::{BucketClient, Error, Result};
 use futures::{StreamExt, TryFutureExt};
 use std::{
     path::{Path, PathBuf},
@@ -43,6 +43,10 @@ impl FileUpload {
         )
     }
 
+    pub async fn from_bucket_client(bucket_client: BucketClient) -> (Self, FileUploadServer) {
+        Self::new(bucket_client.client.clone(), bucket_client.bucket.clone()).await
+    }
+
     pub async fn upload_file(&self, file: &Path) -> Result {
         self.sender
             .send(file.to_path_buf())

file_store/src/lib.rs

@@ -3,7 +3,9 @@ extern crate tls_init;
 mod error;
 
 pub use error::{DecodeError, EncodeError, Error, Result};
+use tokio::sync::Mutex;
 
+pub mod bucket_client;
 pub mod file_info;
 pub mod file_info_poller;
 pub mod file_sink;
@@ -12,7 +14,7 @@ pub mod file_upload;
 mod settings;
 pub mod traits;
 
-use std::path::Path;
+use std::{collections::HashMap, path::Path, sync::OnceLock};
 
 use aws_config::BehaviorVersion;
 use aws_sdk_s3::primitives::ByteStream;
@@ -22,52 +24,82 @@ use chrono::{DateTime, Utc};
 pub use file_info::FileInfo;
 pub use file_sink::{FileSink, FileSinkBuilder};
 
+pub use bucket_client::BucketClient;
 use futures::{
     future,
     stream::{self, BoxStream},
     FutureExt, StreamExt, TryFutureExt, TryStreamExt,
 };
-pub use settings::Settings;
+pub use settings::{BucketSettings, Settings};
 
 pub type Client = aws_sdk_s3::Client;
 pub type Stream<T> = BoxStream<'static, Result<T>>;
 pub type FileInfoStream = Stream<FileInfo>;
 pub type BytesMutStream = Stream<BytesMut>;
 
+static CLIENT_MAP: OnceLock<Mutex<HashMap<ClientKey, Client>>> = OnceLock::new();
+
+#[derive(PartialEq, Eq, Hash, Debug)]
+struct ClientKey {
+    region: Option<String>,
+    endpoint: Option<String>,
+    access_key_id: Option<String>,
+    secret_access_key: Option<String>,
+}
+
 pub async fn new_client(
+    region: Option<String>,
     endpoint: Option<String>,
     _access_key_id: Option<String>,
     _secret_access_key: Option<String>,
-) -> Client {
+) -> aws_sdk_s3::Client {
+    let mut client_map = CLIENT_MAP
+        .get_or_init(|| Mutex::new(HashMap::new()))
+        .lock()
+        .await;
+
+    let key = ClientKey {
+        region: region.clone(),
+        endpoint: endpoint.clone(),
+        access_key_id: _access_key_id.clone(),
+        secret_access_key: _secret_access_key.clone(),
+    };
+
+    if let Some(client) = client_map.get(&key) {
+        tracing::debug!(params = ?key, "Using existing file-store s3 client");
+        return client.clone();
+    }
+
     let config = aws_config::defaults(BehaviorVersion::latest()).load().await;
 
     let mut s3_config = aws_sdk_s3::config::Builder::from(&config);
-    if let Some(endpoint) = endpoint {
-        s3_config = s3_config.endpoint_url(endpoint);
+
+    if let Some(region_str) = region {
+        s3_config = s3_config.region(aws_config::Region::new(region_str));
     }
 
-    #[cfg(feature = "local")]
-    {
+    if let Some(endpoint) = endpoint {
+        s3_config = s3_config.endpoint_url(endpoint);
         // NOTE(mj): If you see something like a DNS error, this is probably
         // the culprit. Need to find a way to make this configurable. It
         // would be nice to allow the "local" feature to be active, but not
         // enforce path style.
         s3_config = s3_config.force_path_style(true);
+    }
 
-        // Set a default region for local development (MinIO doesn't care about the region)
-        s3_config = s3_config.region(aws_config::Region::new("us-east-1"));
-
-        if let Some((access_key_id, secret_access_key)) = _access_key_id.zip(_secret_access_key) {
-            let creds = aws_sdk_s3::config::Credentials::builder()
-                .provider_name("Static")
-                .access_key_id(access_key_id)
-                .secret_access_key(secret_access_key);
+    if let Some((access_key_id, secret_access_key)) = _access_key_id.zip(_secret_access_key) {
+        let creds = aws_sdk_s3::config::Credentials::builder()
+            .provider_name("Static")
+            .access_key_id(access_key_id)
+            .secret_access_key(secret_access_key);
 
-            s3_config = s3_config.credentials_provider(creds.build());
-        }
+        s3_config = s3_config.credentials_provider(creds.build());
     }
 
-    Client::from_conf(s3_config.build())
+    tracing::debug!(params = ?key, "Creating new file-store s3 client");
+    let client = Client::from_conf(s3_config.build());
+    client_map.insert(key, client.clone());
+    client
 }
 
 pub fn list_files<A, B>(

file_store/src/settings.rs

@@ -2,8 +2,11 @@ use config::{Config, File};
 use serde::{Deserialize, Serialize};
 use std::path::Path;
 
+use crate::bucket_client::BucketClient;
+
 #[derive(Debug, Serialize, Deserialize, Clone, Default)]
 pub struct Settings {
+    pub region: Option<String>,
     /// Optional api endpoint for the bucket. Default none
     pub endpoint: Option<String>,
     /// Should only be used for local testing
@@ -13,6 +16,13 @@ pub struct Settings {
     pub secret_access_key: Option<String>,
 }
 
+#[derive(Debug, Serialize, Deserialize, Clone, Default)]
+pub struct BucketSettings {
+    bucket: String,
+    #[serde(flatten)]
+    settings: Settings,
+}
+
 impl Settings {
     /// Load Settings from a given path.
     ///
@@ -26,10 +36,31 @@ impl Settings {
 
     pub async fn connect(&self) -> crate::Client {
         crate::new_client(
+            self.region.clone(),
             self.endpoint.clone(),
             self.access_key_id.clone(),
             self.secret_access_key.clone(),
         )
         .await
     }
 }
+
+impl BucketSettings {
+    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, config::ConfigError> {
+        Config::builder()
+            .add_source(File::with_name(&path.as_ref().to_string_lossy()))
+            .build()
+            .and_then(|config| config.try_deserialize())
+    }
+
+    pub async fn connect(&self) -> BucketClient {
+        BucketClient::new(
+            self.bucket.clone(),
+            self.settings.region.clone(),
+            self.settings.endpoint.clone(),
+            self.settings.access_key_id.clone(),
+            self.settings.secret_access_key.clone(),
+        )
+        .await
+    }
+}

iot_verifier/src/main.rs

@@ -101,8 +101,7 @@ impl Server {
         // *
         // setup the price tracker requirements
         // *
-        let (price_tracker, price_daemon) =
-            PriceTracker::new(&settings.price_tracker, file_store_client.clone()).await?;
+        let (price_tracker, price_daemon) = PriceTracker::new(&settings.price_tracker).await?;
 
         // *
         // setup the loader requirements