Worked on async uploading and switched to futures-util

Author: Kixiron

Cargo.lock

@@ -50,7 +50,7 @@ rusoto_core = "0.43"
 rusoto_credential = "0.43"
 
 # Async
-futures = "0.3"
+futures-util = "0.3"
 tokio = { version = "0.2", features = ["rt-threaded"] }
 
 [target.'cfg(not(windows))'.dependencies]

Cargo.toml

@@ -1,8 +1,7 @@
-use crate::storage::s3::{s3_client, S3_BUCKET_NAME};
+use crate::storage::s3::{s3_client, S3Backend, S3_BUCKET_NAME};
 use failure::{Error, Fail};
-use futures::executor;
 use postgres::Connection;
-use rusoto_s3::{DeleteObjectsRequest, ListObjectsV2Request, ObjectIdentifier, S3Client, S3};
+use rusoto_s3::{DeleteObjectsRequest, ListObjectsV2Request, ObjectIdentifier, S3};
 
 /// List of directories in docs.rs's underlying storage (either the database or S3) containing a
 /// subdirectory named after the crate. Those subdirectories will be deleted.
@@ -23,8 +22,9 @@ pub fn delete_crate(conn: &Connection, name: &str) -> Result<(), Error> {
     };
 
     delete_from_database(conn, name, crate_id)?;
-    if let Some(s3) = s3_client() {
-        delete_from_s3(&s3, name)?;
+    if let Some(client) = s3_client() {
+        let mut backend = S3Backend::new(client, S3_BUCKET_NAME);
+        delete_from_s3(&mut backend, name)?;
     }
 
     Ok(())
@@ -69,22 +69,26 @@ fn delete_from_database(conn: &Connection, name: &str, crate_id: i32) -> Result<
     Ok(())
 }
 
