From fe2f63c4d229a28cac27cf3e240e07cbd52dc2b0 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Thu, 29 Jan 2026 17:22:02 +0000 Subject: [PATCH 1/2] impl(auth): make TokenProviderWithRetry RefUnwindSafe --- Cargo.lock | 1 + src/auth/Cargo.toml | 1 + src/auth/src/retry.rs | 108 ++++++++++++++++++++++++++++++++---------- 3 files changed, 86 insertions(+), 24 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0ce25a785b..3e797f1056 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1503,6 +1503,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "static_assertions", "tempfile", "test-case", "thiserror 2.0.18", diff --git a/src/auth/Cargo.toml b/src/auth/Cargo.toml index 63e7c37d42..556fe5a815 100644 --- a/src/auth/Cargo.toml +++ b/src/auth/Cargo.toml @@ -68,6 +68,7 @@ tokio = { workspace = true, features = ["macros", "rt-multi-thre tokio-test.workspace = true url.workspace = true mutants.workspace = true +static_assertions = { workspace = true } [features] default = ["default-idtoken-backend", "default-rustls-provider"] diff --git a/src/auth/src/retry.rs b/src/auth/src/retry.rs index 6fbd2a5c51..d2d9908fcf 100644 --- a/src/auth/src/retry.rs +++ b/src/auth/src/retry.rs @@ -23,31 +23,85 @@ use gax::retry_loop_internal::retry_loop; use gax::retry_policy::{AlwaysRetry, RetryPolicy, RetryPolicyArg, RetryPolicyExt}; use gax::retry_throttler::{AdaptiveThrottler, RetryThrottlerArg, SharedRetryThrottler}; use std::error::Error; +use std::panic::RefUnwindSafe; use std::sync::{Arc, Mutex}; +/// A wrapper to assert `RefUnwindSafe` on the inner type. +/// +/// This is necessary because adding [TokenProviderWithRetry] to existing, released auth features +/// caused a breaking change. The containing structs would lose their automatically derived +/// `RefUnwindSafe` implementation because the dynamic trait objects (like `dyn RetryPolicy`) +/// are not `RefUnwindSafe` by default. +/// +/// This wrapper solves that by manually implementing `RefUnwindSafe`, allowing us to add +/// retry functionality without triggering a semver-check failure. This is safe because +/// we control the implementation of the inner types and can ensure that they are +/// `RefUnwindSafe`. +#[derive(Debug)] +struct UnwindSafeAdapter(T); + +impl RefUnwindSafe for UnwindSafeAdapter {} + +impl From for UnwindSafeAdapter { + fn from(inner: T) -> Self { + Self(inner) + } +} + +impl RetryPolicy for UnwindSafeAdapter> { + fn on_error( + &self, + state: &gax::retry_state::RetryState, + error: gax::error::Error, + ) -> gax::retry_result::RetryResult { + self.0.on_error(state, error) + } + fn on_throttle( + &self, + state: &gax::retry_state::RetryState, + error: gax::error::Error, + ) -> gax::throttle_result::ThrottleResult { + self.0.on_throttle(state, error) + } + fn remaining_time(&self, state: &gax::retry_state::RetryState) -> Option { + self.0.remaining_time(state) + } +} + +impl BackoffPolicy for UnwindSafeAdapter> { + fn on_failure(&self, state: &gax::retry_state::RetryState) -> std::time::Duration { + self.0.on_failure(state) + } +} + #[derive(Debug)] pub(crate) struct TokenProviderWithRetry { pub(crate) inner: Arc, - retry_policy: Arc, - backoff_policy: Arc, + retry_policy: Arc, + backoff_policy: Arc, retry_throttler: SharedRetryThrottler, } #[derive(Debug, Default)] pub(crate) struct Builder { - retry_policy: Option, - backoff_policy: Option, + retry_policy: Option>, + backoff_policy: Option>, retry_throttler: Option, } impl Builder { - pub(crate) fn with_retry_policy(mut self, retry_policy: RetryPolicyArg) -> Self { - self.retry_policy = Some(retry_policy); + pub(crate) fn with_retry_policy>(mut self, retry_policy: P) -> Self { + let inner: Arc = retry_policy.into().into(); + self.retry_policy = Some(Arc::new(UnwindSafeAdapter::from(inner))); self } - pub(crate) fn with_backoff_policy(mut self, backoff_policy: BackoffPolicyArg) -> Self { - self.backoff_policy = Some(backoff_policy); + pub(crate) fn with_backoff_policy>( + mut self, + backoff_policy: P, + ) -> Self { + let inner: Arc = backoff_policy.into().into(); + self.backoff_policy = Some(Arc::new(UnwindSafeAdapter::from(inner))); self } @@ -57,19 +111,19 @@ impl Builder { } pub(crate) fn build(self, token_provider: T) -> TokenProviderWithRetry { - let backoff_policy: Arc = match self.backoff_policy { - Some(p) => p.into(), - None => Arc::new(ExponentialBackoff::default()), - }; + let backoff_policy = self.backoff_policy.unwrap_or_else(|| { + let p: Arc = Arc::new(ExponentialBackoff::default()); + Arc::new(UnwindSafeAdapter::from(p)) + }); let retry_throttler: SharedRetryThrottler = match self.retry_throttler { Some(p) => p.into(), None => Arc::new(Mutex::new(AdaptiveThrottler::default())), }; - let retry_policy = self - .retry_policy - .unwrap_or_else(|| AlwaysRetry.with_attempt_limit(1).into()) - .into(); + let retry_policy = self.retry_policy.unwrap_or_else(|| { + let p: Arc = Arc::new(AlwaysRetry.with_attempt_limit(1)); + Arc::new(UnwindSafeAdapter::from(p)) + }); TokenProviderWithRetry { inner: Arc::new(token_provider), @@ -142,6 +196,7 @@ mod tests { use gax::retry_state::RetryState; use gax::retry_throttler::RetryThrottler; use mockall::{Sequence, mock}; + use static_assertions::assert_impl_all; use std::error::Error; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; @@ -222,7 +277,7 @@ mod tests { .return_once(|| Ok(token)); let provider = Builder::default() - .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }.into()) + .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }) .build(mock_provider); let token = provider.token().await.unwrap(); @@ -253,7 +308,7 @@ mod tests { }); let provider = Builder::default() - .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }.into()) + .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }) .build(mock_provider); let token = provider.token().await.unwrap(); @@ -269,7 +324,7 @@ mod tests { .returning(|| Err(CredentialsError::from_msg(true, "transient error"))); let provider = Builder::default() - .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }.into()) + .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }) .build(mock_provider); let error = provider.token().await.unwrap_err(); @@ -288,7 +343,7 @@ mod tests { .returning(|| Err(CredentialsError::from_msg(false, "non transient error"))); let provider = Builder::default() - .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }.into()) + .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }) .build(mock_provider); let error = provider.token().await.unwrap_err(); @@ -379,8 +434,8 @@ mod tests { let backoff_policy = TestBackoffPolicy::default(); let retry_throttler = AdaptiveThrottler::new(4.0).unwrap(); builder = builder - .with_retry_policy(retry_policy.into()) - .with_backoff_policy(backoff_policy.into()) + .with_retry_policy(retry_policy) + .with_backoff_policy(backoff_policy) .with_retry_throttler(retry_throttler.into()); } @@ -448,8 +503,8 @@ mod tests { // 4. Build and run let provider = Builder::default() - .with_retry_policy(retry_policy.into()) - .with_backoff_policy(backoff_policy.into()) + .with_retry_policy(retry_policy) + .with_backoff_policy(backoff_policy) .with_retry_throttler(retry_throttler.into()) .build(mock_provider); @@ -479,4 +534,9 @@ mod tests { original_error_string ); } + + #[test] + fn test_unwind_safe() { + assert_impl_all!(Builder: std::panic::UnwindSafe, std::panic::RefUnwindSafe); + } } From 5a0845369e5d3007ff6a26fa2c3f1206937334b6 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Fri, 30 Jan 2026 16:41:45 +0000 Subject: [PATCH 2/2] impl: simplify unwind safety impl and improve docs --- src/auth/src/retry.rs | 114 +++++++++++++----------------------------- 1 file changed, 36 insertions(+), 78 deletions(-) diff --git a/src/auth/src/retry.rs b/src/auth/src/retry.rs index d2d9908fcf..dbc6534d12 100644 --- a/src/auth/src/retry.rs +++ b/src/auth/src/retry.rs @@ -23,85 +23,43 @@ use gax::retry_loop_internal::retry_loop; use gax::retry_policy::{AlwaysRetry, RetryPolicy, RetryPolicyArg, RetryPolicyExt}; use gax::retry_throttler::{AdaptiveThrottler, RetryThrottlerArg, SharedRetryThrottler}; use std::error::Error; -use std::panic::RefUnwindSafe; +use std::panic::{RefUnwindSafe, UnwindSafe}; use std::sync::{Arc, Mutex}; -/// A wrapper to assert `RefUnwindSafe` on the inner type. -/// -/// This is necessary because adding [TokenProviderWithRetry] to existing, released auth features -/// caused a breaking change. The containing structs would lose their automatically derived -/// `RefUnwindSafe` implementation because the dynamic trait objects (like `dyn RetryPolicy`) -/// are not `RefUnwindSafe` by default. -/// -/// This wrapper solves that by manually implementing `RefUnwindSafe`, allowing us to add -/// retry functionality without triggering a semver-check failure. This is safe because -/// we control the implementation of the inner types and can ensure that they are -/// `RefUnwindSafe`. -#[derive(Debug)] -struct UnwindSafeAdapter(T); - -impl RefUnwindSafe for UnwindSafeAdapter {} - -impl From for UnwindSafeAdapter { - fn from(inner: T) -> Self { - Self(inner) - } -} - -impl RetryPolicy for UnwindSafeAdapter> { - fn on_error( - &self, - state: &gax::retry_state::RetryState, - error: gax::error::Error, - ) -> gax::retry_result::RetryResult { - self.0.on_error(state, error) - } - fn on_throttle( - &self, - state: &gax::retry_state::RetryState, - error: gax::error::Error, - ) -> gax::throttle_result::ThrottleResult { - self.0.on_throttle(state, error) - } - fn remaining_time(&self, state: &gax::retry_state::RetryState) -> Option { - self.0.remaining_time(state) - } -} - -impl BackoffPolicy for UnwindSafeAdapter> { - fn on_failure(&self, state: &gax::retry_state::RetryState) -> std::time::Duration { - self.0.on_failure(state) - } -} - #[derive(Debug)] pub(crate) struct TokenProviderWithRetry { pub(crate) inner: Arc, - retry_policy: Arc, - backoff_policy: Arc, + retry_policy: Arc, + backoff_policy: Arc, retry_throttler: SharedRetryThrottler, } #[derive(Debug, Default)] pub(crate) struct Builder { - retry_policy: Option>, - backoff_policy: Option>, + retry_policy: Option, + backoff_policy: Option, retry_throttler: Option, } +// This is necessary because adding [Builder] to existing, released auth features +// caused a breaking change. The containing structs would lose their automatically derived +// `RefUnwindSafe` implementation because the dynamic trait objects (like `dyn RetryPolicy`) +// are not `RefUnwindSafe` nor `UnwindSafe` by default. +// +// This is safe because in the builder, we never call or use the retry policies. We just hold them +// until we create the client. There is no opportunity for a panic to leave them in an inconsistent +// state. +impl RefUnwindSafe for Builder {} +impl UnwindSafe for Builder {} + impl Builder { - pub(crate) fn with_retry_policy>(mut self, retry_policy: P) -> Self { - let inner: Arc = retry_policy.into().into(); - self.retry_policy = Some(Arc::new(UnwindSafeAdapter::from(inner))); + pub(crate) fn with_retry_policy(mut self, retry_policy: RetryPolicyArg) -> Self { + self.retry_policy = Some(retry_policy); self } - pub(crate) fn with_backoff_policy>( - mut self, - backoff_policy: P, - ) -> Self { - let inner: Arc = backoff_policy.into().into(); - self.backoff_policy = Some(Arc::new(UnwindSafeAdapter::from(inner))); + pub(crate) fn with_backoff_policy(mut self, backoff_policy: BackoffPolicyArg) -> Self { + self.backoff_policy = Some(backoff_policy); self } @@ -111,19 +69,19 @@ impl Builder { } pub(crate) fn build(self, token_provider: T) -> TokenProviderWithRetry { - let backoff_policy = self.backoff_policy.unwrap_or_else(|| { - let p: Arc = Arc::new(ExponentialBackoff::default()); - Arc::new(UnwindSafeAdapter::from(p)) - }); + let backoff_policy: Arc = match self.backoff_policy { + Some(p) => p.into(), + None => Arc::new(ExponentialBackoff::default()), + }; let retry_throttler: SharedRetryThrottler = match self.retry_throttler { Some(p) => p.into(), None => Arc::new(Mutex::new(AdaptiveThrottler::default())), }; - let retry_policy = self.retry_policy.unwrap_or_else(|| { - let p: Arc = Arc::new(AlwaysRetry.with_attempt_limit(1)); - Arc::new(UnwindSafeAdapter::from(p)) - }); + let retry_policy = self + .retry_policy + .unwrap_or_else(|| AlwaysRetry.with_attempt_limit(1).into()) + .into(); TokenProviderWithRetry { inner: Arc::new(token_provider), @@ -277,7 +235,7 @@ mod tests { .return_once(|| Ok(token)); let provider = Builder::default() - .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }) + .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }.into()) .build(mock_provider); let token = provider.token().await.unwrap(); @@ -308,7 +266,7 @@ mod tests { }); let provider = Builder::default() - .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }) + .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }.into()) .build(mock_provider); let token = provider.token().await.unwrap(); @@ -324,7 +282,7 @@ mod tests { .returning(|| Err(CredentialsError::from_msg(true, "transient error"))); let provider = Builder::default() - .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }) + .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }.into()) .build(mock_provider); let error = provider.token().await.unwrap_err(); @@ -343,7 +301,7 @@ mod tests { .returning(|| Err(CredentialsError::from_msg(false, "non transient error"))); let provider = Builder::default() - .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }) + .with_retry_policy(AuthRetryPolicy { max_attempts: 2 }.into()) .build(mock_provider); let error = provider.token().await.unwrap_err(); @@ -434,8 +392,8 @@ mod tests { let backoff_policy = TestBackoffPolicy::default(); let retry_throttler = AdaptiveThrottler::new(4.0).unwrap(); builder = builder - .with_retry_policy(retry_policy) - .with_backoff_policy(backoff_policy) + .with_retry_policy(retry_policy.into()) + .with_backoff_policy(backoff_policy.into()) .with_retry_throttler(retry_throttler.into()); } @@ -503,8 +461,8 @@ mod tests { // 4. Build and run let provider = Builder::default() - .with_retry_policy(retry_policy) - .with_backoff_policy(backoff_policy) + .with_retry_policy(retry_policy.into()) + .with_backoff_policy(backoff_policy.into()) .with_retry_throttler(retry_throttler.into()) .build(mock_provider);