Update.

Author: hoytak

Cargo.lock

@@ -14,7 +14,6 @@ use crate::constants::{MAX_CONCURRENT_UPLOADS, MIN_CONCURRENT_UPLOADS, NUM_INITI
 use crate::error::Result;
 #[cfg(not(target_family = "wasm"))]
 use crate::OutputProvider;
-use crate::OutputProvider;
 
 // The upload concurrency controller
 lazy_static::lazy_static! {
@@ -66,28 +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/adaptive_concurrency_control.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/interface.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

@@ -205,7 +205,9 @@ impl RetryWrapper {
                     let (make_request, process_fn, try_count, self_) = retry_info.as_ref();
 
                     if let Some(p) = &self_.connection_permit {
-                        p.lock().await.as_mut().map(|p| p.transfer_starting());
+                        if let Some(p) = p.lock().await.as_mut() {
+                            p.transfer_starting()
+                        }
                     }
 
                     let resp_result = make_request().await;

cas_client/src/remote_client.rs

@@ -349,11 +349,9 @@ impl FileUploadSession {
             async move {
                 let n_bytes_transmitted = session
                     .client
-                    .upload_xorb(&cas_prefix, cas_object, Some(completion_tracker))
+                    .upload_xorb_with_permit(&cas_prefix, cas_object, Some(completion_tracker), upload_permit)
                     .await?;
 
-                drop(upload_permit);
-
                 // Register that the xorb has been uploaded.
                 session.completion_tracker.register_xorb_upload_completion(xorb_hash).await;
 

cas_client/src/retry_wrapper.rs

@@ -281,10 +281,9 @@ impl SessionShardInterface {
                     }
 
                     // Upload the shard.
-                    client.upload_shard(&shard_prefix, &si.shard_hash, false, data, &salt).await?;
-
-                    // Done with the upload, drop the permit.
-                    drop(upload_permit);
+                    client
+                        .upload_shard_with_permit(&shard_prefix, &si.shard_hash, false, data, &salt, upload_permit)
+                        .await?;
 
                     info!("Shard {shard_prefix}/{:?} upload + sync completed successfully.", &si.shard_hash);
 

data/src/file_upload_session.rs

@@ -28,5 +28,8 @@ xet_threadpool = { path = "../xet_threadpool" }
 [target.'cfg(target_family = "wasm")'.dependencies]
 web-time = { workspace = true }
 
+[dev-dependencies]
+rand = { workspace = true }
+
 [features]
 strict = []