From bd080d679d7712316e232f1aaecb15b1684e9042 Mon Sep 17 00:00:00 2001 From: akshayakumar-t Date: Thu, 2 Jul 2026 21:03:46 +0530 Subject: [PATCH] fix(gcp): add exponential backoff to GCP token refresh retries --- lib/vector-common/src/backoff.rs | 214 +++++++++++++++++++++++++++++++ lib/vector-common/src/lib.rs | 2 + lib/vector-core/src/gcp.rs | 10 +- src/common/backoff.rs | 212 +----------------------------- 4 files changed, 226 insertions(+), 212 deletions(-) create mode 100644 lib/vector-common/src/backoff.rs diff --git a/lib/vector-common/src/backoff.rs b/lib/vector-common/src/backoff.rs new file mode 100644 index 000000000..89fd01711 --- /dev/null +++ b/lib/vector-common/src/backoff.rs @@ -0,0 +1,214 @@ +use std::time::Duration; + +// `tokio-retry` crate +// MIT License +// Copyright (c) 2017 Sam Rijs +// +/// A retry strategy driven by exponential back-off. +/// +/// The power corresponds to the number of past attempts. +#[derive(Debug, Clone)] +pub struct ExponentialBackoff { + current: u64, + base: u64, + factor: u64, + max_delay: Option, +} + +impl ExponentialBackoff { + /// Constructs a new exponential back-off strategy, + /// given a base duration in milliseconds. + /// + /// The resulting duration is calculated by taking the base to the `n`-th power, + /// where `n` denotes the number of past attempts. + pub const fn from_millis(base: u64) -> ExponentialBackoff { + ExponentialBackoff { + current: base, + base, + factor: 1u64, + max_delay: None, + } + } + + /// A multiplicative factor that will be applied to the retry delay. + /// + /// For example, using a factor of `1000` will make each delay in units of seconds. + /// + /// Default factor is `1`. + pub const fn factor(mut self, factor: u64) -> ExponentialBackoff { + self.factor = factor; + self + } + + /// Apply a maximum delay. No retry delay will be longer than this `Duration`. + pub const fn max_delay(mut self, duration: Duration) -> ExponentialBackoff { + self.max_delay = Some(duration); + self + } + + /// Resets the exponential back-off strategy to its initial state. + pub fn reset(&mut self) { + self.current = self.base; + } +} + +impl Iterator for ExponentialBackoff { + type Item = Duration; + + fn next(&mut self) -> Option { + let duration = if let Some(duration) = self.current.checked_mul(self.factor) { + Duration::from_millis(duration) + } else { + Duration::from_millis(u64::MAX) + }; + + // check if we reached max delay + if let Some(ref max_delay) = self.max_delay { + if duration > *max_delay { + return Some(*max_delay); + } + } + + if let Some(next) = self.current.checked_mul(self.base) { + self.current = next; + } else { + self.current = u64::MAX; + } + + Some(duration) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_exponential_backoff_initial_delay() { + let mut backoff = ExponentialBackoff::from_millis(2).factor(250); + let first_delay = backoff.next().expect("backoff should produce a delay"); + + assert_eq!(first_delay, Duration::from_millis(500)); + assert!( + first_delay >= Duration::from_millis(100), + "delay should be reasonable for network errors" + ); + } + + #[test] + fn test_exponential_backoff_doubles_each_retry() { + let mut backoff = ExponentialBackoff::from_millis(2).factor(250); + + let delays: Vec = (0..7).map(|_| backoff.next().unwrap()).collect(); + + for i in 1..delays.len() { + let ratio = delays[i].as_millis() as f64 / delays[i - 1].as_millis() as f64; + assert!( + (1.9..=2.1).contains(&ratio), + "delay[{}] ({:?}) should be ~2x delay[{}] ({:?}), got ratio {}", + i, + delays[i], + i - 1, + delays[i - 1], + ratio + ); + } + } + + #[test] + fn test_exponential_backoff_max_delay_cap() { + let mut backoff = ExponentialBackoff::from_millis(2) + .factor(250) + .max_delay(Duration::from_secs(60)); + + for _ in 0..20 { + backoff.next(); + } + + for i in 0..5 { + let delay = backoff.next().expect("backoff should always produce delays"); + assert_eq!( + delay, + Duration::from_secs(60), + "delay {} after cap should be 60s, got {:?}", + i, + delay + ); + } + } + + #[test] + fn test_exponential_backoff_reset() { + let mut backoff = ExponentialBackoff::from_millis(2).factor(250); + + for _ in 0..5 { + backoff.next(); + } + let high_delay = backoff.next().unwrap(); + assert!( + high_delay > Duration::from_secs(1), + "should have advanced beyond initial delay" + ); + + backoff.reset(); + let reset_delay = backoff.next().unwrap(); + assert_eq!( + reset_delay, + Duration::from_millis(500), + "reset should return to initial 500ms delay" + ); + } + + #[test] + fn test_exponential_backoff_progression_prevents_tight_loops() { + let mut backoff = ExponentialBackoff::from_millis(2).factor(250); + + let mut total_wait_time = Duration::ZERO; + for _ in 0..10 { + total_wait_time += backoff.next().unwrap(); + } + + assert!( + total_wait_time >= Duration::from_secs(30), + "10 retries should accumulate significant wait time, got {:?}", + total_wait_time + ); + } + + #[test] + fn test_exponential_backoff_instances_are_independent() { + let mut backoff1 = ExponentialBackoff::from_millis(2).factor(250); + let mut backoff2 = ExponentialBackoff::from_millis(2).factor(250); + + backoff1.next(); + backoff1.next(); + let delay1 = backoff1.next().unwrap(); + + let delay2 = backoff2.next().unwrap(); + + assert_eq!( + delay2, + Duration::from_millis(500), + "new backoff should start at 500ms" + ); + assert!(delay1 > delay2, "advanced backoff should have higher delay"); + } + + #[test] + fn test_exponential_backoff_with_custom_base() { + let mut backoff = ExponentialBackoff::from_millis(10).factor(100); + + assert_eq!(backoff.next(), Some(Duration::from_millis(1000))); + assert_eq!(backoff.next(), Some(Duration::from_millis(10000))); + assert_eq!(backoff.next(), Some(Duration::from_millis(100000))); + } + + #[test] + fn test_exponential_backoff_without_max_delay() { + let mut backoff = ExponentialBackoff::from_millis(2).factor(1000); + + assert_eq!(backoff.next(), Some(Duration::from_millis(2000))); + assert_eq!(backoff.next(), Some(Duration::from_millis(4000))); + assert_eq!(backoff.next(), Some(Duration::from_millis(8000))); + } +} diff --git a/lib/vector-common/src/lib.rs b/lib/vector-common/src/lib.rs index 9264895f8..8cdd29112 100644 --- a/lib/vector-common/src/lib.rs +++ b/lib/vector-common/src/lib.rs @@ -94,4 +94,6 @@ pub mod net; pub mod multiline; +pub mod backoff; + pub mod compression; diff --git a/lib/vector-core/src/gcp.rs b/lib/vector-core/src/gcp.rs index a3d50b9f1..8810301d1 100644 --- a/lib/vector-core/src/gcp.rs +++ b/lib/vector-core/src/gcp.rs @@ -6,6 +6,7 @@ use std::{ use base64::prelude::{Engine as _, BASE64_URL_SAFE}; pub use goauth::scopes::Scope; +use vector_common::backoff::ExponentialBackoff; use goauth::{ auth::{JwtClaims, Token, TokenErr}, credentials::Credentials, @@ -201,12 +202,16 @@ impl GcpAuthenticator { let expires_in = inner.token.read().unwrap().expires_in() as u64; let mut deadline = Duration::from_secs(expires_in - METADATA_TOKEN_EXPIRY_MARGIN_SECS); + let mut error_backoff = ExponentialBackoff::from_millis(2) + .factor(1000) + .max_delay(Duration::from_secs(300)); loop { tokio::time::sleep(deadline).await; debug!("Renewing GCP authentication token."); match inner.regenerate_token(app_info).await { Ok(()) => { sender.send_replace(()); + error_backoff.reset(); let expires_in = inner.token.read().unwrap().expires_in() as u64; // Rather than an expected fresh token, the Metadata Server may return // the same (cached) token during the last 300 seconds of its lifetime. @@ -220,11 +225,14 @@ impl GcpAuthenticator { deadline = Duration::from_secs(new_deadline); } Err(error) => { + deadline = error_backoff + .next() + .unwrap_or(Duration::from_secs(300)); error!( message = "Failed to update GCP authentication token.", + retry_in_secs = deadline.as_secs(), %error ); - deadline = Duration::from_secs(METADATA_TOKEN_ERROR_RETRY_SECS); } } } diff --git a/src/common/backoff.rs b/src/common/backoff.rs index 39268e6be..8ea8a1321 100644 --- a/src/common/backoff.rs +++ b/src/common/backoff.rs @@ -1,211 +1 @@ -use std::time::Duration; - -// `tokio-retry` crate -// MIT License -// Copyright (c) 2017 Sam Rijs -// -/// A retry strategy driven by exponential back-off. -/// -/// The power corresponds to the number of past attempts. -#[derive(Debug, Clone)] -pub(crate) struct ExponentialBackoff { - current: u64, - base: u64, - factor: u64, - max_delay: Option, -} - -impl ExponentialBackoff { - /// Constructs a new exponential back-off strategy, - /// given a base duration in milliseconds. - /// - /// The resulting duration is calculated by taking the base to the `n`-th power, - /// where `n` denotes the number of past attempts. - pub(crate) const fn from_millis(base: u64) -> ExponentialBackoff { - ExponentialBackoff { - current: base, - base, - factor: 1u64, - max_delay: None, - } - } - - /// A multiplicative factor that will be applied to the retry delay. - /// - /// For example, using a factor of `1000` will make each delay in units of seconds. - /// - /// Default factor is `1`. - pub(crate) const fn factor(mut self, factor: u64) -> ExponentialBackoff { - self.factor = factor; - self - } - - /// Apply a maximum delay. No retry delay will be longer than this `Duration`. - pub(crate) const fn max_delay(mut self, duration: Duration) -> ExponentialBackoff { - self.max_delay = Some(duration); - self - } - - /// Resents the exponential back-off strategy to its initial state. - pub(crate) const fn reset(&mut self) { - self.current = self.base; - } -} - -impl Iterator for ExponentialBackoff { - type Item = Duration; - - fn next(&mut self) -> Option { - let duration = if let Some(duration) = self.current.checked_mul(self.factor) { - Duration::from_millis(duration) - } else { - Duration::from_millis(u64::MAX) - }; - - // check if we reached max delay - if let Some(ref max_delay) = self.max_delay { - if duration > *max_delay { - return Some(*max_delay); - } - } - - if let Some(next) = self.current.checked_mul(self.base) { - self.current = next; - } else { - self.current = u64::MAX; - } - - Some(duration) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_exponential_backoff_initial_delay() { - // Validates the initial delay is calculated correctly - let mut backoff = ExponentialBackoff::from_millis(2).factor(250); - let first_delay = backoff.next().expect("backoff should produce a delay"); - - assert_eq!(first_delay, Duration::from_millis(500)); - assert!(first_delay >= Duration::from_millis(100), "delay should be reasonable for network errors"); - } - - #[test] - fn test_exponential_backoff_doubles_each_retry() { - // Validates exponential growth to prevent sustained API hammering - let mut backoff = ExponentialBackoff::from_millis(2).factor(250); - - let delays: Vec = (0..7).map(|_| backoff.next().unwrap()).collect(); - - // Each delay should be roughly double the previous (within rounding) - for i in 1..delays.len() { - let ratio = delays[i].as_millis() as f64 / delays[i-1].as_millis() as f64; - assert!( - (1.9..=2.1).contains(&ratio), - "delay[{}] ({:?}) should be ~2x delay[{}] ({:?}), got ratio {}", - i, delays[i], i-1, delays[i-1], ratio - ); - } - } - - #[test] - fn test_exponential_backoff_max_delay_cap() { - // Validates max delay prevents indefinite waiting - let mut backoff = ExponentialBackoff::from_millis(2) - .factor(250) - .max_delay(Duration::from_secs(60)); - - // Exhaust the exponential growth - for _ in 0..20 { - backoff.next(); - } - - // All subsequent delays should be exactly 60 seconds - for i in 0..5 { - let delay = backoff.next().expect("backoff should always produce delays"); - assert_eq!( - delay, - Duration::from_secs(60), - "delay {} after cap should be 60s, got {:?}", - i, delay - ); - } - } - - #[test] - fn test_exponential_backoff_reset() { - // Validates reset allows fast recovery after transient errors - let mut backoff = ExponentialBackoff::from_millis(2).factor(250); - - // Advance to a high delay - for _ in 0..5 { - backoff.next(); - } - let high_delay = backoff.next().unwrap(); - assert!(high_delay > Duration::from_secs(1), "should have advanced beyond initial delay"); - - // Reset and verify we're back to 500ms - backoff.reset(); - let reset_delay = backoff.next().unwrap(); - assert_eq!(reset_delay, Duration::from_millis(500), "reset should return to initial 500ms delay"); - } - - #[test] - fn test_exponential_backoff_progression_prevents_tight_loops() { - // Validates the backoff prevents tight retry loops that could DOS the service - let mut backoff = ExponentialBackoff::from_millis(2).factor(250); - - let mut total_wait_time = Duration::ZERO; - for _ in 0..10 { - total_wait_time += backoff.next().unwrap(); - } - - // After 10 retries, should have waited at least 30 seconds total - assert!( - total_wait_time >= Duration::from_secs(30), - "10 retries should accumulate significant wait time, got {:?}", - total_wait_time - ); - } - - #[test] - fn test_exponential_backoff_instances_are_independent() { - // Validates that different instances can have independent states - let mut backoff1 = ExponentialBackoff::from_millis(2).factor(250); - let mut backoff2 = ExponentialBackoff::from_millis(2).factor(250); - - // Advance backoff1 - backoff1.next(); - backoff1.next(); - let delay1 = backoff1.next().unwrap(); - - // backoff2 should still be at initial state - let delay2 = backoff2.next().unwrap(); - - assert_eq!(delay2, Duration::from_millis(500), "new backoff should start at 500ms"); - assert!(delay1 > delay2, "advanced backoff should have higher delay"); - } - - #[test] - fn test_exponential_backoff_with_custom_base() { - // Test backoff with custom base and factor - let mut backoff = ExponentialBackoff::from_millis(10).factor(100); - - assert_eq!(backoff.next(), Some(Duration::from_millis(1000))); // 10 * 100, then current = 10 * 10 = 100 - assert_eq!(backoff.next(), Some(Duration::from_millis(10000))); // 100 * 100, then current = 100 * 10 = 1000 - assert_eq!(backoff.next(), Some(Duration::from_millis(100000))); // 1000 * 100 - } - - #[test] - fn test_exponential_backoff_without_max_delay() { - // Test backoff without max delay cap - let mut backoff = ExponentialBackoff::from_millis(2).factor(1000); - - assert_eq!(backoff.next(), Some(Duration::from_millis(2000))); - assert_eq!(backoff.next(), Some(Duration::from_millis(4000))); - assert_eq!(backoff.next(), Some(Duration::from_millis(8000))); - } -} +pub(crate) use vector_common::backoff::ExponentialBackoff;