diff --git a/common/bandwidth-controller/src/traits.rs b/common/bandwidth-controller/src/traits.rs index bb958f482b..1ff14027ef 100644 --- a/common/bandwidth-controller/src/traits.rs +++ b/common/bandwidth-controller/src/traits.rs @@ -25,6 +25,8 @@ pub trait BandwidthTicketProvider: Send + Sync { ) -> Result; async fn get_upgrade_mode_token(&self) -> Result, BandwidthControllerError>; + + async fn close(&self) {} } #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] @@ -56,6 +58,10 @@ where .map_err(|_| BandwidthControllerError::MalformedUpgradeModeToken)?; Ok(Some(token)) } + + async fn close(&self) { + self.storage.close().await; + } } #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] @@ -75,4 +81,8 @@ impl BandwidthTicketProvider for Box async fn get_upgrade_mode_token(&self) -> Result, BandwidthControllerError> { (**self).get_upgrade_mode_token().await } + + async fn close(&self) { + (**self).close().await; + } } diff --git a/common/client-core/src/client/base_client/mod.rs b/common/client-core/src/client/base_client/mod.rs index 7b27d7b1b1..253da7d6c6 100644 --- a/common/client-core/src/client/base_client/mod.rs +++ b/common/client-core/src/client/base_client/mod.rs @@ -1023,6 +1023,16 @@ where let encryption_keys = init_res.client_keys.encryption_keypair(); let identity_keys = init_res.client_keys.identity_keypair(); + let credential_store_for_close = credential_store.clone(); + let close_credential_token = shutdown_tracker.clone_shutdown_token(); + shutdown_tracker.try_spawn_named( + async move { + close_credential_token.cancelled().await; + credential_store_for_close.close().await; + }, + "CredentialStorage::close_on_shutdown", + ); + // the components are started in very specific order. Unless you know what you are doing, // do not change that. let bandwidth_controller = self