Simplify retry logic

2 files changed, 42 insertions, 69 deletions

diff --git a/mullvad-daemon/src/device.rs b/mullvad-daemon/src/device.rs
index 708c17f908..111d0df9e1 100644
--- a/mullvad-daemon/src/device.rs
+++ b/mullvad-daemon/src/device.rs
@@ -424,28 +424,13 @@ impl DeviceService {
         let private_key = PrivateKey::new_from_random();
         let pubkey = private_key.public_key();
 
-        let retry_strategy = Jittered::jitter(
-            ExponentialBackoff::new(
-                RETRY_BACKOFF_INTERVAL_INITIAL,
-                RETRY_BACKOFF_INTERVAL_FACTOR,
-            )
-            .max_delay(RETRY_BACKOFF_INTERVAL_MAX),
-        );
-
         let proxy = self.proxy.clone();
         let api_handle = self.api_availability.clone();
         let token_copy = token.clone();
         let (device, addresses) = retry_future(
-            move || {
-                let wait_online = api_handle.wait_online();
-                let fut = proxy.create(token_copy.clone(), pubkey.clone());
-                async move {
-                    let _ = wait_online.await;
-                    fut.await
-                }
-            },
+            move || api_handle.when_online(proxy.create(token_copy.clone(), pubkey.clone())),
             should_retry_backoff,
-            retry_strategy,
+            retry_strategy(),
         )
         .await
         .map_err(map_rest_error)?;
@@ -491,14 +476,8 @@ impl DeviceService {
         );
 
         retry_future(
-            move || {
-                let wait_online = api_handle.wait_online();
-                let fut = proxy.remove(token.clone(), device.clone());
-                async move {
-                    let _ = wait_online.await;
-                    fut.await
-                }
-            },
+            // NOTE: Not honoring "paused" state, because the account may have no time on it.
+            move || api_handle.when_online(proxy.remove(token.clone(), device.clone())),
             should_retry_backoff,
             retry_strategy,
         )
@@ -545,24 +524,16 @@ impl DeviceService {
         let api_handle = self.api_availability.clone();
         let pubkey = private_key.public_key();
 
-        let retry_strategy = Jittered::jitter(
-            ExponentialBackoff::new(
-                RETRY_BACKOFF_INTERVAL_INITIAL,
-                RETRY_BACKOFF_INTERVAL_FACTOR,
-            )
-            .max_delay(RETRY_BACKOFF_INTERVAL_MAX),
-        );
         let addresses = retry_future(
             move || {
-                let wait_online = api_handle.wait_online();
-                let fut = proxy.replace_wg_key(token.clone(), device.clone(), pubkey.clone());
-                async move {
-                    let _ = wait_online.await;
-                    fut.await
-                }
+                api_handle.when_bg_resumes(proxy.replace_wg_key(
+                    token.clone(),
+                    device.clone(),
+                    pubkey.clone(),
+                ))
             },
             should_retry_backoff,
-            retry_strategy,
+            retry_strategy(),
         )
         .await
         .map_err(map_rest_error)?;
@@ -594,24 +565,10 @@ impl DeviceService {
         let proxy = self.proxy.clone();
         let api_handle = self.api_availability.clone();
 
-        let retry_strategy = Jittered::jitter(
-            ExponentialBackoff::new(
-                RETRY_BACKOFF_INTERVAL_INITIAL,
-                RETRY_BACKOFF_INTERVAL_FACTOR,
-            )
-            .max_delay(RETRY_BACKOFF_INTERVAL_MAX),
-        );
         retry_future(
-            move || {
-                let wait_online = api_handle.wait_online();
-                let fut = proxy.list(token.clone());
-                async move {
-                    let _ = wait_online.await;
-                    fut.await
-                }
-            },
+            move || api_handle.when_online(proxy.list(token.clone())),
             should_retry_backoff,
-            retry_strategy,
+            retry_strategy(),
         )
         .await
         .map_err(map_rest_error)
@@ -788,24 +745,14 @@ impl Account {
                 return;
             };
 
-            let retry_strategy = Jittered::jitter(
-                ExponentialBackoff::new(
-                    RETRY_BACKOFF_INTERVAL_INITIAL,
-                    RETRY_BACKOFF_INTERVAL_FACTOR,
-                )
-                .max_delay(RETRY_BACKOFF_INTERVAL_MAX),
-            );
             let future_generator = move || {
-                let wait_online = api_availability.wait_online();
-                let expiry_fut = accounts_proxy.get_expiry(token.clone());
+                let expiry_fut =
+                    api_availability.when_online(accounts_proxy.get_expiry(token.clone()));
                 let api_availability_copy = api_availability.clone();
-                async move {
-                    let _ = wait_online.await;
-                    handle_expiry_result_inner(&expiry_fut.await, &api_availability_copy)
-                }
+                async move { handle_expiry_result_inner(&expiry_fut.await, &api_availability_copy) }
             };
             let should_retry = move |state_was_updated: &bool| -> bool { !*state_was_updated };
-            retry_future(future_generator, should_retry, retry_strategy).await;
+            retry_future(future_generator, should_retry, retry_strategy()).await;
         });
         runtime.spawn(future);
 
@@ -879,3 +826,13 @@ fn map_rest_error(error: rest::Error) -> Error {
         error => Error::OtherRestError(error),
     }
 }
+
+fn retry_strategy() -> Jittered<ExponentialBackoff> {
+    Jittered::jitter(
+        ExponentialBackoff::new(
+            RETRY_BACKOFF_INTERVAL_INITIAL,
+            RETRY_BACKOFF_INTERVAL_FACTOR,
+        )
+        .max_delay(RETRY_BACKOFF_INTERVAL_MAX),
+    )
+}
diff --git a/mullvad-rpc/src/availability.rs b/mullvad-rpc/src/availability.rs
index da8d624e80..2cf40cf53b 100644
--- a/mullvad-rpc/src/availability.rs
+++ b/mullvad-rpc/src/availability.rs
@@ -122,10 +122,26 @@ impl ApiAvailabilityHandle {
         self.wait_for_state(|state| !state.is_suspended())
     }
 
+    pub fn when_bg_resumes<F: Future<Output = O>, O>(&self, task: F) -> impl Future<Output = O> {
+        let wait_task = self.wait_for_state(|state| !state.is_background_paused());
+        async move {
+            let _ = wait_task.await;
+            task.await
+        }
+    }
+
     pub fn wait_background(&self) -> impl Future<Output = Result<(), Error>> {
         self.wait_for_state(|state| !state.is_background_paused())
     }
 
+    pub fn when_online<F: Future<Output = O>, O>(&self, task: F) -> impl Future<Output = O> {
+        let wait_task = self.wait_for_state(|state| !state.is_offline());
+        async move {
+            let _ = wait_task.await;
+            task.await
+        }
+    }
+
     pub fn wait_online(&self) -> impl Future<Output = Result<(), Error>> {
         self.wait_for_state(|state| !state.is_offline())
     }