From a09aa25eb9da4aac12731a403d2ff4ebcc439a88 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Mon, 12 Jan 2026 17:36:00 +0000 Subject: [PATCH 1/6] impl(auth): add retry to mds idtoken provider --- src/auth/src/credentials/idtoken/mds.rs | 175 ++++++++++++++++++++++-- 1 file changed, 160 insertions(+), 15 deletions(-) diff --git a/src/auth/src/credentials/idtoken/mds.rs b/src/auth/src/credentials/idtoken/mds.rs index 71c7249a68..3fac2bbd89 100644 --- a/src/auth/src/credentials/idtoken/mds.rs +++ b/src/auth/src/credentials/idtoken/mds.rs @@ -71,6 +71,7 @@ use crate::credentials::mds::{ METADATA_ROOT, }; use crate::errors::CredentialsError; +use crate::retry::{Builder as RetryTokenProviderBuilder, TokenProviderWithRetry}; use crate::token::{CachedTokenProvider, Token, TokenProvider}; use crate::token_cache::TokenCache; use crate::{ @@ -79,6 +80,9 @@ use crate::{ credentials::idtoken::{IDTokenCredentials, parse_id_token_from_str}, }; use async_trait::async_trait; +use gax::backoff_policy::BackoffPolicyArg; +use gax::retry_policy::RetryPolicyArg; +use gax::retry_throttler::RetryThrottlerArg; use http::{Extensions, HeaderValue}; use reqwest::Client; use std::sync::Arc; @@ -135,6 +139,7 @@ pub struct Builder { pub(crate) format: Option, licenses: Option, target_audience: String, + retry_builder: RetryTokenProviderBuilder, } impl Builder { @@ -149,6 +154,7 @@ impl Builder { endpoint: None, licenses: None, target_audience: target_audience.into(), + retry_builder: RetryTokenProviderBuilder::default(), } } @@ -222,26 +228,97 @@ impl Builder { self } - fn build_token_provider(self) -> MDSTokenProvider { - let final_endpoint: String; + /// Configure the retry policy for fetching tokens. + /// + /// The retry policy controls how to handle retries, and sets limits on + /// the number of attempts or the total time spent retrying. + /// + /// # Example + /// + /// ```no_run + /// # use google_cloud_auth::credentials::idtoken; + /// use gax::retry_policy::{AlwaysRetry, RetryPolicyExt}; + /// + /// let audience = "https://my-service.a.run.app"; + /// let credentials = idtoken::mds::Builder::new(audience) + /// .with_retry_policy(AlwaysRetry.with_attempt_limit(3)) + /// .build(); + /// ``` + pub fn with_retry_policy>(mut self, v: V) -> Self { + self.retry_builder = self.retry_builder.with_retry_policy(v.into()); + self + } - // Determine the endpoint and whether it was overridden - if let Ok(host_from_env) = std::env::var(GCE_METADATA_HOST_ENV_VAR) { - // Check GCE_METADATA_HOST environment variable first - final_endpoint = format!("http://{host_from_env}"); - } else if let Some(builder_endpoint) = self.endpoint { - // Else, check if an endpoint was provided to the mds::Builder - final_endpoint = builder_endpoint; - } else { - // Else, use the default metadata root - final_endpoint = METADATA_ROOT.to_string(); - }; + /// Configure the retry backoff policy. + /// + /// The backoff policy controls how long to wait in between retry attempts. + /// + /// # Example + /// + /// ```no_run + /// # use google_cloud_auth::credentials::idtoken; + /// use gax::exponential_backoff::ExponentialBackoff; + /// + /// let audience = "https://my-service.a.run.app"; + /// let credentials = idtoken::mds::Builder::new(audience) + /// .with_backoff_policy(ExponentialBackoff::default()) + /// .build(); + /// ``` + pub fn with_backoff_policy>(mut self, v: V) -> Self { + self.retry_builder = self.retry_builder.with_backoff_policy(v.into()); + self + } - MDSTokenProvider { + /// Configure the retry throttler. + /// + /// Advanced applications may want to configure a retry throttler to + /// [Address Cascading Failures] and when [Handling Overload] conditions. + /// The authentication library throttles its retry loop, using a policy to + /// control the throttling algorithm. Use this method to fine tune or + /// customize the default retry throttler. + /// + /// [Handling Overload]: https://sre.google/sre-book/handling-overload/ + /// [Address Cascading Failures]: https://sre.google/sre-book/addressing-cascading-failures/ + /// + /// # Example + /// + /// ```no_run + /// # use google_cloud_auth::credentials::idtoken; + /// use gax::retry_throttler::AdaptiveThrottler; + /// + /// let audience = "https://my-service.a.run.app"; + /// let credentials = idtoken::mds::Builder::new(audience) + /// .with_retry_throttler(AdaptiveThrottler::default()) + /// .build(); + /// ``` + pub fn with_retry_throttler>(mut self, v: V) -> Self { + self.retry_builder = self.retry_builder.with_retry_throttler(v.into()); + self + } + + fn build_token_provider(self) -> TokenProviderWithRetry { + let final_endpoint = self.resolve_endpoint(); + + let tp = MDSTokenProvider { format: self.format, licenses: self.licenses, endpoint: final_endpoint, target_audience: self.target_audience, + }; + self.retry_builder.build(tp) + } + + fn resolve_endpoint(&self) -> String { + // Determine the endpoint + if let Ok(host_from_env) = std::env::var(GCE_METADATA_HOST_ENV_VAR) { + // Check GCE_METADATA_HOST environment variable first + format!("http://{host_from_env}") + } else if let Some(builder_endpoint) = self.endpoint.clone() { + // Else, check if an endpoint was provided to the mds::Builder + builder_endpoint + } else { + // Else, use the default metadata root + METADATA_ROOT.to_string() } } @@ -308,7 +385,11 @@ impl TokenProvider for MDSTokenProvider { mod tests { use super::*; use crate::credentials::idtoken::tests::generate_test_id_token; - use crate::credentials::tests::find_source_error; + use crate::credentials::tests::{ + find_source_error, get_mock_auth_retry_policy, get_mock_backoff_policy, + get_mock_retry_throttler, + }; + use httptest::cycle; use httptest::matchers::{all_of, contains, request, url_decoded}; use httptest::responders::status_code; use httptest::{Expectation, Server}; @@ -319,6 +400,70 @@ mod tests { type TestResult = anyhow::Result<()>; + #[tokio::test] + #[parallel] + async fn test_mds_does_not_retry_on_non_transient_failures() -> TestResult { + let mut server = Server::run(); + let audience = "test-audience"; + server.expect( + Expectation::matching(all_of![ + request::path(format!("{MDS_DEFAULT_URI}/identity")), + request::query(url_decoded(contains(("audience", audience)))), + ]) + .times(1) + .respond_with(status_code(401)), + ); + + let creds = Builder::new(audience) + .with_endpoint(format!("http://{}", server.addr())) + .with_retry_policy(get_mock_auth_retry_policy(3)) + .with_backoff_policy(get_mock_backoff_policy()) + .with_retry_throttler(get_mock_retry_throttler()) + .build()?; + + let err = creds.id_token().await.unwrap_err(); + let source = find_source_error::(&err); + assert!( + matches!(source, Some(e) if e.status() == Some(StatusCode::UNAUTHORIZED)), + "{err:?}" + ); + server.verify_and_clear(); + Ok(()) + } + + #[tokio::test] + #[parallel] + async fn test_mds_retries_for_success() -> TestResult { + let server = Server::run(); + let audience = "test-audience"; + let token_string = generate_test_id_token(audience); + + server.expect( + Expectation::matching(all_of![ + request::path(format!("{MDS_DEFAULT_URI}/identity")), + request::query(url_decoded(contains(("audience", audience)))), + ]) + .times(3) + .respond_with(cycle![ + status_code(503).body("try-again"), + status_code(503).body("try-again"), + status_code(200).body(token_string.clone()), + ]), + ); + + let creds = Builder::new(audience) + .with_endpoint(format!("http://{}", server.addr())) + .with_retry_policy(get_mock_auth_retry_policy(3)) + .with_backoff_policy(get_mock_backoff_policy()) + .with_retry_throttler(get_mock_retry_throttler()) + .build()?; + + let id_token = creds.id_token().await?; + assert_eq!(id_token, token_string); + + Ok(()) + } + #[tokio::test] #[test_case(Format::Standard)] #[test_case(Format::Full)] From 46639816ef672cf8a4f79d11705d69449b2f228a Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Mon, 12 Jan 2026 17:55:55 +0000 Subject: [PATCH 2/6] fix: httptest server doesnt need to be mut --- src/auth/src/credentials/mds.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/auth/src/credentials/mds.rs b/src/auth/src/credentials/mds.rs index bf42de3471..c439dab937 100644 --- a/src/auth/src/credentials/mds.rs +++ b/src/auth/src/credentials/mds.rs @@ -525,7 +525,7 @@ mod tests { #[tokio::test] #[parallel] async fn test_mds_does_not_retry_on_non_transient_failures() -> TestResult { - let mut server = Server::run(); + let server = Server::run(); server.expect( Expectation::matching(request::path(format!("{MDS_DEFAULT_URI}/token"))) .times(1) @@ -541,7 +541,7 @@ mod tests { let err = provider.token().await.unwrap_err(); assert!(!err.is_transient()); - server.verify_and_clear(); + Ok(()) } From d5dfb4c600f3dfa47451476d8667c2473ac6cdd5 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Mon, 12 Jan 2026 17:59:44 +0000 Subject: [PATCH 3/6] fix: httptest changes only on new stuff --- src/auth/src/credentials/idtoken/mds.rs | 4 ++-- src/auth/src/credentials/mds.rs | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/auth/src/credentials/idtoken/mds.rs b/src/auth/src/credentials/idtoken/mds.rs index 3fac2bbd89..278dafc0a8 100644 --- a/src/auth/src/credentials/idtoken/mds.rs +++ b/src/auth/src/credentials/idtoken/mds.rs @@ -403,7 +403,7 @@ mod tests { #[tokio::test] #[parallel] async fn test_mds_does_not_retry_on_non_transient_failures() -> TestResult { - let mut server = Server::run(); + let server = Server::run(); let audience = "test-audience"; server.expect( Expectation::matching(all_of![ @@ -427,7 +427,7 @@ mod tests { matches!(source, Some(e) if e.status() == Some(StatusCode::UNAUTHORIZED)), "{err:?}" ); - server.verify_and_clear(); + Ok(()) } diff --git a/src/auth/src/credentials/mds.rs b/src/auth/src/credentials/mds.rs index c439dab937..bf42de3471 100644 --- a/src/auth/src/credentials/mds.rs +++ b/src/auth/src/credentials/mds.rs @@ -525,7 +525,7 @@ mod tests { #[tokio::test] #[parallel] async fn test_mds_does_not_retry_on_non_transient_failures() -> TestResult { - let server = Server::run(); + let mut server = Server::run(); server.expect( Expectation::matching(request::path(format!("{MDS_DEFAULT_URI}/token"))) .times(1) @@ -541,7 +541,7 @@ mod tests { let err = provider.token().await.unwrap_err(); assert!(!err.is_transient()); - + server.verify_and_clear(); Ok(()) } From d26fe70d034ed35f352c27c8891c52e7cc492dd2 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Thu, 29 Jan 2026 17:22:02 +0000 Subject: [PATCH 4/6] fix: make TokenProviderWithRetry RefUnwindSafe --- Cargo.lock | 1 + src/auth/Cargo.toml | 27 ++++++----- src/auth/src/retry.rs | 108 ++++++++++++++++++++++++++++++++---------- 3 files changed, 99 insertions(+), 37 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cdf5d7cea8..20d3bc9797 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1473,6 +1473,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "static_assertions", "tempfile", "test-case", "thiserror", diff --git a/src/auth/Cargo.toml b/src/auth/Cargo.toml index a12970b2a1..35731eb3b8 100644 --- a/src/auth/Cargo.toml +++ b/src/auth/Cargo.toml @@ -52,19 +52,20 @@ jsonwebtoken = { workspace = true, features = ["rust_crypto"], optional gax.workspace = true [dev-dependencies] -anyhow.workspace = true -httptest.workspace = true -mockall.workspace = true -regex.workspace = true -rsa = { workspace = true, features = ["pem"] } -scoped-env.workspace = true -serial_test.workspace = true -tempfile.workspace = true -test-case.workspace = true -tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] } -tokio-test.workspace = true -url.workspace = true -mutants.workspace = true +anyhow.workspace = true +httptest.workspace = true +mockall.workspace = true +regex.workspace = true +rsa = { workspace = true, features = ["pem"] } +scoped-env.workspace = true +serial_test.workspace = true +tempfile.workspace = true +test-case.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] } +tokio-test.workspace = true +url.workspace = true +mutants.workspace = true +static_assertions.workspace = true [features] default = [] diff --git a/src/auth/src/retry.rs b/src/auth/src/retry.rs index 9d9036964c..e70bb84766 100644 --- a/src/auth/src/retry.rs +++ b/src/auth/src/retry.rs @@ -21,31 +21,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 } @@ -55,19 +109,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), @@ -140,6 +194,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}; @@ -220,7 +275,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(); @@ -251,7 +306,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(); @@ -267,7 +322,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(); @@ -286,7 +341,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(); @@ -377,8 +432,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()); } @@ -446,8 +501,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); @@ -477,4 +532,9 @@ mod tests { original_error_string ); } + + #[test] + fn test_unwind_safe() { + assert_impl_all!(Builder: std::panic::UnwindSafe, std::panic::RefUnwindSafe); + } } From d5a98099b6dbb2ced1e8d59dae08f8cad6abe817 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Thu, 29 Jan 2026 21:04:02 +0000 Subject: [PATCH 5/6] impl: simplify unwind safety --- src/auth/src/retry.rs | 115 ++++++++++++++---------------------------- 1 file changed, 37 insertions(+), 78 deletions(-) diff --git a/src/auth/src/retry.rs b/src/auth/src/retry.rs index d2d9908fcf..16e56fa290 100644 --- a/src/auth/src/retry.rs +++ b/src/auth/src/retry.rs @@ -23,85 +23,44 @@ 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 solves that by manually implementing `RefUnwindSafe` and `UnwindSafe`, 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` and `UnwindSafe`. +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 +70,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 +236,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 +267,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 +283,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 +302,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 +393,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 +462,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); From 71e6995e3c2a8e67235ed0628161030fe350b2ed Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Fri, 30 Jan 2026 16:42:15 +0000 Subject: [PATCH 6/6] fix: address review comments --- src/auth/src/retry.rs | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/auth/src/retry.rs b/src/auth/src/retry.rs index 16e56fa290..4c5f5b1ca9 100644 --- a/src/auth/src/retry.rs +++ b/src/auth/src/retry.rs @@ -41,15 +41,15 @@ pub(crate) struct Builder { 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 solves that by manually implementing `RefUnwindSafe` and `UnwindSafe`, 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` and `UnwindSafe`. +// 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 solves that by manually implementing `RefUnwindSafe` and `UnwindSafe`, 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` and `UnwindSafe`. impl RefUnwindSafe for Builder {} impl UnwindSafe for Builder {}