update chunking size to 10mb and re-use client on upload

Author: gschoeni

src/lib/src/api/client.rs

@@ -3,6 +3,7 @@
 
 use crate::config::AuthConfig;
 use crate::error::OxenError;
+use crate::model::RemoteRepository;
 use crate::view::http;
 use crate::view::OxenResponse;
 
@@ -57,6 +58,11 @@ fn new_for_host<S: AsRef<str>>(host: S, should_add_user_agent: bool) -> Result<C
     }
 }
 
+pub fn builder_for_remote_repo(remote_repo: &RemoteRepository) -> Result<ClientBuilder, OxenError> {
+    let host = get_host_from_url(remote_repo.url())?;
+    builder_for_host(host, true)
+}
+
 pub fn builder_for_url<U: IntoUrl>(url: U) -> Result<ClientBuilder, OxenError> {
     let host = get_host_from_url(url)?;
     builder_for_host(host, true)

src/lib/src/api/client/commits.rs

@@ -29,6 +29,7 @@ use bytesize::ByteSize;
 use flate2::write::GzEncoder;
 use flate2::Compression;
 use futures_util::TryStreamExt;
+use http::header::CONTENT_LENGTH;
 use indicatif::{ProgressBar, ProgressStyle};
 
 pub struct ChunkParams {
@@ -954,6 +955,104 @@ pub async fn upload_data_chunk_to_server_with_retry(
     )))
 }
 