-fn delete_from_s3(s3: &S3Client, name: &str) -> Result<(), Error> {
+fn delete_from_s3(s3: &mut S3Backend<'_>, name: &str) -> Result<(), Error> {
     for prefix in STORAGE_PATHS_TO_DELETE {
         delete_prefix_from_s3(s3, &format!("{}/{}/", prefix, name))?;
     }
+
     Ok(())
 }
 
-fn delete_prefix_from_s3(s3: &S3Client, name: &str) -> Result<(), Error> {
+fn delete_prefix_from_s3(s3: &mut S3Backend<'_>, name: &str) -> Result<(), Error> {
     let mut continuation_token = None;
+
     loop {
-        let list = executor::block_on(s3.list_objects_v2(ListObjectsV2Request {
-            bucket: S3_BUCKET_NAME.into(),
-            prefix: Some(name.into()),
-            continuation_token,
-            ..ListObjectsV2Request::default()
-        }))?;
+        let list =
+            s3.runtime_handle()
+                .block_on(s3.client().list_objects_v2(ListObjectsV2Request {
+                    bucket: S3_BUCKET_NAME.into(),
+                    prefix: Some(name.into()),
+                    continuation_token,
+                    ..ListObjectsV2Request::default()
+                }))?;
 
         let to_delete = list
             .contents
@@ -97,19 +101,22 @@ fn delete_prefix_from_s3(s3: &S3Client, name: &str) -> Result<(), Error> {
             })
             .collect::<Vec<_>>();
 
-        let resp = executor::block_on(s3.delete_objects(DeleteObjectsRequest {
-            bucket: S3_BUCKET_NAME.into(),
-            delete: rusoto_s3::Delete {
-                objects: to_delete,
-                quiet: None,
-            },
-            ..DeleteObjectsRequest::default()
-        }))?;
+        let resp =
+            s3.runtime_handle()
+                .block_on(s3.client().delete_objects(DeleteObjectsRequest {
+                    bucket: S3_BUCKET_NAME.into(),
+                    delete: rusoto_s3::Delete {
+                        objects: to_delete,
+                        quiet: None,
+                    },
+                    ..DeleteObjectsRequest::default()
+                }))?;
 
         if let Some(errs) = resp.errors {
             for err in &errs {
                 log::error!("error deleting file from s3: {:?}", err);
             }
+
             failure::bail!("uploading to s3 failed");
         }
 

src/db/delete_crate.rs

@@ -61,7 +61,7 @@ pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>, Error> {
 
 pub(crate) enum Storage<'a> {
     Database(DatabaseBackend<'a>),
-    S3(S3Backend<'a>),
+    S3(Box<S3Backend<'a>>),
 }
 
 impl<'a> Storage<'a> {
@@ -79,9 +79,9 @@ impl<'a> Storage<'a> {
         }
     }
 
-    fn store_batch(&mut self, batch: &[Blob], trans: &Transaction) -> Result<(), Error> {
+    fn store_batch(&mut self, batch: Vec<Blob>, trans: &Transaction) -> Result<(), Error> {
         match self {
-            Self::Database(db) => db.store_batch(batch, trans),
+            Self::Database(db) => db.store_batch(&batch, trans),
             Self::S3(s3) => s3.store_batch(batch),
         }
     }
@@ -141,7 +141,7 @@ impl<'a> Storage<'a> {
             if batch.is_empty() {
                 break;
             }
-            self.store_batch(&batch, &trans)?;
+            self.store_batch(batch, &trans)?;
         }
 
         trans.commit()?;
@@ -180,7 +180,7 @@ impl<'a> From<DatabaseBackend<'a>> for Storage<'a> {
 
 impl<'a> From<S3Backend<'a>> for Storage<'a> {
     fn from(db: S3Backend<'a>) -> Self {
-        Self::S3(db)
+        Self::S3(Box::new(db))
     }
 }
 

src/storage/mod.rs

@@ -1,16 +1,17 @@
 use super::Blob;
 use failure::Error;
-use futures::future::{self, FutureExt};
+use futures_util::{
+    future::{FutureExt, TryFutureExt},
+    stream::{FuturesUnordered, StreamExt},
+};
 use log::{error, warn};
-use rusoto_core::{region::Region, RusotoError};
+use rusoto_core::region::Region;
 use rusoto_credential::DefaultCredentialsProvider;
-use rusoto_s3::{
-    GetObjectRequest, PutObjectError, PutObjectOutput, PutObjectRequest, S3Client, S3,
-};
+use rusoto_s3::{GetObjectRequest, PutObjectRequest, S3Client, S3};
 use std::convert::TryInto;
 use std::io::Read;
 use time::Timespec;
-use tokio::runtime::Runtime;
+use tokio::runtime::{Handle, Runtime};
 
 #[cfg(test)]
 mod test;
@@ -70,12 +71,16 @@ impl<'a> S3Backend<'a> {
         })
     }
 
-    pub(super) fn store_batch(&mut self, batch: &[Blob]) -> Result<(), Error> {
+    pub(super) fn store_batch(&mut self, mut uploads: Vec<Blob>) -> Result<(), Error> {
         let mut attempts = 0;
 
         loop {
-            let mut futures = Vec::with_capacity(batch.len());
-            for blob in batch {
+            // `FuturesUnordered` is used because the order of execution doesn't
+            // matter, we just want things to execute as fast as possible
+            let futures = FuturesUnordered::new();
+
+            // Drain uploads, filling `futures` with upload requests
+            for blob in uploads.drain(..) {
                 futures.push(
                     self.client
                         .put_object(PutObjectRequest {
@@ -85,33 +90,52 @@ impl<'a> S3Backend<'a> {
                             content_type: Some(blob.mime.clone()),
                             ..Default::default()
                         })
-                        .inspect(|_| {
+                        // Drop the value returned by `put_object` because we don't need it,
+                        // emit an error and replace the error values with the blob that failed
+                        // to upload so that we can retry failed uploads
+                        .map(|resp| match resp {
+                            Ok(..) => Ok(()),
+                            Err(err) => {
+                                error!("failed to upload file to s3: {:?}", err);
+                                Err(blob)
+                            }
+                        })
+                        .inspect_ok(|_| {
+                            // Increment the total uploaded files when a file is uploaded
                             crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                         }),
                 );
             }
             attempts += 1;
 
-            let result: Result<Vec<PutObjectOutput>, RusotoError<PutObjectError>> = self
-                .runtime
-                .block_on(future::join_all(futures))
-                .into_iter()
-                .collect();
-
-            match result {
-                // this batch was successful, start another batch if there are still more files
-                Ok(..) => break,
-                Err(err) => {
-                    error!("failed to upload to s3: {:?}", err);
-                    // if a futures error occurs, retry the batch
-                    if attempts > 2 {
-                        panic!("failed to upload 3 times, exiting");
-                    }
-                }
+            // Collect all the failed uploads so that we can retry them
+            uploads = self.runtime.block_on(
+                futures
+                    .filter_map(|resp| async move { resp.err() })
+                    .collect(),
+            );
+
+            // If there are no further uploads we were successful and can return
+            if uploads.is_empty() {
+                break;
+
+            // If more than three attempts to upload fail, return an error
+            } else if attempts >= 3 {
+                error!("failed to upload to s3, abandoning");
+                failure::bail!("Failed to upload to s3 three times, abandoning");
             }
         }
+
         Ok(())
     }
+
+    pub fn runtime_handle(&self) -> Handle {
+        self.runtime.handle().clone()
+    }
+
+    pub fn client(&self) -> &S3Client {
+        &self.client
+    }
 }
 
 fn parse_timespec(raw: &str) -> Result<Timespec, Error> {
@@ -125,13 +149,15 @@ pub(crate) fn s3_client() -> Option<S3Client> {
     if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
         return None;
     }
+
     let creds = match DefaultCredentialsProvider::new() {
         Ok(creds) => creds,
         Err(err) => {
             warn!("failed to retrieve AWS credentials: {}", err);
             return None;
         }
     };
+
     Some(S3Client::new_with(
         rusoto_core::request::HttpClient::new().unwrap(),
         creds,
@@ -149,7 +175,6 @@ pub(crate) fn s3_client() -> Option<S3Client> {
 pub(crate) mod tests {
     use super::*;
     use crate::test::*;
-    use std::slice;
 
     #[test]
     fn test_parse_timespec() {
@@ -179,7 +204,7 @@ pub(crate) mod tests {
 
             // Add a test file to the database
             let s3 = env.s3();
-            s3.upload(slice::from_ref(&blob)).unwrap();
+            s3.upload(vec![blob.clone()]).unwrap();
 
             // Test that the proper file was returned
             s3.assert_blob(&blob, "dir/foo.txt");
@@ -215,7 +240,7 @@ pub(crate) mod tests {
                 })
                 .collect();
 
-            s3.upload(&blobs).unwrap();
+            s3.upload(blobs.clone()).unwrap();
             for blob in &blobs {
                 s3.assert_blob(blob, &blob.path);
             }

src/storage/s3.rs

@@ -1,5 +1,6 @@
 use super::*;
 use crate::storage::test::assert_blob_eq;
+use rusoto_core::RusotoError;
 use rusoto_s3::{
     CreateBucketRequest, DeleteBucketRequest, DeleteObjectRequest, ListObjectsRequest, S3,
 };
@@ -28,7 +29,7 @@ impl TestS3 {
         )))
     }
 
-    pub(crate) fn upload(&self, blobs: &[Blob]) -> Result<(), Error> {
+    pub(crate) fn upload(&self, blobs: Vec<Blob>) -> Result<(), Error> {
         self.0.borrow_mut().store_batch(blobs)
     }