Merge pull request #45 from oltrep/retry-backend

Closes #43

This PR adds a new backend that wraps another one and performs the retry logic when getting rate limited errors.

I've opened this as a draft PR because I have not tested it yet with real image uploads.

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased Changes
 
+* Added support for automatically retrying uploads when being rate limited. ([#42](https://github.com/Roblox/tarmac/issues/43))
+
 ## 0.6.0 (2020-07-09)
 * Added support for automatically retrying image uploads when image names are moderated. ([#26](https://github.com/Roblox/tarmac/issues/26))
 * Added `create-cache-map` subcommand to aid in prepopulating asset caches. ([#38](https://github.com/Roblox/tarmac/pull/38))

README.md

@@ -113,6 +113,8 @@ Usage:
 ```bash
 tarmac sync [<config-path>] \
 	--target <roblox|debug|none>
+	--retry <number>
+	--retry-delay <60>
 ```
 
 To sync the project in your current working directory with the Roblox cloud, use:
@@ -125,6 +127,11 @@ To validate that all inputs are already synced, use the `none` target:
 tarmac sync --target none
 ```
 
+When tarmac gets rate limited while syncing to Roblox, use the `--retry` argument to automatically attempt to re-upload. This will tell tarmac how many times it can attempt to re-upload each asset. The `--retry-delay` sets the number of seconds to wait between each attempt.
+```bash
+tarmac sync --target roblox --retry 3
+```
+
 ### `tarmac upload-image`
 Uploads a single image as a decal and prints the ID of the resulting image asset to stdout.
 

src/commands/sync.rs

@@ -3,6 +3,7 @@ use std::{
     env,
     io::{self, BufWriter, Write},
     path::{Path, PathBuf},
+    time::Duration,
 };
 
 use fs_err as fs;
@@ -21,14 +22,24 @@ use crate::{
     options::{GlobalOptions, SyncOptions, SyncTarget},
     roblox_web_api::{RobloxApiClient, RobloxApiError},
     sync_backend::{
-        DebugSyncBackend, Error as SyncBackendError, NoneSyncBackend, RobloxSyncBackend,
-        SyncBackend, UploadInfo,
+        DebugSyncBackend, Error as SyncBackendError, NoneSyncBackend, RetryBackend,
+        RobloxSyncBackend, SyncBackend, UploadInfo,
     },
 };
 
+fn sync_session<B: SyncBackend>(session: &mut SyncSession, options: &SyncOptions, mut backend: B) {
+    if let Some(retry) = options.retry {
+        let mut retry_backend =
+            RetryBackend::new(backend, retry, Duration::from_secs(options.retry_delay));
+        session.sync_with_backend(&mut retry_backend);
+    } else {
+        session.sync_with_backend(&mut backend);
+    }
+}
+
 pub fn sync(global: GlobalOptions, options: SyncOptions) -> Result<(), SyncError> {
-    let fuzzy_config_path = match options.config_path {
-        Some(v) => v,
+    let fuzzy_config_path = match &options.config_path {
+        Some(v) => v.to_owned(),
         None => env::current_dir()?,
     };
 
@@ -39,21 +50,20 @@ pub fn sync(global: GlobalOptions, options: SyncOptions) -> Result<(), SyncError
     session.discover_configs()?;
     session.discover_inputs()?;
 
-    match options.target {
+    match &options.target {
         SyncTarget::Roblox => {
-            let mut backend =
-                RobloxSyncBackend::new(&mut api_client, session.root_config().upload_to_group_id);
-
-            session.sync_with_backend(&mut backend);
+            let group_id = session.root_config().upload_to_group_id;
+            sync_session(
+                &mut session,
+                &options,
+                RobloxSyncBackend::new(&mut api_client, group_id),
+            );
         }
         SyncTarget::None => {
-            let mut backend = NoneSyncBackend;
-
-            session.sync_with_backend(&mut backend);
+            sync_session(&mut session, &options, NoneSyncBackend);
         }
         SyncTarget::Debug => {
-            let mut backend = DebugSyncBackend::new();
-            session.sync_with_backend(&mut backend);
+            sync_session(&mut session, &options, DebugSyncBackend::new());
         }
     }
 

src/options.rs

@@ -71,6 +71,15 @@ pub struct SyncOptions {
     #[structopt(long)]
     pub target: SyncTarget,
 
+    /// When provided, Tarmac will upload again at most the given number of times
+    /// when it encounters rate limitation errors.
+    #[structopt(long)]
+    pub retry: Option<usize>,
+
+    /// The number of seconds to wait between each re-upload attempts.
+    #[structopt(long, default_value = "60")]
+    pub retry_delay: u64,
+
     /// The path to a Tarmac config, or a folder containing a Tarmac project.
     pub config_path: Option<PathBuf>,
 }

src/sync_backend.rs

@@ -1,4 +1,4 @@
-use std::{borrow::Cow, io, path::Path};
+use std::{borrow::Cow, io, path::Path, thread, time::Duration};
 
 use fs_err as fs;
 use reqwest::StatusCode;
@@ -10,10 +10,12 @@ pub trait SyncBackend {
     fn upload(&mut self, data: UploadInfo) -> Result<UploadResponse, Error>;
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct UploadResponse {
     pub id: u64,
 }
 
+#[derive(Clone, Debug)]
 pub struct UploadInfo {
     pub name: String,
     pub contents: Vec<u8>,
@@ -105,6 +107,51 @@ impl SyncBackend for DebugSyncBackend {
     }
 }
 
+/// Performs the retry logic for rate limitation errors. The struct wraps a SyncBackend so that
+/// when a RateLimited error occurs, the thread sleeps for a moment and then tries to reupload the
+/// data.
+pub struct RetryBackend<InnerSyncBackend> {
+    inner: InnerSyncBackend,
+    delay: Duration,
+    attempts: usize,
+}
+
+impl<InnerSyncBackend> RetryBackend<InnerSyncBackend> {
+    /// Creates a new backend from another SyncBackend. The max_retries parameter gives the number
+    /// of times the backend will try again (so given 0, it acts just as the original SyncBackend).
+    /// The delay parameter provides the amount of time to wait between each upload attempt.
+    pub fn new(inner: InnerSyncBackend, max_retries: usize, delay: Duration) -> Self {
+        Self {
+            inner,
+            delay,
+            attempts: max_retries + 1,
+        }
+    }
+}
+
+impl<InnerSyncBackend: SyncBackend> SyncBackend for RetryBackend<InnerSyncBackend> {
+    fn upload(&mut self, data: UploadInfo) -> Result<UploadResponse, Error> {
+        for index in 0..self.attempts {
+            if index != 0 {
+                log::info!(
+                    "tarmac is being rate limited, retrying upload ({}/{})",
+                    index,
+                    self.attempts - 1
+                );
+                thread::sleep(self.delay);
+            }
+            let result = self.inner.upload(data.clone());
+
+            match result {
+                Err(Error::RateLimited) => {}
+                _ => return result,
+            }
+        }
+
+        Err(Error::RateLimited)
+    }
+}
+
 #[derive(Debug, Error)]
 pub enum Error {
     #[error("Cannot upload assets with the 'none' target.")]
@@ -125,3 +172,115 @@ pub enum Error {
         source: RobloxApiError,
     },
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[allow(unused_must_use)]
+    mod test_retry_backend {
+        use super::*;
+
+        struct CountUploads<'a> {
+            counter: &'a mut usize,
+            results: Vec<Result<UploadResponse, Error>>,
+        }
+
+        impl<'a> CountUploads<'a> {
+            fn new(counter: &'a mut usize) -> Self {
+                Self {
+                    counter,
+                    results: Vec::new(),
+                }
+            }
+
+            fn with_results(mut self, results: Vec<Result<UploadResponse, Error>>) -> Self {
+                self.results = results;
+                self.results.reverse();
+                self
+            }
+        }
+
+        impl<'a> SyncBackend for CountUploads<'a> {
+            fn upload(&mut self, _data: UploadInfo) -> Result<UploadResponse, Error> {
+                (*self.counter) += 1;
+                self.results.pop().unwrap_or(Err(Error::NoneBackend))
+            }
+        }
+
+        fn any_upload_info() -> UploadInfo {
+            UploadInfo {
+                name: "foo".to_owned(),
+                contents: Vec::new(),
+                hash: "hash".to_owned(),
+            }
+        }
+
+        fn retry_duration() -> Duration {
+            Duration::from_millis(1)
+        }
+
+        #[test]
+        fn upload_at_least_once() {
+            let mut counter = 0;
+            let mut backend =
+                RetryBackend::new(CountUploads::new(&mut counter), 0, retry_duration());
+
+            backend.upload(any_upload_info());
+
+            assert_eq!(counter, 1);
+        }
+
+        #[test]
+        fn upload_again_if_rate_limited() {
+            let mut counter = 0;
+            let inner = CountUploads::new(&mut counter).with_results(vec![
+                Err(Error::RateLimited),
+                Err(Error::RateLimited),
+                Err(Error::NoneBackend),
+            ]);
+            let mut backend = RetryBackend::new(inner, 5, retry_duration());
+
+            backend.upload(any_upload_info());
+
+            assert_eq!(counter, 3);
+        }
+
+        #[test]
+        fn upload_returns_first_success_result() {
+            let mut counter = 0;
+            let success = UploadResponse { id: 10 };
+            let inner = CountUploads::new(&mut counter).with_results(vec![
+                Err(Error::RateLimited),
+                Err(Error::RateLimited),
+                Ok(success.clone()),
+            ]);
+            let mut backend = RetryBackend::new(inner, 5, retry_duration());
+
+            let upload_result = backend.upload(any_upload_info()).unwrap();
+
+            assert_eq!(counter, 3);
+            assert_eq!(upload_result, success);
+        }
+
+        #[test]
+        fn upload_returns_rate_limited_when_retries_exhausted() {
+            let mut counter = 0;
+            let inner = CountUploads::new(&mut counter).with_results(vec![
+                Err(Error::RateLimited),
+                Err(Error::RateLimited),
+                Err(Error::RateLimited),
+                Err(Error::RateLimited),
+            ]);
+            let mut backend = RetryBackend::new(inner, 2, retry_duration());
+
+            let upload_result = backend.upload(any_upload_info()).unwrap_err();
+
+            assert_eq!(counter, 3);
+            assert!(match upload_result {
+                Error::RateLimited => true,
+                _ => false,
+            });
+        }
+    }
+}