upload multipart in parallel

Author: parmesant

src/storage/object_storage.rs

@@ -844,7 +844,11 @@ pub trait ObjectStorage: Debug + Send + Sync + 'static {
                 let stream_relative_path = format!("{stream_name}/{file_suffix}");
 
                 // Try uploading the file, handle potential errors without breaking the loop
-                if let Err(e) = self.upload_file(&stream_relative_path, &path).await {
+                // if let Err(e) = self.upload_multipart(key, path)
+                if let Err(e) = self
+                    .upload_multipart(&RelativePathBuf::from(&stream_relative_path), &path)
+                    .await
+                {
                     error!("Failed to upload file {filename:?}: {e}");
                     continue; // Skip to the next file
                 }

src/storage/s3.rs

@@ -521,53 +521,50 @@ impl S3 {
 
         let mut async_writer = self.client.put_multipart(location).await?;
 
-        // /* `abort_multipart()` has been removed */
-        // let close_multipart = |err| async move {
-        //     error!("multipart upload failed. {:?}", err);
-        //     self.client
-        //         .abort_multipart(&key.into(), &multipart_id)
-        //         .await
-        // };
-
         let meta = file.metadata().await?;
         let total_size = meta.len() as usize;
         if total_size < MIN_MULTIPART_UPLOAD_SIZE {
             let mut data = Vec::new();
             file.read_to_end(&mut data).await?;
             self.client.put(location, data.into()).await?;
-            // async_writer.put_part(data.into()).await?;
-            // async_writer.complete().await?;
+
             return Ok(());
         } else {
-            let mut data = Vec::new();
-            file.read_to_end(&mut data).await?;
+            let mut buf = [0; MIN_MULTIPART_UPLOAD_SIZE * 2];
 
-            // let mut upload_parts = Vec::new();
+            let mut position = 0;
+            let mut tasks = FuturesUnordered::new();
 
-            let has_final_partial_part = total_size % MIN_MULTIPART_UPLOAD_SIZE > 0;
-            let num_full_parts = total_size / MIN_MULTIPART_UPLOAD_SIZE;
-            let total_parts = num_full_parts + if has_final_partial_part { 1 } else { 0 };
+            while position < total_size {
+                // Calculate how much to read in this iteration
+                let bytes_to_read = std::cmp::min(buf.len(), total_size - position);
+                let buffer_slice = &mut buf[..bytes_to_read];
 
-            // Upload each part
-            for part_number in 0..(total_parts) {
-                let start_pos = part_number * MIN_MULTIPART_UPLOAD_SIZE;
-                let end_pos = if part_number == num_full_parts && has_final_partial_part {
-                    // Last part might be smaller than 5MB (which is allowed)
-                    total_size
-                } else {
-                    // All other parts must be at least 5MB
-                    start_pos + MIN_MULTIPART_UPLOAD_SIZE
-                };
+                // Read exactly that many bytes
+                match file.read_exact(buffer_slice).await {
+                    Ok(_) => {}
+                    Err(e) => return Err(e.into()),
+                }
 
-                // Extract this part's data
-                let part_data = data[start_pos..end_pos].to_vec();
+                position += bytes_to_read;
 
                 // Upload the part
-                async_writer.put_part(part_data.into()).await?;
+                tasks.push(tokio::spawn(
+                    async_writer.put_part(buffer_slice.to_vec().into()),
+                ));
+            }
 
-                // upload_parts.push(part_number as u64 + 1);
+            while let Some(res) = tasks.next().await {
+                res??;
             }
-            async_writer.complete().await?;
+
+            match async_writer.complete().await {
+                Ok(_) => {}
+                Err(e) => {
+                    error!("Failed to complete multipart upload: {:?}", e);
+                    async_writer.abort().await?;
+                }
+            };
         }
         Ok(())
     }