Adaptive Concurrancy Controller.

Author: hoytak

Cargo.lock

@@ -11,4 +11,39 @@ utils::configurable_constants! {
     /// no more retries are attempted.
     ref CLIENT_RETRY_MAX_DURATION_MS: u64 = 6 * 60 * 1000; // 6m
 
+    /// The target time for a small transfer to complete.
+    ref CONCURRENCY_CONTROL_TARGET_TIME_SMALL_TRANSFER_MS : u64 = 5 * 1000;
+
+    /// The target time for a large transfer to complete.  Default is 20 seconds.
+    ref CONCURRENCY_CONTROL_TARGET_TIME_LARGE_TRANSFER_MS : u64 = 20 * 1000;
+
+    /// The size of a large transfer.
+    ref CONCURRENCY_CONTROL_LARGE_TRANSFER_NUM_BYTES : u64 = 64_000_000;
+
+    /// The minimum time in milliseconds between adjustments when increasing the concurrency.
+    ref CONCURRENCY_CONTROL_MIN_INCREASE_WINDOW_MS : u64 = 500;
+
+    /// The minimum time in milliseconds between adjustments when decreasing the concurrency.
+    ref CONCURRENCY_CONTROL_MIN_DECREASE_WINDOW_MS : u64 = 250;
+
+    /// The maximum number of connection successes and failures to examine when adjusting the concurrancy.
+    ref CONCURRENCY_CONTROL_TRACKING_SIZE : usize = 32;
+
+    /// The maximum time window within which to examine successes and failures when adjusting the concurrancy.
+    ref CONCURRENCY_CONTROL_TRACKING_WINDOW_MS : u64 = 30 * 1000;
+
+    /// Log the concurrency on this interval.
+    ref CONCURRENCY_CONTROL_LOGGING_INTERVAL_MS: u64 = 10 * 1000;
+
+    /// The maximum number of simultaneous xorb and/or shard upload streams permitted by
+    /// the adaptive concurrency control. Can be overwritten by environment variable "HF_XET_MAX_CONCURRENT_UPLOADS".
+    ref MAX_CONCURRENT_UPLOADS: usize = 100;
+
+    /// The minimum number of simultaneous xorb and/or shard upload streams that the
+    /// the adaptive concurrency control may reduce the concurrancy down to on slower connections.
+    ref MIN_CONCURRENT_UPLOADS: usize = 2;
+
+    /// The starting number of concurrent upload streams, which will increase up to MAX_CONCURRENT_UPLOADS
+    /// on successful completions.
+    ref NUM_INITIAL_CONCURRENT_UPLOADS: usize = 16;
 }

cas_client/src/adaptive_concurrency_control.rs

@@ -9,10 +9,18 @@ use merklehash::MerkleHash;
 use progress_tracking::item_tracking::SingleItemProgressUpdater;
 use progress_tracking::upload_tracking::CompletionTracker;
 
+use crate::adaptive_concurrency_control::{AdaptiveConcurrencyController, ConnectionPermit};
+use crate::constants::{MAX_CONCURRENT_UPLOADS, MIN_CONCURRENT_UPLOADS, NUM_INITIAL_CONCURRENT_UPLOADS};
 use crate::error::Result;
 #[cfg(not(target_family = "wasm"))]
 use crate::OutputProvider;
 
+// The upload concurrency controller
+lazy_static::lazy_static! {
+    static ref UPLOAD_CONCURRENCY_CONTROLLER : Arc<AdaptiveConcurrencyController>
+        = AdaptiveConcurrencyController::new("uploads", *NUM_INITIAL_CONCURRENT_UPLOADS, (*MIN_CONCURRENT_UPLOADS, *MAX_CONCURRENT_UPLOADS));
+}
+
 /// A Client to the Shard service. The shard service
 /// provides for
 /// 1. upload shard to the shard service
@@ -57,23 +65,56 @@ pub trait Client {
     ) -> Result<Option<Bytes>>;
 
     /// Upload a new shard.
