[FEAT] Add flag to limit number of connections to S3 (#1360)

* Adds flag to S3Config that limits the number of connections to S3,
default to 25 (recommended in the
[aws-cpp-sdk](https://github.com/aws/aws-sdk-cpp/blob/c859aa7a1d85c5137f0416e4aa75b86ceb32e216/docs/ClientConfiguration_Parameters.md?plain=1#L46))

src/common/io-config/src/python.rs

@@ -11,6 +11,7 @@ use crate::config;
 ///     endpoint_url: URL to the S3 endpoint, defaults to endpoints to AWS
 ///     key_id: AWS Access Key ID, defaults to auto-detection from the current environment
 ///     access_key: AWS Secret Access Key, defaults to auto-detection from the current environment
+///     max_connections: Maximum number of connections to S3 at any time, defaults to 25
 ///     session_token: AWS Session Token, required only if `key_id` and `access_key` are temporary credentials
 ///     retry_initial_backoff_ms: Initial backoff duration in milliseconds for an S3 retry, defaults to 1000ms
 ///     connect_timeout_ms: Timeout duration to wait to make a connection to S3 in milliseconds, defaults to 60 seconds
@@ -142,6 +143,7 @@ impl S3Config {
         key_id: Option<String>,
         session_token: Option<String>,
         access_key: Option<String>,
+        max_connections: Option<u32>,
         retry_initial_backoff_ms: Option<u64>,
         connect_timeout_ms: Option<u64>,
         read_timeout_ms: Option<u64>,
@@ -157,6 +159,7 @@ impl S3Config {
                 key_id: key_id.or(def.key_id),
                 session_token: session_token.or(def.session_token),
                 access_key: access_key.or(def.access_key),
+                max_connections: max_connections.unwrap_or(def.max_connections),
                 retry_initial_backoff_ms: retry_initial_backoff_ms
                     .unwrap_or(def.retry_initial_backoff_ms),
                 connect_timeout_ms: connect_timeout_ms.unwrap_or(def.connect_timeout_ms),
@@ -202,6 +205,12 @@ impl S3Config {
         Ok(self.config.access_key.clone())
     }
 
+    /// AWS max connections
+    #[getter]
+    pub fn max_connections(&self) -> PyResult<u32> {
+        Ok(self.config.max_connections)
+    }
+
     /// AWS Retry Initial Backoff Time in Milliseconds
     #[getter]
     pub fn retry_initial_backoff_ms(&self) -> PyResult<u64> {

src/common/io-config/src/s3.rs

@@ -11,6 +11,7 @@ pub struct S3Config {
     pub key_id: Option<String>,
     pub session_token: Option<String>,
     pub access_key: Option<String>,
+    pub max_connections: u32,
     pub retry_initial_backoff_ms: u64,
     pub connect_timeout_ms: u64,
     pub read_timeout_ms: u64,
@@ -27,6 +28,7 @@ impl Default for S3Config {
             key_id: None,
             session_token: None,
             access_key: None,
+            max_connections: 25,
             retry_initial_backoff_ms: 1000,
             connect_timeout_ms: 60_000,
             read_timeout_ms: 60_000,
@@ -47,6 +49,7 @@ impl Display for S3Config {
     key_id: {:?}
     session_token: {:?},
     access_key: {:?}
+    max_connections: {},
     retry_initial_backoff_ms: {},
     connect_timeout_ms: {},
     read_timeout_ms: {},
@@ -59,6 +62,7 @@ impl Display for S3Config {
             self.session_token,
             self.access_key,
             self.retry_initial_backoff_ms,
+            self.max_connections,
             self.connect_timeout_ms,
             self.read_timeout_ms,
             self.num_tries,

src/daft-io/src/azure_blob.rs

@@ -152,7 +152,7 @@ impl ObjectSource for AzureBlobSource {
                 .into_error(e)
                 .into()
             });
-        Ok(GetResult::Stream(stream.boxed(), None))
+        Ok(GetResult::Stream(stream.boxed(), None, None))
     }
 
     async fn get_size(&self, uri: &str) -> super::Result<usize> {

src/daft-io/src/google_cloud.rs

@@ -146,7 +146,7 @@ impl GCSClientWrapper {
                     .into_error(e)
                     .into()
                 });
-                Ok(GetResult::Stream(response.boxed(), size))
+                Ok(GetResult::Stream(response.boxed(), size, None))
             }
             GCSClientWrapper::S3Compat(client) => {
                 let uri = format!("s3://{}/{}", bucket, key);

src/daft-io/src/http.rs

@@ -119,7 +119,7 @@ impl ObjectSource for HttpSource {
             .into_error(e)
             .into()
         });
-        Ok(GetResult::Stream(stream.boxed(), size_bytes))
+        Ok(GetResult::Stream(stream.boxed(), size_bytes, None))
     }
 
     async fn get_size(&self, uri: &str) -> super::Result<usize> {

src/daft-io/src/object_io.rs

@@ -4,12 +4,17 @@ use async_trait::async_trait;
 use bytes::Bytes;
 use futures::stream::{BoxStream, Stream};
 use futures::StreamExt;
+use tokio::sync::OwnedSemaphorePermit;
 
 use crate::local::{collect_file, LocalFile};
 
 pub enum GetResult {
     File(LocalFile),
-    Stream(BoxStream<'static, super::Result<Bytes>>, Option<usize>),
+    Stream(
+        BoxStream<'static, super::Result<Bytes>>,
+        Option<usize>,
+        Option<OwnedSemaphorePermit>,
+    ),
 }
 
 async fn collect_bytes<S>(mut stream: S, size_hint: Option<usize>) -> super::Result<Bytes>
@@ -40,7 +45,7 @@ impl GetResult {
         use GetResult::*;
         match self {
             File(f) => collect_file(f).await,
-            Stream(stream, size) => collect_bytes(stream, size).await,
+            Stream(stream, size, _permit) => collect_bytes(stream, size).await,
         }
     }
 }

src/daft-io/src/s3_like.rs

@@ -5,6 +5,7 @@ use aws_smithy_async::rt::sleep::TokioSleep;
 use reqwest::StatusCode;
 use s3::operation::head_object::HeadObjectError;
 use s3::operation::list_objects_v2::ListObjectsV2Error;
+use tokio::sync::{OwnedSemaphorePermit, SemaphorePermit};
 
 use crate::object_io::{FileMetadata, FileType, LSResult};
 use crate::{InvalidArgumentSnafu, SourceType};
@@ -34,6 +35,7 @@ use std::sync::Arc;
 use std::time::Duration;
 pub(crate) struct S3LikeSource {
     region_to_client_map: tokio::sync::RwLock<HashMap<Region, Arc<s3::Client>>>,
+    connection_pool_sema: Arc<tokio::sync::Semaphore>,
     default_region: Region,
     s3_config: S3Config,
     anonymous: bool,
@@ -91,6 +93,9 @@ enum Error {
     #[snafu(display("Unable to create http client. {}", source))]
     UnableToCreateClient { source: reqwest::Error },
 
+    #[snafu(display("Unable to grab semaphore. {}", source))]
+    UnableToGrabSemaphore { source: tokio::sync::AcquireError },
+
     #[snafu(display(
         "Unable to parse data as Utf8 while reading header for file: {path}. {source}"
     ))]
@@ -252,6 +257,9 @@ async fn build_client(config: &S3Config) -> super::Result<S3LikeSource> {
     client_map.insert(default_region.clone(), client.into());
     Ok(S3LikeSource {
         region_to_client_map: tokio::sync::RwLock::new(client_map),
+        connection_pool_sema: Arc::new(tokio::sync::Semaphore::new(
+            config.max_connections as usize,
+        )),
         s3_config: config.clone(),
         default_region,
         anonymous,
@@ -290,6 +298,7 @@ impl S3LikeSource {
     #[async_recursion]
     async fn _get_impl(
         &self,
+        permit: OwnedSemaphorePermit,
         uri: &str,
         range: Option<Range<usize>>,
         region: &Region,
@@ -357,7 +366,11 @@ impl S3LikeSource {
                             .into()
                         })
                         .boxed();
-                    Ok(GetResult::Stream(stream, Some(v.content_length as usize)))
+                    Ok(GetResult::Stream(
+                        stream,
+                        Some(v.content_length as usize),
+                        Some(permit),
+                    ))
                 }
 
                 Err(SdkError::ServiceError(err)) => {
@@ -378,7 +391,7 @@ impl S3LikeSource {
 
                             let new_region = Region::new(region_name);
                             log::debug!("S3 Region of {uri} different than client {:?} vs {:?} Attempting GET in that region with new client", new_region, region);
-                            self._get_impl(uri, range, &new_region).await
+                            self._get_impl(permit, uri, range, &new_region).await
                         }
                         _ => Err(UnableToOpenFileSnafu { path: uri }
                             .into_error(SdkError::ServiceError(err))
@@ -393,7 +406,12 @@ impl S3LikeSource {
     }
 
     #[async_recursion]
-    async fn _head_impl(&self, uri: &str, region: &Region) -> super::Result<usize> {
+    async fn _head_impl(
+        &self,
+        _permit: SemaphorePermit<'async_recursion>,
+        uri: &str,
+        region: &Region,
+    ) -> super::Result<usize> {
         let parsed = url::Url::parse(uri).with_context(|_| InvalidUrlSnafu { path: uri })?;
 
         let bucket = match parsed.host_str() {
@@ -456,7 +474,7 @@ impl S3LikeSource {
 
                             let new_region = Region::new(region_name);
                             log::debug!("S3 Region of {uri} different than client {:?} vs {:?} Attempting HEAD in that region with new client", new_region, region);
-                            self._head_impl(uri, &new_region).await
+                            self._head_impl(_permit, uri, &new_region).await
                         }
                         _ => Err(UnableToHeadFileSnafu { path: uri }
                             .into_error(SdkError::ServiceError(err))
@@ -473,6 +491,7 @@ impl S3LikeSource {
     #[async_recursion]
     async fn _list_impl(
         &self,
+        _permit: SemaphorePermit<'async_recursion>,
         bucket: &str,
         key: &str,
         delimiter: String,
@@ -572,6 +591,7 @@ impl S3LikeSource {
                         let new_region = Region::new(region_name);
                         log::debug!("S3 Region of {uri} different than client {:?} vs {:?} Attempting List in that region with new client", new_region, region);
                         self._list_impl(
+                            _permit,
                             bucket,
                             key,
                             delimiter,
@@ -595,11 +615,23 @@ impl S3LikeSource {
 #[async_trait]
 impl ObjectSource for S3LikeSource {
     async fn get(&self, uri: &str, range: Option<Range<usize>>) -> super::Result<GetResult> {
-        self._get_impl(uri, range, &self.default_region).await
+        let permit = self
+            .connection_pool_sema
+            .clone()
+            .acquire_owned()
+            .await
+            .context(UnableToGrabSemaphoreSnafu)?;
+        self._get_impl(permit, uri, range, &self.default_region)
+            .await
     }
 
     async fn get_size(&self, uri: &str) -> super::Result<usize> {
-        self._head_impl(uri, &self.default_region).await
+        let permit = self
+            .connection_pool_sema
+            .acquire()
+            .await
+            .context(UnableToGrabSemaphoreSnafu)?;
+        self._head_impl(permit, uri, &self.default_region).await
     }
     async fn ls(
         &self,
@@ -621,21 +653,34 @@ impl ObjectSource for S3LikeSource {
         if let Some(key) = key.strip_prefix('/') {
             // assume its a directory first
             let key = format!("{}/", key.strip_suffix('/').unwrap_or(key));
-            let lsr = self
-                ._list_impl(
+            let lsr = {
+                let permit = self
+                    .connection_pool_sema
+                    .acquire()
+                    .await
+                    .context(UnableToGrabSemaphoreSnafu)?;
+                self._list_impl(
+                    permit,
                     bucket,
                     &key,
                     delimiter.into(),
                     continuation_token.map(String::from),
                     &self.default_region,
                 )
-                .await?;
+                .await?
+            };
             if lsr.files.is_empty() && key.contains('/') {
+                let permit = self
+                    .connection_pool_sema
+                    .acquire()
+                    .await
+                    .context(UnableToGrabSemaphoreSnafu)?;
                 // Might be a File
                 let split = key.rsplit_once('/');
                 let (new_key, _) = split.unwrap();
                 let mut lsr = self
                     ._list_impl(
+                        permit,
                         bucket,
                         new_key,
                         delimiter.into(),