Flush device cache correctly on shutdown

2 files changed, 15 insertions, 3 deletions

diff --git a/mullvad-daemon/src/device.rs b/mullvad-daemon/src/device.rs
index 0e934626ef..470cfd7c89 100644
--- a/mullvad-daemon/src/device.rs
+++ b/mullvad-daemon/src/device.rs
@@ -295,14 +295,24 @@ impl AccountManager {
             abort_handle.abort();
         }
     }
+
+    /// Consumes the object and completes when there is nothing left to write to
+    /// the cache file.
+    pub fn finalize(mut self) -> impl Future<Output = ()> {
+        let join_handle = self.cache_task_join_handle.take();
+        drop(self);
+
+        async move {
+            if let Some(join_handle) = join_handle {
+                let _ = join_handle.await;
+            }
+        }
+    }
 }
 
 impl Drop for AccountManager {
     fn drop(&mut self) {
         self.stop_key_rotation();
-        if let Some(cache_task_join_handle) = self.cache_task_join_handle.take() {
-            let _ = self.runtime.block_on(cache_task_join_handle);
-        }
     }
 }
 
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index c51cb8510e..464824fc59 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -880,6 +880,7 @@ where
         let Daemon {
             event_listener,
             mut shutdown_tasks,
+            account_manager,
             rpc_runtime,
             tunnel_state_machine_handle,
             target_state,
@@ -887,6 +888,7 @@ where
         } = self;
 
         shutdown_tasks.push(Box::pin(target_state.finalize()));
+        shutdown_tasks.insert(0, Box::pin(account_manager.finalize()));
 
         (
             event_listener,