Use S3 multipart uploads for large files (#85)

Author: saethlin

src/client.rs

@@ -1,6 +1,5 @@
 use crate::{Crate, Status, Tool, Version};
-use aws_sdk_s3::client::fluent_builders::PutObject;
-use aws_sdk_s3::model::Object;
+use aws_sdk_s3::model::{CompletedMultipartUpload, CompletedPart, Object};
 use backoff::Error;
 use backoff::ExponentialBackoff;
 use color_eyre::Result;
@@ -16,6 +15,8 @@ pub struct Client {
     tool: Tool,
 }
 
+const CHUNK_SIZE: usize = 5 * 1024 * 1024;
+
 impl Client {
     pub async fn new(tool: Tool, bucket: &str) -> Result<Self> {
         let config = aws_config::load_from_env().await;
@@ -31,22 +32,75 @@ impl Client {
         self.tool
     }
 
-    pub fn put_object(&self) -> PutObject {
-        self.inner.put_object().bucket(&self.bucket)
+    pub async fn upload(&self, key: &str, data: &[u8], content_type: &str) -> Result<()> {
+        retry(|| self._upload(key, data, content_type)).await
     }
 
-    pub async fn upload_raw(&self, krate: &Crate, data: Vec<u8>) -> Result<()> {
-        retry(|| {
-            self.put_object()
-                .key(self.tool.raw_crate_path(krate))
-                .body(data.clone().into())
-                .content_type("text/plain")
+    async fn _upload(&self, key: &str, data: &[u8], content_type: &str) -> Result<()> {
+        // S3 has a minimum multipart upload size of 5 MB. If we are below that, we need to use
+        // PutObject.
+        if data.len() < CHUNK_SIZE {
+            self.inner
+                .put_object()
+                .bucket(&self.bucket)
+                .key(key)
+                .body(data.to_vec().into())
+                .content_type(content_type)
                 .send()
-        })
-        .await?;
+                .await?;
+            return Ok(());
+        }
+
+        let res = self
+            .inner
+            .create_multipart_upload()
+            .bucket(&self.bucket)
+            .key(key)
+            .content_type(content_type)
+            .send()
+            .await?;
+        let upload_id = res.upload_id().unwrap();
+        let mut parts = Vec::new();
+        for (part_number, chunk) in data.chunks(CHUNK_SIZE).enumerate() {
+            // part numbers must start at 1
+            let part_number = part_number as i32 + 1;
+            let upload_part_res = self
+                .inner
+                .upload_part()
+                .key(key)
+                .bucket(&self.bucket)
+                .upload_id(upload_id)
+                .body(chunk.to_vec().into())
+                .part_number(part_number)
+                .send()
+                .await?;
+            parts.push(
+                CompletedPart::builder()
+                    .e_tag(upload_part_res.e_tag.unwrap_or_default())
+                    .part_number(part_number)
+                    .build(),
+            )
+        }
+        let completed_multipart_upload = CompletedMultipartUpload::builder()
+            .set_parts(Some(parts))
+            .build();
+        self.inner
+            .complete_multipart_upload()
+            .bucket(&self.bucket)
+            .key(key)
+            .multipart_upload(completed_multipart_upload)
+            .upload_id(upload_id)
+            .send()
+            .await?;
+
         Ok(())
     }
 
+    pub async fn upload_raw(&self, krate: &Crate, data: Vec<u8>) -> Result<()> {
+        self.upload(&self.tool.raw_crate_path(krate), &data, "text/plain")
+            .await
+    }
+
     pub async fn download_raw(&self, krate: &Crate) -> Result<Vec<u8>> {
         retry(|| self._download_raw(krate)).await
     }
@@ -69,17 +123,8 @@ impl Client {
     }
 
     pub async fn upload_html(&self, krate: &Crate, data: Vec<u8>) -> Result<()> {
-        retry(move || {
-            self.inner
-                .put_object()
-                .bucket(&self.bucket)
-                .key(self.tool.rendered_crate_path(krate))
-                .body(data.clone().into())
-                .content_type("text/html")
-                .send()
-        })
-        .await?;
-        Ok(())
+        let key = self.tool.rendered_crate_path(krate);
+        retry(|| self.upload(&key, &data, "text/html")).await
     }
 
     pub async fn get_crate_downloads(&self) -> Result<HashMap<String, Option<u64>>> {
@@ -174,17 +219,8 @@ impl Client {
     }
 
     pub async fn upload_landing_page(&self, data: Vec<u8>) -> Result<()> {
-        retry(move || {
-            self.inner
-                .put_object()
-                .bucket(&self.bucket)
-                .key(self.tool.landing_page_path())
-                .body(data.clone().into())
-                .content_type("text/html")
-                .send()
-        })
-        .await?;
-        Ok(())
+        self.upload(self.tool.landing_page_path(), &data, "text/html")
+            .await
     }
 
     pub async fn list_db(&self) -> Result<Option<Object>> {

src/sync.rs

@@ -23,11 +23,11 @@ pub async fn run(args: Args) -> Result<()> {
 
     log::info!("Uploading the error page");
     client
-        .put_object()
-        .key(format!("{}/403", args.tool))
-        .body(ERROR_PAGE.as_bytes().to_vec().into())
-        .content_type("text/html")
-        .send()
+        .upload(
+            &format!("{}/403", args.tool),
+            ERROR_PAGE.as_bytes(),
+            "text/html",
+        )
         .await?;
 
     let should_refresh_db = client
@@ -52,19 +52,11 @@ pub async fn run(args: Args) -> Result<()> {
 
         let serialized = serde_json::to_string(&versions).unwrap();
         client
-            .put_object()
-            .key("crates.json")
-            .body(serialized.into_bytes().into())
-            .content_type("application/json")
-            .send()
+            .upload("crates.json", serialized.as_bytes(), "application/json")
             .await?;
         let serialized = serde_json::to_string(&name_to_downloads).unwrap();
         client
-            .put_object()
-            .key("downloads.json")
-            .body(serialized.into_bytes().into())
-            .content_type("application/json")
-            .send()
+            .upload("downloads.json", serialized.as_bytes(), "application/json")
             .await?;
         name_to_downloads
     } else {
@@ -87,11 +79,11 @@ pub async fn run(args: Args) -> Result<()> {
 
     let ub_page = crate::render::render_ub(&crates)?;
     client
-        .put_object()
-        .key(format!("{}/ub", args.tool))
-        .body(ub_page.into_bytes().into())
-        .content_type("text/html")
-        .send()
+        .upload(
+            &format!("{}/ub", args.tool),
+            ub_page.as_bytes(),
+            "text/html",
+        )
         .await?;
 
     Ok(())
@@ -167,11 +159,11 @@ async fn sync_all_html(client: Arc<Client>) -> Result<Vec<Crate>> {
         .finish()
         .unwrap();
     client
-        .put_object()
-        .key(format!("{}/raw.tar.xz", client.tool()))
-        .body(raw.into())
-        .content_type("application/octet-stream")
-        .send()
+        .upload(
+            &format!("{}/raw.tar.xz", client.tool()),
+            &raw,
+            "application/octet-stream",
+        )
         .await?;
 
     let rendered: Vec<u8> = Arc::into_inner(all_rendered)
@@ -182,11 +174,11 @@ async fn sync_all_html(client: Arc<Client>) -> Result<Vec<Crate>> {
         .finish()
         .unwrap();
     client
-        .put_object()
-        .key(format!("{}/html.tar.xz", client.tool()))
-        .body(rendered.into())
-        .content_type("application/octet-stream")
-        .send()
+        .upload(
+            &format!("{}/html.tar.xz", client.tool()),
+            &rendered,
+            "application/octet-stream",
+        )
         .await?;
 
     Ok(crates)