-    async fn upload_shard(
+    async fn upload_shard_with_permit(
         &self,
         prefix: &str,
         hash: &MerkleHash,
         force_sync: bool,
         shard_data: bytes::Bytes,
         salt: &[u8; 32],
+        upload_permit: ConnectionPermit,
     ) -> Result<bool>;
 
     /// Upload a new xorb.
-    async fn upload_xorb(
+    async fn upload_xorb_with_permit(
         &self,
         prefix: &str,
         serialized_cas_object: SerializedCasObject,
         upload_tracker: Option<Arc<CompletionTracker>>,
+        upload_permit: ConnectionPermit,
     ) -> Result<u64>;
 
+    /// Acquire an upload permit.
+    async fn acquire_upload_permit(&self) -> Result<ConnectionPermit> {
+        UPLOAD_CONCURRENCY_CONTROLLER.acquire_connection_permit().await
+    }
+
+    /// Upload a new shard, acquiring the permit.
+    async fn upload_shard(
+        &self,
+        prefix: &str,
+        hash: &MerkleHash,
+        force_sync: bool,
+        shard_data: bytes::Bytes,
+        salt: &[u8; 32],
+    ) -> Result<bool> {
+        let permit = self.acquire_upload_permit().await?;
+        self.upload_shard_with_permit(prefix, hash, force_sync, shard_data, salt, permit)
+            .await
+    }
+
+    /// Upload a new xorb, acquiring the permit.
+    async fn upload_xorb(
+        &self,
+        prefix: &str,
+        serialized_cas_object: SerializedCasObject,
+        upload_tracker: Option<Arc<CompletionTracker>>,
+    ) -> Result<u64> {
+        let permit = self.acquire_upload_permit().await?;
+        self.upload_xorb_with_permit(prefix, serialized_cas_object, upload_tracker, permit)
+            .await
+    }
+
     /// Check if a XORB already exists.
     async fn exists(&self, prefix: &str, hash: &MerkleHash) -> Result<bool>;
 

cas_client/src/constants.rs

@@ -25,3 +25,5 @@ mod output_provider;
 pub mod remote_client;
 mod retry_wrapper;
 mod upload_progress_stream;
+
+mod adaptive_concurrency_control;

cas_client/src/interface.rs

@@ -22,6 +22,7 @@ use tempfile::TempDir;
 use tokio::runtime::Handle;
 use tracing::{debug, error, info, warn};
 
+use crate::adaptive_concurrency_control::ConnectionPermit;
 use crate::error::{CasClientError, Result};
 use crate::output_provider::OutputProvider;
 use crate::Client;
@@ -214,11 +215,12 @@ impl LocalClient {
 /// LocalClient is responsible for writing/reading Xorbs on local disk.
 #[async_trait]
 impl Client for LocalClient {
-    async fn upload_xorb(
+    async fn upload_xorb_with_permit(
         &self,
         _prefix: &str,
         serialized_cas_object: SerializedCasObject,
         upload_tracker: Option<Arc<CompletionTracker>>,
+        _upload_permit: ConnectionPermit,
     ) -> Result<u64> {
         // moved hash validation into [CasObject::serialize], so removed from here.
         let hash = &serialized_cas_object.hash;
@@ -302,13 +304,14 @@ impl Client for LocalClient {
         true
     }
 
-    async fn upload_shard(
+    async fn upload_shard_with_permit(
         &self,
         _prefix: &str, // Prefix not used in current implementation
         shard_hash: &MerkleHash,
         _force_sync: bool,
         shard_data: bytes::Bytes,
         salt: &[u8; 32],
+        _upload_permit: ConnectionPermit,
     ) -> Result<bool> {
         // Write out the shard to the shard directory.
         let shard = MDBShardFile::write_out_from_reader(&self.shard_dir, &mut Cursor::new(&shard_data))?;

cas_client/src/lib.rs

@@ -28,6 +28,7 @@ use utils::auth::AuthConfig;
 #[cfg(not(target_family = "wasm"))]
 use utils::singleflight::Group;
 
+use crate::adaptive_concurrency_control::ConnectionPermit;
 #[cfg(not(target_family = "wasm"))]
 use crate::download_utils::*;
 use crate::error::{CasClientError, Result};
@@ -596,11 +597,12 @@ impl Client for RemoteClient {
     #[instrument(skip_all, name = "RemoteClient::upload_xorb", fields(key = Key{prefix : prefix.to_string(), hash : serialized_cas_object.hash}.to_string(),
                  xorb.len = serialized_cas_object.serialized_data.len(), xorb.num_chunks = serialized_cas_object.num_chunks
     ))]
-    async fn upload_xorb(
+    async fn upload_xorb_with_permit(
         &self,
         prefix: &str,
         serialized_cas_object: SerializedCasObject,
         upload_tracker: Option<Arc<CompletionTracker>>,
+        upload_permit: ConnectionPermit,
     ) -> Result<u64> {
         let key = Key {
             prefix: prefix.to_string(),
@@ -639,6 +641,7 @@ impl Client for RemoteClient {
                 let api_tag = "cas::upload_xorb";
 
                 let response: UploadXorbResponse = RetryWrapper::new(api_tag)
+                    .with_connection_permit(upload_permit)
                     .run_and_extract_json(move || {
                         let upload_stream = upload_stream.clone_with_reset();
                         let url = url.clone();
@@ -773,13 +776,14 @@ impl Client for RemoteClient {
 
     #[instrument(skip_all, name = "RemoteClient::upload_shard", fields(shard.hash = hash.hex(), shard.len = shard_data.len()
     ))]
-    async fn upload_shard(
+    async fn upload_shard_with_permit(
         &self,
         prefix: &str,
         hash: &MerkleHash,
         force_sync: bool,
         shard_data: bytes::Bytes,
         _salt: &[u8; 32],
+        upload_permit: ConnectionPermit,
     ) -> Result<bool> {
         if self.dry_run {
             return Ok(true);
@@ -801,6 +805,7 @@ impl Client for RemoteClient {
         let url = Url::parse(&format!("{}/shard/{key}", self.endpoint))?;
 
         let response: UploadShardResponse = RetryWrapper::new(api_tag)
+            .with_connection_permit(upload_permit)
             .run_and_extract_json(move || {
                 client
                     .request(method.clone(), url.clone())

cas_client/src/local_client.rs

@@ -1,12 +1,15 @@
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 
+use bytes::Bytes;
 use reqwest::{Response, StatusCode};
 use reqwest_retry::{default_on_request_failure, default_on_request_success, Retryable};
+use tokio::sync::Mutex;
 use tokio_retry::strategy::{jitter, ExponentialBackoff};
 use tokio_retry::RetryIf;
 use tracing::{error, info};
 
+use crate::adaptive_concurrency_control::ConnectionPermit;
 use crate::constants::{CLIENT_RETRY_BASE_DELAY_MS, CLIENT_RETRY_MAX_ATTEMPTS};
 use crate::error::CasClientError;
 use crate::http_client::request_id_from_response;
@@ -23,6 +26,7 @@ pub struct RetryWrapper {
     no_retry_on_429: bool,
     log_errors_as_info: bool,
     api_tag: &'static str,
+    connection_permit: Option<Mutex<Option<ConnectionPermit>>>,
 }
 
 impl RetryWrapper {
@@ -33,6 +37,7 @@ impl RetryWrapper {
             no_retry_on_429: false,
             log_errors_as_info: false,
             api_tag,
+            connection_permit: None,
         }
     }
 
@@ -56,6 +61,11 @@ impl RetryWrapper {
         self
     }
 
+    pub fn with_connection_permit(mut self, permit: ConnectionPermit) -> Self {
+        self.connection_permit = Some(Mutex::new(Some(permit)));
+        self
+    }
+
     fn process_error_response(&self, try_idx: usize, err: reqwest_middleware::Error) -> RetryableReqwestError {
         let api = &self.api_tag;
 
@@ -194,6 +204,12 @@ impl RetryWrapper {
                 async move {
                     let (make_request, process_fn, try_count, self_) = retry_info.as_ref();
 
+                    if let Some(p) = &self_.connection_permit {
+                        if let Some(p) = p.lock().await.as_mut() {
+                            p.transfer_starting()
+                        }
+                    }
+
                     let resp_result = make_request().await;
                     let try_idx = try_count.fetch_add(1, Ordering::Relaxed);
 
@@ -204,10 +220,35 @@ impl RetryWrapper {
                         Ok(resp) => self_.process_ok_response(try_idx, resp),
                     };
 
-                    match checked_result {
-                        Ok(ok_response) => process_fn(ok_response).await,
-                        Err(e) => Err(e),
+                    let (n_bytes, processing_result) = match checked_result {
+                        Ok(ok_response) => (ok_response.content_length().unwrap_or(0), process_fn(ok_response).await),
+                        Err(e) => (0, Err(e)),
+                    };
+
+                    // Now, possibly adjust the connection permit.
+                    if let Some(permit_holder) = &self_.connection_permit {
+                        let mut maybe_permit = permit_holder.lock().await;
+
+                        match &processing_result {
+                            Ok(_) => {
+                                if let Some(permit) = maybe_permit.take() {
+                                    permit.report_completion(n_bytes, true).await;
+                                }
+                            },
+                            Err(RetryableReqwestError::FatalError(_)) => {
+                                if let Some(permit) = maybe_permit.take() {
+                                    permit.report_completion(0, false).await;
+                                }
+                            },
+                            Err(RetryableReqwestError::RetryableError(_)) => {
+                                if let Some(permit) = maybe_permit.as_ref() {
+                                    permit.report_retryable_failure().await;
+                                }
+                            },
+                        }
                     }
+
+                    processing_result
                 }
             },
             |err: &RetryableReqwestError| matches!(err, RetryableReqwestError::RetryableError(_)),
@@ -278,6 +319,47 @@ impl RetryWrapper {
         .await
     }
 
+    /// Run a connection and process the result as bytes, retrying on transient errors or on issues not getting the
+    /// full object.
+    ///
+    /// The `make_request` function returns a future that resolves to a Result<Response> object as is returned by the
+    /// client middleware.  For example, `|| client.clone().get(url).send()` returns a future (as `send()` is async)
+    /// that will then be evaluatated to get the response.
+    ///
+    /// This functions acts just like the json() function on a client response, but retries the entire connection on
+    /// transient errors.  
+    pub async fn run_and_extract_bytes<ReqFut, ReqFn>(self, make_request: ReqFn) -> Result<Bytes, CasClientError>
+    where
+        ReqFn: Fn() -> ReqFut + Send + 'static,
+        ReqFut: std::future::Future<Output = Result<Response, reqwest_middleware::Error>> + 'static,
+    {
+        self.run_and_process(make_request, |resp: Response| {
+            async move {
+                // Extract the bytes from the final result.
+                let r: Result<Bytes, reqwest::Error> = resp.bytes().await;
+
+                match r {
+                    Ok(v) => Ok(v),
+                    Err(e) => {
+                        #[cfg(not(target_arch = "wasm32"))]
+                        let is_connect = e.is_connect();
+                        #[cfg(target_arch = "wasm32")]
+                        let is_connect = false;
+
+                        if is_connect || e.is_decode() || e.is_body() || e.is_timeout() {
+                            // We got an incomplete or corrupted response from the server, possibly due to a dropped
+                            // connection.  Presumably this error is transient.
+                            Err(RetryableReqwestError::RetryableError(e.into()))
+                        } else {
+                            Err(RetryableReqwestError::FatalError(e.into()))
+                        }
+                    },
+                }
+            }
+        })
+        .await
+    }
+
     /// Run a connection and process the result object, retrying on transient errors.
     ///
     /// The `make_request` function returns a future that resolves to a Result<Response> object as is returned by the

cas_client/src/remote_client.rs

@@ -13,14 +13,6 @@ utils::configurable_constants! {
     /// set to 3 weeks.
     ref MDB_SHARD_LOCAL_CACHE_EXPIRATION_SECS: u64 = 3 * 7 * 24 * 3600;
 
-    /// The maximum number of simultaneous xorb upload streams.
-    /// can be overwritten by environment variable "HF_XET_MAX_CONCURRENT_UPLOADS".
-    /// The default value changes from 8 to 100 when "High Performance Mode" is enabled
-    ref MAX_CONCURRENT_UPLOADS: usize = GlobalConfigMode::HighPerformanceOption {
-        standard: 8,
-        high_performance: 100,
-    };
-
     /// The maximum number of files to ingest at once on the upload path
     ref MAX_CONCURRENT_FILE_INGESTION: usize =  GlobalConfigMode::HighPerformanceOption {
         standard: 8,