+pub async fn upload_data_chunk_to_server_with_client_with_retry(
+    client: &reqwest::Client,
+    remote_repo: &RemoteRepository,
+    chunk: &[u8],
+    hash: &str,
+    params: &ChunkParams,
+    is_compressed: bool,
+    filename: &Option<String>,
+) -> Result<(), OxenError> {
+    let mut total_tries = 0;
+    let mut last_error = String::from("");
+    while total_tries < constants::NUM_HTTP_RETRIES {
+        match upload_data_chunk_to_server_with_client(
+            client,
+            remote_repo,
+            chunk,
+            hash,
+            params,
+            is_compressed,
+            filename,
+        )
+        .await
+        {
+            Ok(_) => {
+                return Ok(());
+            }
+            Err(err) => {
+                total_tries += 1;
+                // Exponentially back off
+                let sleep_time = total_tries * total_tries;
+                log::debug!(
+                    "upload_data_chunk_to_server_with_client_with_retry upload failed sleeping {}: {}",
+                    sleep_time,
+                    err
+                );
+                last_error = format!("{}", err);
+                std::thread::sleep(std::time::Duration::from_secs(sleep_time));
+            }
+        }
+    }
+
+    Err(OxenError::basic_str(format!(
+        "Upload chunk retry failed. {}",
+        last_error
+    )))
+}
+
+async fn upload_data_chunk_to_server_with_client(
+    client: &reqwest::Client,
+    remote_repo: &RemoteRepository,
+    chunk: &[u8],
+    hash: &str,
+    params: &ChunkParams,
+    is_compressed: bool,
+    filename: &Option<String>,
+) -> Result<StatusMessage, OxenError> {
+    let maybe_filename = if !is_compressed {
+        format!(
+            "&filename={}",
+            urlencoding::encode(
+                filename
+                    .as_ref()
+                    .expect("Must provide filename if !compressed")
+            )
+        )
+    } else {
+        String::from("")
+    };
+
+    let uri = format!(
+        "/commits/upload_chunk?chunk_num={}&total_size={}&hash={}&total_chunks={}&is_compressed={}{}",
+        params.chunk_num, params.total_size, hash, params.total_chunks, is_compressed, maybe_filename);
+    let url = api::endpoint::url_from_repo(remote_repo, &uri)?;
+    let total_size = chunk.len() as u64;
+    log::debug!(
+        "upload_data_chunk_to_server posting {} to url {}",
+        ByteSize::b(total_size),
+        url
+    );
+
+    let res = client
+        .post(&url)
+        .header(CONTENT_LENGTH, total_size.to_string())
+        .body(chunk.to_owned())
+        .send()
+        .await?;
+    let body = client::parse_json_body(&url, res).await?;
+
+    log::debug!("upload_data_chunk_to_server got response {}", body);
+    let response: Result<StatusMessage, serde_json::Error> = serde_json::from_str(&body);
+    match response {
+        Ok(response) => Ok(response),
+        Err(err) => Err(OxenError::basic_str(format!(
+            "upload_data_chunk_to_server Err deserializing: {err}"
+        ))),
+    }
+}
+
 async fn upload_data_chunk_to_server(
     remote_repo: &RemoteRepository,
     chunk: &[u8],
@@ -990,23 +1089,16 @@ async fn upload_data_chunk_to_server(
         .timeout(time::Duration::from_secs(120))
         .build()?;
 
-    match client.post(&url).body(chunk.to_owned()).send().await {
-        Ok(res) => {
-            let body = client::parse_json_body(&url, res).await?;
+    let res = client.post(&url).body(chunk.to_owned()).send().await?;
+    let body = client::parse_json_body(&url, res).await?;
 
-            log::debug!("upload_data_chunk_to_server got response {}", body);
-            let response: Result<StatusMessage, serde_json::Error> = serde_json::from_str(&body);
-            match response {
-                Ok(response) => Ok(response),
-                Err(err) => Err(OxenError::basic_str(format!(
-                    "upload_data_chunk_to_server Err deserializing: {err}"
-                ))),
-            }
-        }
-        Err(e) => {
-            let err_str = format!("Err upload_data_chunk_to_server: {e:?}");
-            Err(OxenError::basic_str(err_str))
-        }
+    log::debug!("upload_data_chunk_to_server got response {}", body);
+    let response: Result<StatusMessage, serde_json::Error> = serde_json::from_str(&body);
+    match response {
+        Ok(response) => Ok(response),
+        Err(err) => Err(OxenError::basic_str(format!(
+            "upload_data_chunk_to_server Err deserializing: {err}"
+        ))),
     }
 }
 

src/lib/src/constants.rs

@@ -173,10 +173,10 @@ pub const EVAL_ERROR_COL: &str = "_oxen_eval_error";
 pub const EVAL_DURATION_COL: &str = "_oxen_eval_duration";
 
 // Data transfer
-// Average chunk size of ~4mb
-/// Average chunk size of ~4mb when chunking and sending data
+// Average chunk size of ~10mb
+/// Average chunk size of ~10mb when chunking and sending data
 // pub const AVG_CHUNK_SIZE: u64 = 1024 * 1024 * 4;
-pub const AVG_CHUNK_SIZE: u64 = 1024 * 1024 * 4;
+pub const AVG_CHUNK_SIZE: u64 = 1024 * 1024 * 10;
 // Retry and back off of requests N times
 /// Retry and back off of requests N times
 #[cfg(test)]

src/lib/src/core/v_latest/push.rs

@@ -9,8 +9,8 @@ use std::sync::Arc;
 use tokio::time::Duration;
 
 use crate::api::client::commits::ChunkParams;
+use crate::constants::AVG_CHUNK_SIZE;
 use crate::constants::DEFAULT_REMOTE_NAME;
-use crate::constants::{self, AVG_CHUNK_SIZE};
 use crate::core::progress::push_progress::PushProgress;
 use crate::error::OxenError;
 use crate::model::entry::commit_entry::Entry;
@@ -446,13 +446,20 @@ async fn upload_large_file_chunks(
     let mut total_bytes_read = 0;
     let mut chunk_size = chunk_size;
 
+    // Create a client for uploading chunks
+    let client = api::client::builder_for_remote_repo(&remote_repo)
+        .unwrap()
+        .build()
+        .unwrap();
+
     // Create queues for sending data to workers
     type PieceOfWork = (
         Vec<u8>,
         u64,   // chunk size
         usize, // chunk num
         usize, // total chunks
         u64,   // total size
+        reqwest::Client,
         RemoteRepository,
         String, // entry hash
         Commit,
@@ -462,13 +469,13 @@ async fn upload_large_file_chunks(
     // In order to upload chunks in parallel
     // We should only read N chunks at a time so that
     // the whole file does not get read into memory
-    let sub_chunk_size = constants::DEFAULT_NUM_WORKERS;
+    let sub_chunk_size = concurrency::num_threads_for_items(total_chunks);
 
     let mut total_chunk_idx = 0;
     let mut processed_chunk_idx = 0;
     let num_sub_chunks = (total_chunks / sub_chunk_size) + 1;
     log::debug!(
-        "upload_large_file_chunks {:?} proccessing file in {} subchunks of size {} from total {} chunk size {} file size {}",
+        "upload_large_file_chunks {:?} processing file in {} subchunks of size {} from total {} chunk size {} file size {}",
         entry.path(),
         num_sub_chunks,
         sub_chunk_size,
@@ -530,6 +537,7 @@ async fn upload_large_file_chunks(
                 processed_chunk_idx, // Needs to be the overall chunk num
                 total_chunks,
                 total_bytes,
+                client.clone(),
                 remote_repo.to_owned(),
                 entry.hash().to_owned(),
                 commit.to_owned(),
@@ -548,6 +556,7 @@ async fn upload_large_file_chunks(
                     chunk_num,
                     total_chunks,
                     total_size,
+                    client,
                     remote_repo,
                     entry_hash,
                     _commit,
@@ -568,7 +577,8 @@ async fn upload_large_file_chunks(
                 };
 
                 let is_compressed = false;
-                match api::client::commits::upload_data_chunk_to_server_with_retry(
+                match api::client::commits::upload_data_chunk_to_server_with_client_with_retry(
+                    &client,
                     &remote_repo,
                     &buffer,
                     &entry_hash,