Add with_tokio_runtime to HTTP stores

Author: tustvold

object_store/src/aws/client.rs

@@ -19,13 +19,13 @@ use crate::aws::checksum::Checksum;
 use crate::aws::credential::{AwsCredential, CredentialExt, CredentialProvider};
 use crate::aws::STRICT_PATH_ENCODE_SET;
 use crate::client::pagination::stream_paginated;
-use crate::client::retry::RetryExt;
+use crate::client::retry::{ClientConfig, RetryExt};
 use crate::multipart::UploadPart;
 use crate::path::DELIMITER;
 use crate::util::{format_http_range, format_prefix};
 use crate::{
     BoxStream, ClientOptions, ListResult, MultipartId, ObjectMeta, Path, Result,
-    RetryConfig, StreamExt,
+    StreamExt,
 };
 use base64::prelude::BASE64_STANDARD;
 use base64::Engine;
@@ -208,7 +208,7 @@ pub struct S3Config {
     pub bucket: String,
     pub bucket_endpoint: String,
     pub credentials: Box<dyn CredentialProvider>,
-    pub retry_config: RetryConfig,
+    pub client_config: ClientConfig,
     pub client_options: ClientOptions,
     pub sign_payload: bool,
     pub checksum: Option<Checksum>,
@@ -271,7 +271,7 @@ impl S3Client {
                 self.config.sign_payload,
                 None,
             )
-            .send_retry(&self.config.retry_config)
+            .send_retry(&self.config.client_config)
             .await
             .context(GetRequestSnafu {
                 path: path.as_ref(),
@@ -317,7 +317,7 @@ impl S3Client {
                 self.config.sign_payload,
                 payload_sha256,
             )
-            .send_retry(&self.config.retry_config)
+            .send_retry(&self.config.client_config)
             .await
             .context(PutRequestSnafu {
                 path: path.as_ref(),
@@ -345,7 +345,7 @@ impl S3Client {
                 self.config.sign_payload,
                 None,
             )
-            .send_retry(&self.config.retry_config)
+            .send_retry(&self.config.client_config)
             .await
             .context(DeleteRequestSnafu {
                 path: path.as_ref(),
@@ -370,7 +370,7 @@ impl S3Client {
                 self.config.sign_payload,
                 None,
             )
-            .send_retry(&self.config.retry_config)
+            .send_retry(&self.config.client_config)
             .await
             .context(CopyRequestSnafu {
                 path: from.as_ref(),
@@ -422,7 +422,7 @@ impl S3Client {
                 self.config.sign_payload,
                 None,
             )
-            .send_retry(&self.config.retry_config)
+            .send_retry(&self.config.client_config)
             .await
             .context(ListRequestSnafu)?
             .bytes()
@@ -476,7 +476,7 @@ impl S3Client {
                 self.config.sign_payload,
                 None,
             )
-            .send_retry(&self.config.retry_config)
+            .send_retry(&self.config.client_config)
             .await
             .context(CreateMultipartRequestSnafu)?
             .bytes()
@@ -521,7 +521,7 @@ impl S3Client {
                 self.config.sign_payload,
                 None,
             )
-            .send_retry(&self.config.retry_config)
+            .send_retry(&self.config.client_config)
             .await
             .context(CompleteMultipartRequestSnafu)?;
 

object_store/src/aws/credential.rs

@@ -16,10 +16,10 @@
 // under the License.
 
 use crate::aws::STRICT_ENCODE_SET;
-use crate::client::retry::RetryExt;
+use crate::client::retry::{ClientConfig, RetryExt};
 use crate::client::token::{TemporaryToken, TokenCache};
 use crate::util::hmac_sha256;
-use crate::{Result, RetryConfig};
+use crate::Result;
 use bytes::Buf;
 use chrono::{DateTime, Utc};
 use futures::future::BoxFuture;
@@ -328,7 +328,7 @@ impl CredentialProvider for StaticCredentialProvider {
 pub struct InstanceCredentialProvider {
     pub cache: TokenCache<Arc<AwsCredential>>,
     pub client: Client,
-    pub retry_config: RetryConfig,
+    pub client_config: ClientConfig,
     pub imdsv1_fallback: bool,
     pub metadata_endpoint: String,
 }
@@ -338,7 +338,7 @@ impl CredentialProvider for InstanceCredentialProvider {
         Box::pin(self.cache.get_or_insert_with(|| {
             instance_creds(
                 &self.client,
-                &self.retry_config,
+                &self.client_config,
                 &self.metadata_endpoint,
                 self.imdsv1_fallback,
             )
@@ -361,15 +361,15 @@ pub struct WebIdentityProvider {
     pub session_name: String,
     pub endpoint: String,
     pub client: Client,
-    pub retry_config: RetryConfig,
+    pub client_config: ClientConfig,
 }
 
 impl CredentialProvider for WebIdentityProvider {
     fn get_credential(&self) -> BoxFuture<'_, Result<Arc<AwsCredential>>> {
         Box::pin(self.cache.get_or_insert_with(|| {
             web_identity(
                 &self.client,
-                &self.retry_config,
+                &self.client_config,
                 &self.token_path,
                 &self.role_arn,
                 &self.session_name,
@@ -405,7 +405,7 @@ impl From<InstanceCredentials> for AwsCredential {
 /// <https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html#instance-metadata-security-credentials>
 async fn instance_creds(
     client: &Client,
-    retry_config: &RetryConfig,
+    config: &ClientConfig,
     endpoint: &str,
     imdsv1_fallback: bool,
 ) -> Result<TemporaryToken<Arc<AwsCredential>>, StdError> {
@@ -417,7 +417,7 @@ async fn instance_creds(
     let token_result = client
         .request(Method::PUT, token_url)
         .header("X-aws-ec2-metadata-token-ttl-seconds", "600") // 10 minute TTL
-        .send_retry(retry_config)
+        .send_retry(config)
         .await;
 
     let token = match token_result {
@@ -438,7 +438,7 @@ async fn instance_creds(
         role_request = role_request.header(AWS_EC2_METADATA_TOKEN_HEADER, token);
     }
 
-    let role = role_request.send_retry(retry_config).await?.text().await?;
+    let role = role_request.send_retry(config).await?.text().await?;
 
     let creds_url = format!("{endpoint}/{CREDENTIALS_PATH}/{role}");
     let mut creds_request = client.request(Method::GET, creds_url);
@@ -447,7 +447,7 @@ async fn instance_creds(
     }
 
     let creds: InstanceCredentials =
-        creds_request.send_retry(retry_config).await?.json().await?;
+        creds_request.send_retry(config).await?.json().await?;
 
     let now = Utc::now();
     let ttl = (creds.expiration - now).to_std().unwrap_or_default();
@@ -491,7 +491,7 @@ impl From<AssumeRoleCredentials> for AwsCredential {
 /// <https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts-technical-overview.html>
 async fn web_identity(
     client: &Client,
-    retry_config: &RetryConfig,
+    config: &ClientConfig,
     token_path: &str,
     role_arn: &str,
     session_name: &str,
@@ -510,7 +510,7 @@ async fn web_identity(
             ("Version", "2011-06-15"),
             ("WebIdentityToken", &token),
         ])
-        .send_retry(retry_config)
+        .send_retry(config)
         .await?
         .bytes()
         .await?;
@@ -722,7 +722,7 @@ mod tests {
         // For example https://github.com/aws/amazon-ec2-metadata-mock
         let endpoint = env::var("EC2_METADATA_ENDPOINT").unwrap();
         let client = Client::new();
-        let retry_config = RetryConfig::default();
+        let config = ClientConfig::default();
 
         // Verify only allows IMDSv2
         let resp = client
@@ -737,7 +737,7 @@ mod tests {
             "Ensure metadata endpoint is set to only allow IMDSv2"
         );
 
-        let creds = instance_creds(&client, &retry_config, &endpoint, false)
+        let creds = instance_creds(&client, &config, &endpoint, false)
             .await
             .unwrap();
 
@@ -762,7 +762,7 @@ mod tests {
 
         let endpoint = server.url();
         let client = Client::new();
-        let retry_config = RetryConfig::default();
+        let config = ClientConfig::default();
 
         // Test IMDSv2
         server.push_fn(|req| {
@@ -788,7 +788,7 @@ mod tests {
             Response::new(Body::from(r#"{"AccessKeyId":"KEYID","Code":"Success","Expiration":"2022-08-30T10:51:04Z","LastUpdated":"2022-08-30T10:21:04Z","SecretAccessKey":"SECRET","Token":"TOKEN","Type":"AWS-HMAC"}"#))
         });
 
-        let creds = instance_creds(&client, &retry_config, endpoint, true)
+        let creds = instance_creds(&client, &config, endpoint, true)
             .await
             .unwrap();
 
@@ -821,7 +821,7 @@ mod tests {
             Response::new(Body::from(r#"{"AccessKeyId":"KEYID","Code":"Success","Expiration":"2022-08-30T10:51:04Z","LastUpdated":"2022-08-30T10:21:04Z","SecretAccessKey":"SECRET","Token":"TOKEN","Type":"AWS-HMAC"}"#))
         });
 
-        let creds = instance_creds(&client, &retry_config, endpoint, true)
+        let creds = instance_creds(&client, &config, endpoint, true)
             .await
             .unwrap();
 
@@ -838,7 +838,7 @@ mod tests {
         );
 
         // Should fail
-        instance_creds(&client, &retry_config, endpoint, false)
+        instance_creds(&client, &config, endpoint, false)
             .await
             .unwrap_err();
     }

object_store/src/aws/mod.rs

@@ -44,6 +44,7 @@ use std::ops::Range;
 use std::str::FromStr;
 use std::sync::Arc;
 use tokio::io::AsyncWrite;
+use tokio::runtime::Handle;
 use tracing::info;
 use url::Url;
 
@@ -53,6 +54,7 @@ use crate::aws::credential::{
     AwsCredential, CredentialProvider, InstanceCredentialProvider,
     StaticCredentialProvider, WebIdentityProvider,
 };
+use crate::client::retry::ClientConfig;
 use crate::multipart::{CloudMultiPartUpload, CloudMultiPartUploadImpl, UploadPart};
 use crate::util::str_is_truthy;
 use crate::{
@@ -407,7 +409,7 @@ pub struct AmazonS3Builder {
     endpoint: Option<String>,
     token: Option<String>,
     url: Option<String>,
-    retry_config: RetryConfig,
+    client_config: ClientConfig,
     imdsv1_fallback: bool,
     virtual_hosted_style_request: bool,
     unsigned_payload: bool,
@@ -850,7 +852,17 @@ impl AmazonS3Builder {
 
     /// Set the retry configuration
     pub fn with_retry(mut self, retry_config: RetryConfig) -> Self {
-        self.retry_config = retry_config;
+        self.client_config.retry = retry_config;
+        self
+    }
+
+    /// Set the tokio runtime to use to perform IO
+    ///
+    /// This allows isolating IO into a dedicated [`Runtime`](tokio::runtime::Runtime) either
+    /// to ensure acceptable scheduling jitter in the presence of CPU-bound tasks, or to allow
+    /// using `object_store` outside of a tokio context
+    pub fn with_tokio_runtime(mut self, runtime: Handle) -> Self {
+        self.client_config.runtime = Some(runtime);
         self
     }
 
@@ -978,7 +990,7 @@ impl AmazonS3Builder {
                         role_arn,
                         endpoint,
                         client,
-                        retry_config: self.retry_config.clone(),
+                        client_config: self.client_config.clone(),
                     }) as _
                 }
                 _ => match self.profile {
@@ -996,7 +1008,7 @@ impl AmazonS3Builder {
                         Box::new(InstanceCredentialProvider {
                             cache: Default::default(),
                             client: client_options.client()?,
-                            retry_config: self.retry_config.clone(),
+                            client_config: self.client_config.clone(),
                             imdsv1_fallback: self.imdsv1_fallback,
                             metadata_endpoint: self
                                 .metadata_endpoint
@@ -1031,7 +1043,7 @@ impl AmazonS3Builder {
             bucket,
             bucket_endpoint,
             credentials,
-            retry_config: self.retry_config,
+            client_config: self.client_config,
             client_options: self.client_options,
             sign_payload: !self.unsigned_payload,
             checksum: self.checksum_algorithm,
@@ -1063,8 +1075,8 @@ fn profile_credentials(
 mod tests {
     use super::*;
     use crate::tests::{
-        get_nonexistent_object, list_uses_directories_correctly, list_with_delimiter,
-        put_get_delete_list_opts, rename_and_copy, stream_get,
+        dedicated_tokio, get_nonexistent_object, list_uses_directories_correctly,
+        list_with_delimiter, put_get_delete_list_opts, rename_and_copy, stream_get,
     };
     use bytes::Bytes;
     use std::collections::HashMap;
@@ -1285,6 +1297,21 @@ mod tests {
         assert!(builder.is_err());
     }
 
+    #[test]
+    fn s3_test_non_tokio() {
+        let (handle, shutdown) = dedicated_tokio();
+        let config = maybe_skip_integration!();
+        let integration = config.with_tokio_runtime(handle).build().unwrap();
+        futures::executor::block_on(async move {
+            put_get_delete_list_opts(&integration, true).await;
+            list_uses_directories_correctly(&integration).await;
+            list_with_delimiter(&integration).await;
+            rename_and_copy(&integration).await;
+            stream_get(&integration).await;
+        });
+        shutdown();
+    }
+
     #[tokio::test]
     async fn s3_test() {
         let config = maybe_skip_integration!();