// Copyright 2024 RustFS Team // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. //! Operation execution policy for external KMS backend calls. //! //! Vault-backed operations leave the process boundary, so every call needs a //! per-attempt timeout, a total operation deadline, and classification-driven //! bounded retries. The Vault backends and the credential provider wire every //! outbound `vaultrs` call through [`execute`]. //! //! Retry safety is driven by two orthogonal classifications: //! - [`OpClass`] states whether replaying the operation is safe at all. //! - [`ErrorClass`] states whether the observed failure is worth replaying. //! //! Mutating operations without an idempotency key or CAS precondition are never //! retried automatically: a response lost after the server applied the write //! would otherwise be replayed into duplicate side effects (extra key versions, //! repeated deletes). //! //! Every execution also records operation metrics (attempt failures by retry //! class, terminal outcome, attempts used, wall-clock duration) through the //! process-global `metrics` recorder. Metric labels carry only static enum //! values — operation names, classes, outcomes — never key identifiers, key //! material, ciphertext, or tokens. use std::collections::HashMap; use std::future::Future; use std::sync::{Arc, LazyLock, Mutex, Weak}; use std::time::Duration; use rand::{RngExt, SeedableRng, rngs::StdRng}; use sha2::{Digest, Sha256}; use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore, TryAcquireError}; use tokio::time::Instant; use tokio_util::sync::CancellationToken; use crate::config::KmsConfig; use crate::error::{KmsError, Result}; /// Default backoff cap before the first retry; doubles per retry. const DEFAULT_BASE_BACKOFF: Duration = Duration::from_millis(100); /// Default upper bound for a single backoff sleep. const DEFAULT_MAX_BACKOFF: Duration = Duration::from_secs(2); const DEFAULT_MAX_CONCURRENT_OPERATIONS: usize = 64; const RESERVED_CREDENTIAL_OPERATIONS: usize = 1; const DEFAULT_MAX_QUEUED_OPERATIONS: usize = 64; const DEFAULT_CIRCUIT_FAILURE_THRESHOLD: u32 = 5; const DEFAULT_CIRCUIT_OPEN_DURATION: Duration = Duration::from_secs(30); const CIRCUIT_OPEN_MESSAGE: &str = "KMS backend circuit is open; waiting for the half-open recovery probe"; const BACKPRESSURE_MESSAGE: &str = "KMS backend capacity and wait queue are full"; /// Replay safety of a backend operation. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum OpClass { /// No persistent side effects (encrypt/decrypt/generate/describe/list/health); /// safe to retry on any retryable failure. ReadIdempotent, /// External write without an idempotency key or CAS precondition /// (create/rotate/delete/configure); executed at most once. MutatingNonIdempotent, /// Authentication exchange (login, token renewal); safe to resend. Auth, } impl OpClass { fn retryable(self) -> bool { !matches!(self, OpClass::MutatingNonIdempotent) } } /// Retry-relevant classification of a failed attempt. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum ErrorClass { /// Connection-level failure (connect/send/response read); the server may or /// may not have observed the request. RetryableConn, /// Retryable HTTP status: 429 throttling or a recoverable 5xx. RetryableStatus, /// Deterministic failure (auth, validation, not-found, malformed data); /// retrying cannot help and may mask the real problem. Fatal, } /// Classify a `vaultrs` client error for retry purposes. /// /// Status codes are inspected both on `ClientError::APIError` (JSON error body) /// and on a wrapped rustify `ServerResponseError` (non-JSON body, e.g. an HTML /// page from an intermediate load balancer). Everything that is not throttling, /// a recoverable 5xx, or a connection-level failure is fatal; in particular /// 400/401/403/404 must never be retried. pub(crate) fn classify_vaultrs(error: &vaultrs::error::ClientError) -> ErrorClass { use rustify::errors::ClientError as RestError; use vaultrs::error::ClientError; match error { ClientError::APIError { code, .. } => classify_status(*code), ClientError::RestClientError { source } => match source { RestError::ServerResponseError { code, .. } => classify_status(*code), RestError::RequestError { .. } | RestError::ResponseError { .. } => ErrorClass::RetryableConn, _ => ErrorClass::Fatal, }, _ => ErrorClass::Fatal, } } /// Retry classification of an HTTP status, shared by every backend that talks /// to an external KMS over HTTP. pub(crate) fn classify_status(code: u16) -> ErrorClass { match code { 429 | 500 | 502 | 503 | 504 => ErrorClass::RetryableStatus, _ => ErrorClass::Fatal, } } /// Failure of a single attempt, carrying its retry classification. #[derive(Debug)] pub(crate) struct AttemptError { pub(crate) class: ErrorClass, pub(crate) error: KmsError, } impl AttemptError { /// A failure that must never be retried, regardless of operation class. pub(crate) fn fatal(error: KmsError) -> Self { Self { class: ErrorClass::Fatal, error, } } /// Classify a `vaultrs` failure and map it onto a domain error. /// /// Classification reads the raw error before `map` consumes it, so call /// sites keep their site-specific error mapping (404 to key-not-found and /// so on) without losing the status code the retry decision needs. pub(crate) fn from_vaultrs( error: vaultrs::error::ClientError, map: impl FnOnce(vaultrs::error::ClientError) -> KmsError, ) -> Self { let class = classify_vaultrs(&error); Self { class, error: map(error), } } } /// Budgets applied by [`execute`]. #[derive(Debug, Clone)] pub(crate) struct RetryPolicy { /// Upper bound for one backend attempt. pub(crate) attempt_timeout: Duration, /// Upper bound for the whole operation, including retries and backoff. pub(crate) op_deadline: Duration, /// Maximum attempts for retryable operation classes; non-idempotent /// mutations always run exactly once regardless of this value. pub(crate) max_attempts: u32, /// Backoff cap before the first retry; doubles per retry. pub(crate) base_backoff: Duration, /// Upper bound for a single backoff sleep. pub(crate) max_backoff: Duration, runtime: Arc, } impl RetryPolicy { /// Derive the policy from the KMS configuration. /// /// `timeout` and `retry_attempts` are taken with the config-level clamps /// applied. The operation deadline covers the worst-case budget of all /// attempts plus backoff, so it bounds runaway loops without cutting any /// attempt short; an independently configurable deadline is left to the /// admin-API follow-up. pub(crate) fn for_backend( config: &KmsConfig, backend: &'static str, endpoint: &str, namespace: Option<&str>, scope: &'static str, ) -> Self { Self::for_capacity(config, backend, endpoint, namespace, scope, CapacityClass::Operations) } pub(crate) fn for_credentials( config: &KmsConfig, backend: &'static str, endpoint: &str, namespace: Option<&str>, scope: &'static str, ) -> Self { Self::for_capacity(config, backend, endpoint, namespace, scope, CapacityClass::Credentials) } fn for_capacity( config: &KmsConfig, backend: &'static str, endpoint: &str, namespace: Option<&str>, scope: &'static str, capacity: CapacityClass, ) -> Self { let endpoint = url::Url::parse(endpoint).map_or_else(|_| endpoint.to_owned(), |url| url.to_string()); let mut identity = Sha256::new(); for part in [backend, endpoint.as_str(), namespace.unwrap_or_default()] { identity.update(part.len().to_be_bytes()); identity.update(part.as_bytes()); } let mut policy = Self { attempt_timeout: config.effective_timeout(), op_deadline: Duration::ZERO, max_attempts: config.effective_retry_attempts(), base_backoff: DEFAULT_BASE_BACKOFF, max_backoff: DEFAULT_MAX_BACKOFF, runtime: BackendRuntime::shared(identity.finalize().into(), backend, scope, capacity), }; policy.op_deadline = policy.worst_case_budget(); policy } /// Total worst-case duration: every attempt hits `attempt_timeout` and /// every backoff sleeps its full cap. fn worst_case_budget(&self) -> Duration { let attempts = self.max_attempts.max(1); let mut budget = self.attempt_timeout.saturating_mul(attempts); for completed in 1..attempts { budget = budget.saturating_add(backoff_cap(self, completed)); } budget } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum CapacityClass { Operations, Credentials, } static ACTIVE_REGISTRY: LazyLock>>> = LazyLock::new(|| Mutex::new(HashMap::new())); #[derive(Debug)] struct BackendCapacity { total: Arc, operations: Arc, } #[derive(Debug)] struct BackendRuntime { /// Static backend name, carried so per-operation metrics can be attributed /// to the backend that served them rather than aggregated across all of them. backend: &'static str, capacity: Arc, capacity_class: CapacityClass, queued: Arc, breaker: Mutex, in_flight: metrics::Gauge, circuit_open: metrics::Gauge, } impl BackendRuntime { fn shared(identity: [u8; 32], backend: &'static str, scope: &'static str, capacity: CapacityClass) -> Arc { let mut registry = ACTIVE_REGISTRY.lock().unwrap_or_else(|poisoned| poisoned.into_inner()); registry.retain(|_, active| active.strong_count() > 0); let active = registry.get(&identity).and_then(Weak::upgrade).unwrap_or_else(|| { Arc::new(BackendCapacity { total: Arc::new(Semaphore::new(DEFAULT_MAX_CONCURRENT_OPERATIONS)), operations: Arc::new(Semaphore::new(DEFAULT_MAX_CONCURRENT_OPERATIONS - RESERVED_CREDENTIAL_OPERATIONS)), }) }); registry.insert(identity, Arc::downgrade(&active)); Arc::new(Self::new(backend, scope, active, capacity)) } fn new(backend: &'static str, scope: &'static str, active: Arc, capacity: CapacityClass) -> Self { let in_flight = metrics::gauge!(METRIC_IN_FLIGHT, "backend" => backend, "scope" => scope); let circuit_open = metrics::gauge!(METRIC_CIRCUIT_OPEN, "backend" => backend, "scope" => scope); in_flight.increment(0.0); circuit_open.increment(0.0); Self { backend, capacity: active, capacity_class: capacity, queued: Arc::new(Semaphore::new(DEFAULT_MAX_QUEUED_OPERATIONS)), breaker: Mutex::new(CircuitState::default()), in_flight, circuit_open, } } fn try_acquire_active(&self) -> std::result::Result { let operations = matches!(self.capacity_class, CapacityClass::Operations) .then(|| Arc::clone(&self.capacity.operations).try_acquire_owned()) .transpose()?; match Arc::clone(&self.capacity.total).try_acquire_owned() { Ok(total) => Ok(ActivePermits { _total: total, _operations: operations, }), Err(error) => { drop(operations); Err(error) } } } async fn acquire_active(&self) -> std::result::Result { let operations = match self.capacity_class { CapacityClass::Operations => Some(Arc::clone(&self.capacity.operations).acquire_owned().await?), CapacityClass::Credentials => None, }; let total = Arc::clone(&self.capacity.total).acquire_owned().await?; Ok(ActivePermits { _total: total, _operations: operations, }) } fn breaker_state(&self) -> Result> { self.breaker .lock() .map_err(|_| KmsError::backend_error("KMS backend circuit state is unavailable")) } fn rejects(&self, now: Instant) -> Result { let state = self.breaker_state()?; Ok(matches!(state.open_until, Some(open_until) if now < open_until || state.probe_in_flight)) } fn admit(&self, now: Instant) -> Result> { let mut state = self.breaker_state()?; Ok(match state.open_until { None => Some(CircuitAdmission { generation: state.generation, probe: false, }), Some(open_until) if now >= open_until && !state.probe_in_flight => { state.probe_in_flight = true; Some(CircuitAdmission { generation: state.generation, probe: true, }) } Some(_) => None, }) } fn transition(&self, state: &mut CircuitState, open: bool) { let was_open = state.open_until.is_some(); state.generation = state.generation.saturating_add(1); state.consecutive_failures = 0; state.open_until = open.then(|| Instant::now() + DEFAULT_CIRCUIT_OPEN_DURATION); state.probe_in_flight = false; if was_open != open { if open { self.circuit_open.increment(1.0); } else { self.circuit_open.decrement(1.0); } } } fn complete(&self, admission: CircuitAdmission, failure: Option) -> bool { let Ok(mut state) = self.breaker_state() else { return true; }; if state.generation != admission.generation { return state.open_until.is_some(); } match failure { None | Some(ErrorClass::Fatal) => { state.consecutive_failures = 0; if admission.probe { self.transition(&mut state, false); } } Some(ErrorClass::RetryableConn | ErrorClass::RetryableStatus) => { state.consecutive_failures = state.consecutive_failures.saturating_add(1); if admission.probe || state.consecutive_failures >= DEFAULT_CIRCUIT_FAILURE_THRESHOLD { self.transition(&mut state, true); } } } state.open_until.is_some() } fn abandon(&self, admission: CircuitAdmission) { if let Ok(mut state) = self.breaker_state() && state.generation == admission.generation && admission.probe { self.transition(&mut state, true); } } } impl Drop for BackendRuntime { fn drop(&mut self) { let state = self.breaker.get_mut().unwrap_or_else(|poisoned| poisoned.into_inner()); if state.open_until.is_some() { self.circuit_open.decrement(1.0); } } } #[derive(Debug, Default)] struct CircuitState { generation: u64, consecutive_failures: u32, open_until: Option, probe_in_flight: bool, } #[derive(Debug, Clone, Copy)] struct CircuitAdmission { generation: u64, probe: bool, } struct RuntimePermit { runtime: Arc, admission: Option, _active: ActivePermits, } struct ActivePermits { _total: OwnedSemaphorePermit, _operations: Option, } impl RuntimePermit { #[cfg(test)] fn new(runtime: Arc, admission: CircuitAdmission, active: OwnedSemaphorePermit) -> Self { Self::with_active( runtime, admission, ActivePermits { _total: active, _operations: None, }, ) } fn with_active(runtime: Arc, admission: CircuitAdmission, active: ActivePermits) -> Self { runtime.in_flight.increment(1.0); Self { runtime, admission: Some(admission), _active: active, } } fn complete(&mut self, failure: Option) -> bool { self.admission .take() .is_some_and(|admission| self.runtime.complete(admission, failure)) } } impl Drop for RuntimePermit { fn drop(&mut self) { if let Some(admission) = self.admission.take() { self.runtime.abandon(admission); } self.runtime.in_flight.decrement(1.0); } } /// Exponential backoff cap after `completed_attempts` failed attempts. fn backoff_cap(policy: &RetryPolicy, completed_attempts: u32) -> Duration { let doublings = completed_attempts.saturating_sub(1).min(31); policy.base_backoff.saturating_mul(1u32 << doublings).min(policy.max_backoff) } /// Equal jitter: sleep within `[cap / 2, cap]`, keeping at least half the cap /// so backoff still backs off while decorrelating retry bursts. fn equal_jitter(rng: &mut impl RngExt, cap: Duration) -> Duration { let half = cap / 2; let spread = u64::try_from(half.as_nanos()).unwrap_or(u64::MAX); half + Duration::from_nanos(rng.random_range(0..=spread)) } // --------------------------------------------------------------------------- // Metrics // // Every execution is recorded here, at the single choke point all backend // calls flow through, so instrumenting a new call site costs nothing beyond // naming its operation. Label values are exclusively static enum strings // (operation names, classes, outcomes) — key identifiers, key material, // ciphertext, and tokens must never reach a metric label. // --------------------------------------------------------------------------- /// Counter: operations executed, by `backend`, `operation`, `op_class`, and `outcome`. const METRIC_OPERATIONS_TOTAL: &str = "rustfs_kms_backend_operations_total"; /// Counter: failed attempts, by `backend`, `operation` and `error_class` /// (including `attempt_timeout` for attempts cut off by the per-attempt timeout). const METRIC_ATTEMPT_FAILURES_TOTAL: &str = "rustfs_kms_backend_attempt_failures_total"; /// Histogram: wall-clock duration of a whole operation (attempts plus /// backoff), in seconds, by `backend`, `operation` and `outcome`. const METRIC_OPERATION_DURATION_SECONDS: &str = "rustfs_kms_backend_operation_duration_seconds"; /// Histogram: attempts one operation used before completing, by `backend`, /// `operation` and `outcome`. const METRIC_OPERATION_ATTEMPTS: &str = "rustfs_kms_backend_operation_attempts"; /// Gauge: backend attempts currently in flight. const METRIC_IN_FLIGHT: &str = "rustfs_kms_backend_in_flight"; /// Gauge: number of open or half-open backend circuits. const METRIC_CIRCUIT_OPEN: &str = "rustfs_kms_backend_circuit_open"; impl OpClass { fn as_label(self) -> &'static str { match self { OpClass::ReadIdempotent => "read_idempotent", OpClass::MutatingNonIdempotent => "mutating_non_idempotent", OpClass::Auth => "auth", } } } impl ErrorClass { fn as_label(self) -> &'static str { match self { ErrorClass::RetryableConn => "retryable_conn", ErrorClass::RetryableStatus => "retryable_status", ErrorClass::Fatal => "fatal", } } } /// How one policy execution terminated, for the `outcome` metric label. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum Outcome { Success, /// A fatal-classified failure ended the operation on its first observation. Fatal, /// The attempt budget ran out; the last failure was retryable (including a /// timed-out final attempt). BudgetExhausted, /// The operation deadline ran out before another attempt could complete. DeadlineExceeded, /// The operation deadline ran out while waiting for backend capacity. BackpressureTimeout, /// Backend capacity and the bounded wait queue were both full. BackpressureRejected, /// A retryable failure opened the circuit, or an open circuit rejected the operation. CircuitOpen, Cancelled, } impl Outcome { fn as_label(self) -> &'static str { match self { Outcome::Success => "success", Outcome::Fatal => "fatal", Outcome::BudgetExhausted => "budget_exhausted", Outcome::DeadlineExceeded => "deadline_exceeded", Outcome::BackpressureTimeout => "backpressure_timeout", Outcome::BackpressureRejected => "backpressure_rejected", Outcome::CircuitOpen => "circuit_open", Outcome::Cancelled => "cancelled", } } } /// Register metric descriptions once per process. fn describe_metrics() { static DESCRIBE: std::sync::Once = std::sync::Once::new(); DESCRIBE.call_once(|| { metrics::describe_counter!( METRIC_OPERATIONS_TOTAL, "Total KMS backend operations executed under the operation policy, by backend, operation, operation class, and outcome" ); metrics::describe_counter!( METRIC_ATTEMPT_FAILURES_TOTAL, "Total failed KMS backend attempts, by backend, operation and retry classification" ); metrics::describe_histogram!( METRIC_OPERATION_DURATION_SECONDS, "Wall-clock duration of KMS backend operations including retries and backoff, in seconds, by backend and operation" ); metrics::describe_histogram!( METRIC_OPERATION_ATTEMPTS, "Number of attempts a KMS backend operation used before completing, by backend and operation" ); metrics::describe_gauge!(METRIC_IN_FLIGHT, "KMS backend attempts currently in flight, by backend and policy scope"); metrics::describe_gauge!( METRIC_CIRCUIT_OPEN, "Number of open or half-open KMS backend circuits, by backend and policy scope" ); }); } /// Record one failed attempt with its retry classification. fn record_attempt_failure(backend: &'static str, operation: &'static str, error_class: &'static str) { metrics::counter!( METRIC_ATTEMPT_FAILURES_TOTAL, "backend" => backend, "operation" => operation, "error_class" => error_class ) .increment(1); } /// Record the terminal outcome of one policy execution. fn record_operation( backend: &'static str, operation: &'static str, class: OpClass, outcome: Outcome, attempts: u32, elapsed: Duration, ) { metrics::counter!( METRIC_OPERATIONS_TOTAL, "backend" => backend, "operation" => operation, "op_class" => class.as_label(), "outcome" => outcome.as_label() ) .increment(1); metrics::histogram!( METRIC_OPERATION_DURATION_SECONDS, "backend" => backend, "operation" => operation, "outcome" => outcome.as_label() ) .record(elapsed.as_secs_f64()); metrics::histogram!( METRIC_OPERATION_ATTEMPTS, "backend" => backend, "operation" => operation, "outcome" => outcome.as_label() ) .record(f64::from(attempts)); } async fn acquire_attempt( runtime: Arc, operation: &'static str, cancel: &CancellationToken, deadline: Instant, ) -> std::result::Result { if cancel.is_cancelled() { return Err(( Outcome::Cancelled, KmsError::operation_cancelled(format!("{operation} cancelled before backend admission")), )); } match runtime.rejects(Instant::now()) { Ok(false) => {} Ok(true) => return Err((Outcome::CircuitOpen, KmsError::backend_error(CIRCUIT_OPEN_MESSAGE))), Err(error) => return Err((Outcome::CircuitOpen, error)), } let active = match runtime.try_acquire_active() { Ok(active) => active, Err(TryAcquireError::Closed) => { return Err(( Outcome::BackpressureRejected, KmsError::backend_error("KMS backend admission is unavailable"), )); } Err(TryAcquireError::NoPermits) => { let queued = Arc::clone(&runtime.queued) .try_acquire_owned() .map_err(|_| (Outcome::BackpressureRejected, KmsError::backend_error(BACKPRESSURE_MESSAGE)))?; let acquire = runtime.acquire_active(); let active = tokio::select! { biased; _ = cancel.cancelled() => { return Err((Outcome::Cancelled, KmsError::operation_cancelled(format!( "{operation} cancelled while waiting for backend capacity" )))); } result = tokio::time::timeout_at(deadline, acquire) => match result { Ok(Ok(active)) => active, Ok(Err(_)) => { return Err((Outcome::BackpressureRejected, KmsError::backend_error( "KMS backend admission is unavailable" ))); } Err(_) => { return Err((Outcome::BackpressureTimeout, KmsError::operation_timed_out(format!( "{operation} exceeded its operation deadline while waiting for backend capacity" )))); } } }; drop(queued); active } }; if cancel.is_cancelled() { return Err(( Outcome::Cancelled, KmsError::operation_cancelled(format!("{operation} cancelled before backend attempt")), )); } if Instant::now() >= deadline { return Err(( Outcome::BackpressureTimeout, KmsError::operation_timed_out(format!("{operation} exceeded its operation deadline before backend attempt")), )); } let admission = match runtime.admit(Instant::now()) { Ok(Some(admission)) => admission, Ok(None) => return Err((Outcome::CircuitOpen, KmsError::backend_error(CIRCUIT_OPEN_MESSAGE))), Err(error) => return Err((Outcome::CircuitOpen, error)), }; Ok(RuntimePermit::with_active(runtime, admission, active)) } /// Run `attempt` under the policy. /// /// Each attempt is bounded by `attempt_timeout` (further capped by whatever is /// left of `op_deadline`), and retryable failures are replayed with exponential /// backoff and jitter when the operation class allows it. Cancellation aborts /// both in-flight attempts and backoff sleeps. /// /// A timed-out attempt counts as a connection-class failure: the server may /// have processed the request, which is exactly why non-idempotent mutations /// are never replayed. pub(crate) async fn execute( operation: &'static str, class: OpClass, policy: &RetryPolicy, cancel: &CancellationToken, attempt: F, ) -> Result where F: FnMut() -> Fut, Fut: Future>, { // Seed an owned RNG up front: the thread-local RNG is not Send and must // not be held across await points. let mut rng = StdRng::from_rng(&mut rand::rng()); execute_with_jitter(operation, class, policy, cancel, move |cap| equal_jitter(&mut rng, cap), attempt).await } /// [`execute`] with an injectable jitter source so tests can pin deterministic /// backoff durations instead of asserting around random sleeps. pub(crate) async fn execute_with_jitter( operation: &'static str, class: OpClass, policy: &RetryPolicy, cancel: &CancellationToken, jitter: J, attempt: F, ) -> Result where F: FnMut() -> Fut, Fut: Future>, J: FnMut(Duration) -> Duration, { describe_metrics(); let started = Instant::now(); let mut attempts_made = 0u32; let (outcome, result) = drive_attempts(operation, class, policy, cancel, jitter, attempt, &mut attempts_made).await; record_operation(policy.runtime.backend, operation, class, outcome, attempts_made, started.elapsed()); result } /// The attempt loop behind [`execute_with_jitter`], returning the terminal /// outcome alongside the result so the caller can record it exactly once. async fn drive_attempts( operation: &'static str, class: OpClass, policy: &RetryPolicy, cancel: &CancellationToken, mut jitter: J, mut attempt: F, attempts_made: &mut u32, ) -> (Outcome, Result) where F: FnMut() -> Fut, Fut: Future>, J: FnMut(Duration) -> Duration, { let deadline = Instant::now() + policy.op_deadline; let max_attempts = if class.retryable() { policy.max_attempts.max(1) } else { 1 }; let mut attempt_no = 0u32; loop { if cancel.is_cancelled() { return ( Outcome::Cancelled, Err(KmsError::operation_cancelled(format!( "{operation} cancelled before attempt {}", attempt_no + 1 ))), ); } let remaining = deadline.saturating_duration_since(Instant::now()); if remaining.is_zero() { return ( Outcome::DeadlineExceeded, Err(KmsError::operation_timed_out(format!( "{operation} exceeded operation deadline of {:?}", policy.op_deadline ))), ); } let mut runtime_permit = match acquire_attempt(Arc::clone(&policy.runtime), operation, cancel, deadline).await { Ok(permit) => permit, Err((outcome, error)) => return (outcome, Err(error)), }; attempt_no += 1; *attempts_made = attempt_no; let remaining = deadline.saturating_duration_since(Instant::now()); let attempt_budget = policy.attempt_timeout.min(remaining); let outcome = tokio::select! { biased; _ = cancel.cancelled() => { return ( Outcome::Cancelled, Err(KmsError::operation_cancelled(format!("{operation} cancelled during attempt {attempt_no}"))), ); } outcome = tokio::time::timeout(attempt_budget, attempt()) => outcome, }; let failure = match outcome { Ok(Ok(value)) => { runtime_permit.complete(None); return (Outcome::Success, Ok(value)); } Ok(Err(failure)) => { record_attempt_failure(policy.runtime.backend, operation, failure.class.as_label()); failure } Err(_) => { record_attempt_failure(policy.runtime.backend, operation, "attempt_timeout"); AttemptError { class: ErrorClass::RetryableConn, error: KmsError::operation_timed_out(format!( "{operation} attempt {attempt_no} timed out after {attempt_budget:?}" )), } } }; let circuit_open = runtime_permit.complete(Some(failure.class)); if failure.class == ErrorClass::Fatal { return (Outcome::Fatal, Err(failure.error)); } drop(runtime_permit); if circuit_open { return (Outcome::CircuitOpen, Err(failure.error)); } if attempt_no >= max_attempts { return (Outcome::BudgetExhausted, Err(failure.error)); } let backoff = jitter(backoff_cap(policy, attempt_no)); if backoff >= deadline.saturating_duration_since(Instant::now()) { // Not enough deadline budget left for another attempt. return (Outcome::DeadlineExceeded, Err(failure.error)); } tracing::warn!( operation, attempt = attempt_no, error_class = ?failure.class, backoff = ?backoff, "KMS backend attempt failed with a retryable error; backing off before retry" ); tokio::select! { biased; _ = cancel.cancelled() => { return ( Outcome::Cancelled, Err(KmsError::operation_cancelled(format!("{operation} cancelled during retry backoff"))), ); } _ = tokio::time::sleep(backoff) => {} } } } #[cfg(test)] fn test_runtime(max_concurrent: usize, max_queued: usize) -> Arc { named_test_runtime("test-backend", max_concurrent, max_queued) } #[cfg(test)] fn named_test_runtime(backend: &'static str, max_concurrent: usize, max_queued: usize) -> Arc { let capacity = Arc::new(BackendCapacity { total: Arc::new(Semaphore::new(max_concurrent)), operations: Arc::new(Semaphore::new(max_concurrent)), }); let mut runtime = BackendRuntime::new(backend, "test-scope", capacity, CapacityClass::Credentials); runtime.queued = Arc::new(Semaphore::new(max_queued)); Arc::new(runtime) } #[cfg(test)] impl RetryPolicy { pub(crate) fn for_test( attempt_timeout: Duration, op_deadline: Duration, max_attempts: u32, base_backoff: Duration, max_backoff: Duration, ) -> Self { Self { attempt_timeout, op_deadline, max_attempts, base_backoff, max_backoff, runtime: test_runtime(DEFAULT_MAX_CONCURRENT_OPERATIONS, DEFAULT_MAX_QUEUED_OPERATIONS), } } pub(crate) fn shares_active_capacity_with(&self, other: &Self) -> bool { Arc::ptr_eq(&self.runtime.capacity, &other.runtime.capacity) } pub(crate) fn uses_credential_reserve(&self) -> bool { matches!(self.runtime.capacity_class, CapacityClass::Credentials) } } #[cfg(test)] mod tests { use super::*; use std::sync::Arc; use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; use tokio::sync::Notify; type AttemptResult = std::result::Result; fn policy_of(attempt_timeout_ms: u64, op_deadline_ms: u64, max_attempts: u32, base_ms: u64, max_ms: u64) -> RetryPolicy { RetryPolicy::for_test( Duration::from_millis(attempt_timeout_ms), Duration::from_millis(op_deadline_ms), max_attempts, Duration::from_millis(base_ms), Duration::from_millis(max_ms), ) } fn policy_with_limit(max_concurrent: usize, max_queued: usize) -> RetryPolicy { let mut policy = policy_of(1_000, 60_000, 1, 10, 10); policy.runtime = test_runtime(max_concurrent, max_queued); policy } fn unique_endpoint(label: &str) -> String { static NEXT_ENDPOINT_ID: AtomicU64 = AtomicU64::new(0); format!("https://{label}-{}.example.invalid", NEXT_ENDPOINT_ID.fetch_add(1, Ordering::Relaxed)) } /// Deterministic jitter: always sleep the full backoff cap. fn full_jitter(cap: Duration) -> Duration { cap } fn retryable_conn_error() -> AttemptError { AttemptError { class: ErrorClass::RetryableConn, error: KmsError::backend_error("connection reset by peer"), } } fn is_circuit_open(error: &KmsError) -> bool { matches!(error, KmsError::BackendError { message } if message.contains("circuit is open")) } fn trip_breaker(runtime: &BackendRuntime) { for attempt in 1..=DEFAULT_CIRCUIT_FAILURE_THRESHOLD { let admission = runtime.admit(Instant::now()).expect("state").expect("closed"); assert_eq!( runtime.complete(admission, Some(ErrorClass::RetryableConn)), attempt == DEFAULT_CIRCUIT_FAILURE_THRESHOLD ); } } #[tokio::test(start_paused = true)] async fn hung_attempt_fails_within_attempt_timeout() { let policy = policy_of(5_000, 60_000, 1, 100, 2_000); let cancel = CancellationToken::new(); let started = Instant::now(); let result: Result<()> = execute_with_jitter("encrypt", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, || { std::future::pending::>() }) .await; assert!(matches!(result, Err(KmsError::OperationTimedOut { .. })), "got {result:?}"); assert_eq!(started.elapsed(), Duration::from_millis(5_000)); } #[tokio::test(start_paused = true)] async fn retryable_status_retries_until_success() { let policy = policy_of(1_000, 60_000, 3, 100, 2_000); let cancel = CancellationToken::new(); let calls = Arc::new(AtomicU32::new(0)); let calls_in_attempt = calls.clone(); let started = Instant::now(); let result = execute_with_jitter("generate_data_key", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, move || { let calls = calls_in_attempt.clone(); async move { if calls.fetch_add(1, Ordering::SeqCst) < 2 { Err(AttemptError { class: ErrorClass::RetryableStatus, error: KmsError::backend_error("throttled (429)"), }) } else { Ok(7u32) } } }) .await; assert_eq!(result.expect("retries within budget must succeed"), 7); assert_eq!(calls.load(Ordering::SeqCst), 3); // Full-cap backoff: 100ms after attempt 1, 200ms after attempt 2. assert_eq!(started.elapsed(), Duration::from_millis(300)); } #[tokio::test(start_paused = true)] async fn fatal_error_is_not_retried() { let policy = policy_of(1_000, 60_000, 5, 100, 2_000); let cancel = CancellationToken::new(); let calls = Arc::new(AtomicU32::new(0)); let calls_in_attempt = calls.clone(); let started = Instant::now(); let result: Result<()> = execute_with_jitter("decrypt", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, move || { let calls = calls_in_attempt.clone(); async move { calls.fetch_add(1, Ordering::SeqCst); Err(AttemptError { class: ErrorClass::Fatal, error: KmsError::access_denied("permission denied (403)"), }) } }) .await; assert!(matches!(result, Err(KmsError::AccessDenied { .. })), "got {result:?}"); assert_eq!(calls.load(Ordering::SeqCst), 1); assert_eq!(started.elapsed(), Duration::ZERO); } #[tokio::test(start_paused = true)] async fn mutating_non_idempotent_runs_exactly_once() { let policy = policy_of(1_000, 60_000, 5, 100, 2_000); let cancel = CancellationToken::new(); let calls = Arc::new(AtomicU32::new(0)); let calls_in_attempt = calls.clone(); let result: Result<()> = execute_with_jitter("rotate_key", OpClass::MutatingNonIdempotent, &policy, &cancel, full_jitter, move || { let calls = calls_in_attempt.clone(); async move { calls.fetch_add(1, Ordering::SeqCst); Err(retryable_conn_error()) } }) .await; // Even a retryable failure must not replay a non-idempotent mutation. assert!(matches!(result, Err(KmsError::BackendError { .. })), "got {result:?}"); assert_eq!(calls.load(Ordering::SeqCst), 1); } #[tokio::test(start_paused = true)] async fn cancel_aborts_backoff_immediately() { // Long backoff (10s cap) so a prompt return can only come from cancellation. let policy = policy_of(1_000, 600_000, 5, 10_000, 10_000); let cancel = CancellationToken::new(); let canceller = cancel.clone(); tokio::spawn(async move { tokio::time::sleep(Duration::from_millis(500)).await; canceller.cancel(); }); let calls = Arc::new(AtomicU32::new(0)); let calls_in_attempt = calls.clone(); let started = Instant::now(); let result: Result<()> = execute_with_jitter("decrypt", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, move || { let calls = calls_in_attempt.clone(); async move { calls.fetch_add(1, Ordering::SeqCst); Err(retryable_conn_error()) } }) .await; assert!(matches!(result, Err(KmsError::OperationCancelled { .. })), "got {result:?}"); assert_eq!(calls.load(Ordering::SeqCst), 1); assert_eq!(started.elapsed(), Duration::from_millis(500)); } #[tokio::test(start_paused = true)] async fn cancelled_token_short_circuits_before_first_attempt() { let policy = policy_of(1_000, 60_000, 5, 100, 2_000); let cancel = CancellationToken::new(); cancel.cancel(); let calls = Arc::new(AtomicU32::new(0)); let calls_in_attempt = calls.clone(); let result: Result<()> = execute_with_jitter("list_keys", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, move || { let calls = calls_in_attempt.clone(); async move { calls.fetch_add(1, Ordering::SeqCst); Ok(()) } }) .await; assert!(matches!(result, Err(KmsError::OperationCancelled { .. })), "got {result:?}"); assert_eq!(calls.load(Ordering::SeqCst), 0); } #[test] fn backend_generation_reuses_capacity_but_resets_queue_and_breaker() { let endpoint = unique_endpoint("fresh-generation"); let config = KmsConfig::default(); let first = RetryPolicy::for_backend(&config, "vault", &endpoint, Some("namespace"), "operations"); let second = RetryPolicy::for_backend(&config, "vault", &endpoint, Some("namespace"), "operations"); assert!(!Arc::ptr_eq(&first.runtime, &second.runtime)); assert!(Arc::ptr_eq(&first.runtime.capacity, &second.runtime.capacity)); assert!(!Arc::ptr_eq(&first.runtime.queued, &second.runtime.queued)); let _first_queue = Arc::clone(&first.runtime.queued) .try_acquire_owned() .expect("first generation queue capacity"); assert_eq!(first.runtime.queued.available_permits(), DEFAULT_MAX_QUEUED_OPERATIONS - 1); assert_eq!(second.runtime.queued.available_permits(), DEFAULT_MAX_QUEUED_OPERATIONS); trip_breaker(&first.runtime); assert!(first.runtime.rejects(Instant::now()).expect("first generation breaker state")); assert!( !second .runtime .rejects(Instant::now()) .expect("second generation breaker state") ); } #[test] fn backend_capacity_domains_isolate_credentials_queue_and_breaker() { let endpoint = unique_endpoint("shared-capacity"); let config = KmsConfig::default(); let operations = RetryPolicy::for_backend(&config, "vault", &endpoint, Some("namespace"), "operations"); let login = RetryPolicy::for_credentials(&config, "vault", &endpoint, Some("namespace"), "login"); let renew = RetryPolicy::for_credentials(&config, "vault", &endpoint, Some("namespace"), "renew"); assert!(!Arc::ptr_eq(&operations.runtime, &login.runtime)); assert!(Arc::ptr_eq(&operations.runtime.capacity, &login.runtime.capacity)); assert!(matches!(operations.runtime.capacity_class, CapacityClass::Operations)); assert!(matches!(login.runtime.capacity_class, CapacityClass::Credentials)); assert!(Arc::ptr_eq(&login.runtime.capacity, &renew.runtime.capacity)); assert!(!Arc::ptr_eq(&operations.runtime.queued, &login.runtime.queued)); let _operations_queue = Arc::clone(&operations.runtime.queued) .try_acquire_owned() .expect("operations queue capacity"); assert_eq!(operations.runtime.queued.available_permits(), DEFAULT_MAX_QUEUED_OPERATIONS - 1); assert_eq!(login.runtime.queued.available_permits(), DEFAULT_MAX_QUEUED_OPERATIONS); trip_breaker(&operations.runtime); assert!(operations.runtime.rejects(Instant::now()).expect("operations breaker state")); assert!(!login.runtime.rejects(Instant::now()).expect("login breaker state")); } #[tokio::test(start_paused = true)] async fn credential_refresh_is_not_starved_by_operations_capacity() { let endpoint = unique_endpoint("credential-capacity"); let config = KmsConfig::default(); let operations = RetryPolicy::for_backend(&config, "vault", &endpoint, Some("namespace"), "operations"); let login = RetryPolicy::for_credentials(&config, "vault", &endpoint, Some("namespace"), "credentials-login"); let _operation_permits = (0..DEFAULT_MAX_CONCURRENT_OPERATIONS - RESERVED_CREDENTIAL_OPERATIONS) .map(|_| operations.runtime.try_acquire_active().expect("operations capacity")) .collect::>(); assert_eq!(login.runtime.capacity.total.available_permits(), RESERVED_CREDENTIAL_OPERATIONS); execute_with_jitter( "credential_login", OpClass::Auth, &login, &CancellationToken::new(), full_jitter, || async { Ok(()) }, ) .await .expect("credential login must use reserved capacity"); } #[test] fn backend_identity_isolates_endpoint_and_namespace() { let config = KmsConfig::default(); let endpoint_a = unique_endpoint("endpoint-a"); let endpoint_b = unique_endpoint("endpoint-b"); let first_endpoint = RetryPolicy::for_backend(&config, "vault", &endpoint_a, Some("namespace"), "operations"); let second_endpoint = RetryPolicy::for_backend(&config, "vault", &endpoint_b, Some("namespace"), "operations"); assert!(!Arc::ptr_eq(&first_endpoint.runtime.capacity, &second_endpoint.runtime.capacity)); let namespace_endpoint = unique_endpoint("namespace"); let first_namespace = RetryPolicy::for_backend(&config, "vault", &namespace_endpoint, Some("tenant-a"), "operations"); let second_namespace = RetryPolicy::for_backend(&config, "vault", &namespace_endpoint, Some("tenant-b"), "operations"); assert!(!Arc::ptr_eq(&first_namespace.runtime.capacity, &second_namespace.runtime.capacity)); } #[tokio::test(start_paused = true)] async fn admission_bounds_active_queue_and_deadline() { let mut policy = policy_with_limit(1, 1); policy.op_deadline = Duration::from_millis(100); let active = Arc::clone(&policy.runtime.capacity.total) .acquire_owned() .await .expect("active permit"); let queued = Arc::clone(&policy.runtime.queued) .acquire_owned() .await .expect("queued permit"); let excess: Result<()> = execute_with_jitter( "queue_full", OpClass::ReadIdempotent, &policy, &CancellationToken::new(), full_jitter, || async { Ok(()) }, ) .await; assert!(matches!(&excess, Err(KmsError::BackendError { message }) if message == BACKPRESSURE_MESSAGE)); drop(queued); tokio::spawn(async move { tokio::time::sleep(Duration::from_millis(90)).await; drop(active); }); let started = Instant::now(); let timed_out: Result<()> = execute_with_jitter( "queue_deadline", OpClass::ReadIdempotent, &policy, &CancellationToken::new(), full_jitter, std::future::pending::>, ) .await; assert!(matches!(timed_out, Err(KmsError::OperationTimedOut { .. }))); assert_eq!(started.elapsed(), Duration::from_millis(100)); let active = Arc::clone(&policy.runtime.capacity.total) .acquire_owned() .await .expect("active permit"); tokio::spawn(async move { tokio::time::sleep(Duration::from_millis(100)).await; drop(active); }); let boundary = acquire_attempt( Arc::clone(&policy.runtime), "deadline_boundary", &CancellationToken::new(), Instant::now() + Duration::from_millis(100), ) .await; assert!(matches!(boundary, Err((Outcome::BackpressureTimeout, _)))); } #[tokio::test(start_paused = true)] async fn queued_admission_cancellation_releases_queue_capacity() { let policy = Arc::new(policy_with_limit(1, 1)); let active = Arc::clone(&policy.runtime.capacity.total) .acquire_owned() .await .expect("active permit"); let cancel = CancellationToken::new(); let calls = Arc::new(AtomicU32::new(0)); let queued = tokio::spawn({ let policy = Arc::clone(&policy); let cancel = cancel.clone(); let calls = Arc::clone(&calls); async move { execute_with_jitter("cancel_queued", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, move || { calls.fetch_add(1, Ordering::SeqCst); std::future::ready(Ok(())) }) .await } }); for _ in 0..100 { if policy.runtime.queued.available_permits() == 0 { break; } tokio::task::yield_now().await; } assert_eq!(policy.runtime.queued.available_permits(), 0, "request must be queued"); cancel.cancel(); let result = queued.await.expect("queued task join"); assert!(matches!(result, Err(KmsError::OperationCancelled { .. })), "got {result:?}"); assert_eq!(calls.load(Ordering::SeqCst), 0); assert_eq!(policy.runtime.queued.available_permits(), 1); drop(active); } #[tokio::test(start_paused = true)] async fn queued_admission_observes_circuit_opened_while_waiting() { let policy = Arc::new(policy_with_limit(1, 1)); let active = Arc::clone(&policy.runtime.capacity.total) .acquire_owned() .await .expect("active permit"); let calls = Arc::new(AtomicU32::new(0)); let queued = tokio::spawn({ let policy = Arc::clone(&policy); let calls = Arc::clone(&calls); async move { execute_with_jitter( "open_while_queued", OpClass::ReadIdempotent, &policy, &CancellationToken::new(), full_jitter, move || { calls.fetch_add(1, Ordering::SeqCst); std::future::ready(Ok(())) }, ) .await } }); for _ in 0..100 { if policy.runtime.queued.available_permits() == 0 { break; } tokio::task::yield_now().await; } assert_eq!(policy.runtime.queued.available_permits(), 0, "request must be queued"); trip_breaker(&policy.runtime); drop(active); let result = queued.await.expect("queued task join"); assert!(result.as_ref().is_err_and(is_circuit_open), "got {result:?}"); assert_eq!(calls.load(Ordering::SeqCst), 0); assert_eq!(policy.runtime.queued.available_permits(), 1); assert_eq!(policy.runtime.capacity.total.available_permits(), 1); } #[tokio::test(start_paused = true)] async fn breaker_bounds_attempts_and_recovers_with_one_probe() { let mut policy = policy_with_limit(2, 2); policy.max_attempts = 10; let policy = Arc::new(policy); let calls = Arc::new(AtomicU32::new(0)); let cancel = CancellationToken::new(); let failed: Result<()> = execute_with_jitter("breaker_open", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, { let calls = Arc::clone(&calls); move || { let calls = Arc::clone(&calls); async move { calls.fetch_add(1, Ordering::SeqCst); Err(retryable_conn_error()) } } }) .await; failed.expect_err("breaker threshold must stop retries"); assert_eq!(calls.load(Ordering::SeqCst), 5); for _ in 0..3 { let rejected: Result<()> = execute_with_jitter("breaker_rejected", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, || async { Ok(()) }) .await; assert!(rejected.as_ref().is_err_and(is_circuit_open), "got {rejected:?}"); } let already_cancelled = CancellationToken::new(); already_cancelled.cancel(); let cancelled: Result<()> = execute_with_jitter( "breaker_cancelled", OpClass::ReadIdempotent, &policy, &already_cancelled, full_jitter, || async { Ok(()) }, ) .await; assert!(matches!(cancelled, Err(KmsError::OperationCancelled { .. }))); tokio::time::advance(DEFAULT_CIRCUIT_OPEN_DURATION).await; let started = Arc::new(Notify::new()); let release = Arc::new(Notify::new()); let probe = tokio::spawn({ let policy = Arc::clone(&policy); let calls = Arc::clone(&calls); let started = Arc::clone(&started); let release = Arc::clone(&release); async move { execute_with_jitter( "breaker_probe", OpClass::ReadIdempotent, &policy, &CancellationToken::new(), full_jitter, move || { let calls = Arc::clone(&calls); let started = Arc::clone(&started); let release = Arc::clone(&release); async move { calls.fetch_add(1, Ordering::SeqCst); started.notify_one(); release.notified().await; Ok(()) } }, ) .await } }); started.notified().await; let concurrent: Result<()> = execute_with_jitter( "breaker_concurrent_probe", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, || async { Ok(()) }, ) .await; assert!(concurrent.as_ref().is_err_and(is_circuit_open), "got {concurrent:?}"); release.notify_one(); probe.await.expect("probe join").expect("successful probe must close circuit"); execute_with_jitter("breaker_recovered", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, || async { Ok(()) }) .await .expect("closed circuit must recover"); assert_eq!(calls.load(Ordering::SeqCst), 6); } #[tokio::test(start_paused = true)] async fn breaker_stale_and_failed_probes_preserve_state() { let runtime = test_runtime(2, 2); let stale = runtime.admit(Instant::now()).expect("state").expect("closed"); trip_breaker(&runtime); tokio::time::advance(DEFAULT_CIRCUIT_OPEN_DURATION).await; let active = Arc::clone(&runtime.capacity.total).try_acquire_owned().expect("capacity"); let probe = runtime.admit(Instant::now()).expect("state").expect("probe"); drop(RuntimePermit::new(Arc::clone(&runtime), probe, active)); tokio::time::advance(DEFAULT_CIRCUIT_OPEN_DURATION - Duration::from_millis(1)).await; assert!(runtime.rejects(Instant::now()).expect("state")); tokio::time::advance(Duration::from_millis(1)).await; let probe = runtime.admit(Instant::now()).expect("state").expect("probe"); assert!(runtime.complete(probe, Some(ErrorClass::RetryableStatus))); tokio::time::advance(DEFAULT_CIRCUIT_OPEN_DURATION).await; let recovery = runtime.admit(Instant::now()).expect("state").expect("probe"); assert!(!runtime.complete(recovery, Some(ErrorClass::Fatal))); assert!(!runtime.complete(stale, Some(ErrorClass::RetryableConn))); let reset = test_runtime(2, 2); for _ in 0..4 { let admission = reset.admit(Instant::now()).expect("state").expect("closed"); assert!(!reset.complete(admission, Some(ErrorClass::RetryableConn))); } let fatal = reset.admit(Instant::now()).expect("state").expect("closed"); assert!(!reset.complete(fatal, Some(ErrorClass::Fatal))); for _ in 0..4 { let admission = reset.admit(Instant::now()).expect("state").expect("closed"); assert!(!reset.complete(admission, Some(ErrorClass::RetryableConn))); } } #[tokio::test(start_paused = true)] async fn total_duration_never_exceeds_deadline() { // Worst case without a deadline would be 5 * 10s + backoff; the 25s // deadline must cut both the attempt count and the final attempt short. let policy = policy_of(10_000, 25_000, 5, 100, 2_000); let cancel = CancellationToken::new(); let calls = Arc::new(AtomicU32::new(0)); let calls_in_attempt = calls.clone(); let started = Instant::now(); let result: Result<()> = execute_with_jitter("describe_key", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, move || { calls_in_attempt.fetch_add(1, Ordering::SeqCst); std::future::pending::>() }) .await; assert!(matches!(result, Err(KmsError::OperationTimedOut { .. })), "got {result:?}"); // attempt 1: 10s + 100ms backoff; attempt 2: 10s + 200ms backoff; // attempt 3 runs with the residual 4.7s budget, then the deadline is // exhausted and no further backoff or attempt happens. assert_eq!(calls.load(Ordering::SeqCst), 3); assert_eq!(started.elapsed(), Duration::from_millis(25_000)); } #[tokio::test(start_paused = true)] async fn default_jitter_stays_within_equal_jitter_bounds() { let policy = policy_of(1_000, 60_000, 3, 100, 2_000); let cancel = CancellationToken::new(); let calls = Arc::new(AtomicU32::new(0)); let calls_in_attempt = calls.clone(); let started = Instant::now(); let result = execute("health_check", OpClass::ReadIdempotent, &policy, &cancel, move || { let calls = calls_in_attempt.clone(); async move { if calls.fetch_add(1, Ordering::SeqCst) < 2 { Err(retryable_conn_error()) } else { Ok(()) } } }) .await; result.expect("retries within budget must succeed"); let elapsed = started.elapsed(); // Equal jitter sleeps within [cap / 2, cap]; caps are 100ms then 200ms. assert!( elapsed >= Duration::from_millis(150) && elapsed <= Duration::from_millis(300), "elapsed {elapsed:?} outside equal-jitter bounds" ); } #[test] fn equal_jitter_is_seed_deterministic_and_bounded() { let caps = [Duration::from_millis(100), Duration::from_millis(250), Duration::from_secs(2)]; let mut first = StdRng::seed_from_u64(1569); let mut second = StdRng::seed_from_u64(1569); for cap in caps { let a = equal_jitter(&mut first, cap); let b = equal_jitter(&mut second, cap); assert_eq!(a, b, "same seed must yield the same jitter sequence"); assert!(a >= cap / 2 && a <= cap, "jitter {a:?} outside [{:?}, {cap:?}]", cap / 2); } } #[test] fn backoff_caps_double_up_to_max() { let policy = policy_of(1_000, 60_000, 6, 100, 700); let caps: Vec = (1..=5).map(|completed| backoff_cap(&policy, completed)).collect(); assert_eq!( caps, vec![ Duration::from_millis(100), Duration::from_millis(200), Duration::from_millis(400), Duration::from_millis(700), Duration::from_millis(700), ] ); } #[test] fn retry_policy_from_config_applies_clamps() { let config = KmsConfig { timeout: Duration::from_secs(3_600), retry_attempts: 50, ..KmsConfig::default() }; let policy = RetryPolicy::for_backend(&config, "test-backend", "https://example.invalid", None, "config-clamp"); assert_eq!(policy.attempt_timeout, Duration::from_secs(300)); assert_eq!(policy.max_attempts, 10); // The deadline must cover the full worst case so it never cuts a // policy-conformant operation short. assert!(policy.op_deadline >= policy.attempt_timeout.saturating_mul(policy.max_attempts)); let in_range = KmsConfig::default(); let policy = RetryPolicy::for_backend(&in_range, "test-backend", "https://example.invalid", None, "config-in-range"); assert_eq!(policy.attempt_timeout, in_range.timeout); assert_eq!(policy.max_attempts, in_range.retry_attempts); } #[test] fn classify_vaultrs_matrix() { use rustify::errors::ClientError as RestError; use vaultrs::error::ClientError; let api = |code: u16| ClientError::APIError { code, errors: vec![] }; for code in [429u16, 500, 502, 503, 504] { assert_eq!(classify_vaultrs(&api(code)), ErrorClass::RetryableStatus, "status {code}"); } for code in [400u16, 401, 403, 404, 405, 412, 501] { assert_eq!(classify_vaultrs(&api(code)), ErrorClass::Fatal, "status {code}"); } // Status errors whose body could not be parsed stay wrapped in the // rustify error; classification must still see the code. let raw_status = |code: u16| ClientError::RestClientError { source: RestError::ServerResponseError { code, content: None }, }; assert_eq!(classify_vaultrs(&raw_status(503)), ErrorClass::RetryableStatus); assert_eq!(classify_vaultrs(&raw_status(403)), ErrorClass::Fatal); let send_failure = ClientError::RestClientError { source: RestError::RequestError { source: anyhow::anyhow!("connection refused"), url: "http://127.0.0.1:8200/v1/sys/health".to_string(), method: "GET".to_string(), }, }; assert_eq!(classify_vaultrs(&send_failure), ErrorClass::RetryableConn); let read_failure = ClientError::RestClientError { source: RestError::ResponseError { source: anyhow::anyhow!("connection reset by peer"), }, }; assert_eq!(classify_vaultrs(&read_failure), ErrorClass::RetryableConn); let malformed = ClientError::JsonParseError { source: serde_json::from_str::("{").expect_err("truncated JSON must not parse"), }; assert_eq!(classify_vaultrs(&malformed), ErrorClass::Fatal); assert_eq!(classify_vaultrs(&ClientError::ResponseEmptyError), ErrorClass::Fatal); assert_eq!(classify_vaultrs(&ClientError::InvalidLoginMethodError), ErrorClass::Fatal); } // -- Metric emission ---------------------------------------------------- // // Each test installs a thread-local debugging recorder and drives a // paused-clock current-thread runtime inside it, so the emitted metrics // (including virtual-clock durations) are fully deterministic. use metrics_util::MetricKind; use metrics_util::debugging::{DebugValue, DebuggingRecorder}; type MetricEntry = ( metrics_util::CompositeKey, Option, Option, DebugValue, ); /// Run `test` on a paused current-thread runtime under a debugging /// recorder and return one snapshot of everything it emitted. /// /// A single snapshot per test on purpose: `Snapshotter::snapshot` drains /// the recorded state, so taking it per assertion would only show the /// first assertion any data. fn record_metrics(test: impl FnOnce() -> std::pin::Pin>>) -> (Vec, Out) { let recorder = DebuggingRecorder::new(); let snapshotter = recorder.snapshotter(); let out = metrics::with_local_recorder(&recorder, || { let runtime = tokio::runtime::Builder::new_current_thread() .enable_time() .start_paused(true) .build() .expect("current-thread runtime must build"); runtime.block_on(test()) }); (snapshotter.snapshot().into_vec(), out) } fn labels_match(key: &metrics::Key, labels: &[(&str, &str)]) -> bool { labels .iter() .all(|(label, expected)| key.labels().any(|l| l.key() == *label && l.value() == *expected)) } fn counter_value(snapshot: &[MetricEntry], name: &str, labels: &[(&str, &str)]) -> u64 { snapshot .iter() .filter_map(|(composite, _unit, _description, value)| { let matches = composite.kind() == MetricKind::Counter && composite.key().name() == name && labels_match(composite.key(), labels); match (matches, value) { (true, DebugValue::Counter(count)) => Some(*count), _ => None, } }) .sum() } fn gauge_value(snapshot: &[MetricEntry], name: &str, labels: &[(&str, &str)]) -> Option { snapshot.iter().find_map(|(composite, _unit, _description, value)| { let matches = composite.kind() == MetricKind::Gauge && composite.key().name() == name && labels_match(composite.key(), labels); match (matches, value) { (true, DebugValue::Gauge(value)) => Some(value.into_inner()), _ => None, } }) } fn histogram_values(snapshot: &[MetricEntry], name: &str, labels: &[(&str, &str)]) -> Vec { snapshot .iter() .filter_map(|(composite, _unit, _description, value)| { let matches = composite.kind() == MetricKind::Histogram && composite.key().name() == name && labels_match(composite.key(), labels); match (matches, value) { (true, DebugValue::Histogram(values)) => Some(values), _ => None, } }) .flatten() .map(|value| value.into_inner()) .collect() } #[test] fn metrics_record_retried_success_with_attempts_and_duration() { let calls_in_test = Arc::new(AtomicU32::new(0)); let (snapshot, ()) = record_metrics(move || { Box::pin(async move { let policy = policy_of(1_000, 60_000, 3, 100, 2_000); let cancel = CancellationToken::new(); let calls_in_attempt = calls_in_test.clone(); execute_with_jitter("metrics_read", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, move || { let calls = calls_in_attempt.clone(); async move { if calls.fetch_add(1, Ordering::SeqCst) < 2 { Err(AttemptError { class: ErrorClass::RetryableStatus, error: KmsError::backend_error("throttled (429)"), }) } else { Ok(()) } } }) .await .expect("retries within budget must succeed"); }) }); assert_eq!( counter_value( &snapshot, METRIC_OPERATIONS_TOTAL, &[ ("operation", "metrics_read"), ("op_class", "read_idempotent"), ("outcome", "success") ] ), 1 ); assert_eq!( counter_value( &snapshot, METRIC_ATTEMPT_FAILURES_TOTAL, &[("operation", "metrics_read"), ("error_class", "retryable_status")] ), 2 ); assert_eq!( histogram_values( &snapshot, METRIC_OPERATION_ATTEMPTS, &[("operation", "metrics_read"), ("outcome", "success")] ), vec![3.0] ); // Full-cap backoffs of 100ms and 200ms on the paused clock. let durations = histogram_values( &snapshot, METRIC_OPERATION_DURATION_SECONDS, &[("operation", "metrics_read"), ("outcome", "success")], ); assert_eq!(durations.len(), 1); assert!((durations[0] - 0.3).abs() < 1e-9, "expected 0.3s of virtual backoff, got {durations:?}"); } /// Request, error and latency series must be attributable to one backend. /// /// Operation names are shared across backends (`decrypt` exists on every /// one of them), so without a `backend` label a Vault Transit latency /// regression is indistinguishable from an AWS one in the same series. #[test] fn operation_metrics_separate_series_per_backend() { fn policy_for(backend: &'static str) -> RetryPolicy { let mut policy = policy_of(1_000, 60_000, 2, 100, 2_000); policy.runtime = named_test_runtime(backend, DEFAULT_MAX_CONCURRENT_OPERATIONS, DEFAULT_MAX_QUEUED_OPERATIONS); policy } let (snapshot, ()) = record_metrics(|| { Box::pin(async { let cancel = CancellationToken::new(); execute_with_jitter( "decrypt", OpClass::ReadIdempotent, &policy_for("vault-transit"), &cancel, full_jitter, || async { Ok(()) }, ) .await .expect("transit attempt succeeds"); let result: Result<()> = execute_with_jitter("decrypt", OpClass::ReadIdempotent, &policy_for("aws"), &cancel, full_jitter, || async { Err(AttemptError { class: ErrorClass::Fatal, error: KmsError::access_denied("permission denied (403)"), }) }) .await; assert!(matches!(result, Err(KmsError::AccessDenied { .. })), "got {result:?}"); }) }); let succeeded = [("operation", "decrypt"), ("outcome", "success")]; let failed = [("operation", "decrypt"), ("outcome", "fatal")]; assert_eq!(counter_value(&snapshot, METRIC_OPERATIONS_TOTAL, &[("backend", "vault-transit")]), 1); assert_eq!(counter_value(&snapshot, METRIC_OPERATIONS_TOTAL, &[("backend", "aws")]), 1); assert_eq!( counter_value(&snapshot, METRIC_OPERATIONS_TOTAL, &[("backend", "aws"), succeeded[0], succeeded[1]]), 0, "the AWS failure must not be counted under the transit success series" ); assert_eq!( counter_value( &snapshot, METRIC_ATTEMPT_FAILURES_TOTAL, &[("backend", "aws"), ("operation", "decrypt"), ("error_class", "fatal")] ), 1 ); assert_eq!( counter_value( &snapshot, METRIC_ATTEMPT_FAILURES_TOTAL, &[("backend", "vault-transit"), ("operation", "decrypt")] ), 0 ); for (backend, labels) in [("vault-transit", succeeded), ("aws", failed)] { let with_backend = [("backend", backend), labels[0], labels[1]]; assert_eq!( histogram_values(&snapshot, METRIC_OPERATION_DURATION_SECONDS, &with_backend).len(), 1, "{backend} must own a latency series of its own" ); assert_eq!(histogram_values(&snapshot, METRIC_OPERATION_ATTEMPTS, &with_backend), vec![1.0]); } } #[test] fn metrics_record_fatal_outcome_with_single_attempt() { let (snapshot, ()) = record_metrics(|| { Box::pin(async { let policy = policy_of(1_000, 60_000, 5, 100, 2_000); let cancel = CancellationToken::new(); let result: Result<()> = execute_with_jitter("metrics_fatal", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, || async { Err(AttemptError { class: ErrorClass::Fatal, error: KmsError::access_denied("permission denied (403)"), }) }) .await; result.expect_err("a fatal failure must end the operation"); }) }); assert_eq!( counter_value( &snapshot, METRIC_OPERATIONS_TOTAL, &[("operation", "metrics_fatal"), ("outcome", "fatal")] ), 1 ); assert_eq!( counter_value( &snapshot, METRIC_ATTEMPT_FAILURES_TOTAL, &[("operation", "metrics_fatal"), ("error_class", "fatal")] ), 1 ); assert_eq!( histogram_values( &snapshot, METRIC_OPERATION_ATTEMPTS, &[("operation", "metrics_fatal"), ("outcome", "fatal")] ), vec![1.0] ); } #[test] fn metrics_record_mutating_budget_exhausted_after_one_attempt() { let (snapshot, ()) = record_metrics(|| { Box::pin(async { let policy = policy_of(1_000, 60_000, 5, 100, 2_000); let cancel = CancellationToken::new(); let result: Result<()> = execute_with_jitter( "metrics_rotate", OpClass::MutatingNonIdempotent, &policy, &cancel, full_jitter, || async { Err(retryable_conn_error()) }, ) .await; result.expect_err("a mutating operation must not retry a retryable failure"); }) }); assert_eq!( counter_value( &snapshot, METRIC_OPERATIONS_TOTAL, &[ ("operation", "metrics_rotate"), ("op_class", "mutating_non_idempotent"), ("outcome", "budget_exhausted") ] ), 1 ); assert_eq!( histogram_values( &snapshot, METRIC_OPERATION_ATTEMPTS, &[("operation", "metrics_rotate"), ("outcome", "budget_exhausted")] ), vec![1.0] ); } #[test] fn metrics_record_timeouts_and_deadline_outcome() { let (snapshot, ()) = record_metrics(|| { Box::pin(async { // Hung attempts: 10s each against a 25s deadline (see // total_duration_never_exceeds_deadline for the timeline). let policy = policy_of(10_000, 25_000, 5, 100, 2_000); let cancel = CancellationToken::new(); let result: Result<()> = execute_with_jitter("metrics_hung", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, || { std::future::pending::>() }) .await; result.expect_err("hung attempts must exhaust the deadline"); }) }); assert_eq!( counter_value( &snapshot, METRIC_OPERATIONS_TOTAL, &[("operation", "metrics_hung"), ("outcome", "deadline_exceeded")] ), 1 ); assert_eq!( counter_value( &snapshot, METRIC_ATTEMPT_FAILURES_TOTAL, &[("operation", "metrics_hung"), ("error_class", "attempt_timeout")] ), 3 ); let durations = histogram_values( &snapshot, METRIC_OPERATION_DURATION_SECONDS, &[("operation", "metrics_hung"), ("outcome", "deadline_exceeded")], ); assert_eq!(durations.len(), 1); assert!((durations[0] - 25.0).abs() < 1e-9, "expected the full 25s deadline, got {durations:?}"); } #[test] fn metrics_record_cancelled_outcome() { let calls_in_test = Arc::new(AtomicU32::new(0)); let (snapshot, ()) = record_metrics(move || { Box::pin(async move { let policy = policy_of(1_000, 600_000, 5, 10_000, 10_000); let cancel = CancellationToken::new(); let canceller = cancel.clone(); tokio::spawn(async move { tokio::time::sleep(Duration::from_millis(500)).await; canceller.cancel(); }); let calls_in_attempt = calls_in_test.clone(); let result: Result<()> = execute_with_jitter("metrics_cancel", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, move || { let calls = calls_in_attempt.clone(); async move { calls.fetch_add(1, Ordering::SeqCst); Err(retryable_conn_error()) } }) .await; result.expect_err("cancellation must abort the backoff"); }) }); assert_eq!( counter_value( &snapshot, METRIC_OPERATIONS_TOTAL, &[("operation", "metrics_cancel"), ("outcome", "cancelled")] ), 1 ); assert_eq!( histogram_values( &snapshot, METRIC_OPERATION_ATTEMPTS, &[("operation", "metrics_cancel"), ("outcome", "cancelled")] ), vec![1.0] ); } #[test] fn metrics_record_backpressure_and_circuit_outcomes() { let (snapshot, ()) = record_metrics(|| { Box::pin(async { let rejected = policy_with_limit(1, 0); let rejected_active = Arc::clone(&rejected.runtime.capacity.total) .try_acquire_owned() .expect("active capacity"); let result: Result<()> = execute_with_jitter( "metrics_backpressure_rejected", OpClass::ReadIdempotent, &rejected, &CancellationToken::new(), full_jitter, || async { Ok(()) }, ) .await; assert!(matches!(result, Err(KmsError::BackendError { .. }))); drop(rejected_active); let mut timed_out = policy_with_limit(1, 1); timed_out.op_deadline = Duration::from_millis(100); let timeout_active = Arc::clone(&timed_out.runtime.capacity.total) .try_acquire_owned() .expect("active capacity"); let result: Result<()> = execute_with_jitter( "metrics_backpressure_timeout", OpClass::ReadIdempotent, &timed_out, &CancellationToken::new(), full_jitter, || async { Ok(()) }, ) .await; assert!(matches!(result, Err(KmsError::OperationTimedOut { .. }))); drop(timeout_active); let open = policy_with_limit(1, 1); trip_breaker(&open.runtime); let result: Result<()> = execute_with_jitter( "metrics_circuit_open", OpClass::ReadIdempotent, &open, &CancellationToken::new(), full_jitter, || async { Ok(()) }, ) .await; assert!(result.as_ref().is_err_and(is_circuit_open), "got {result:?}"); let mut trip = policy_with_limit(1, 1); trip.max_attempts = DEFAULT_CIRCUIT_FAILURE_THRESHOLD + 1; let result: Result<()> = execute_with_jitter( "metrics_circuit_trip", OpClass::ReadIdempotent, &trip, &CancellationToken::new(), full_jitter, || async { Err(retryable_conn_error()) }, ) .await; assert!(matches!(result, Err(KmsError::BackendError { .. }))); }) }); for (operation, outcome, attempts) in [ ("metrics_backpressure_rejected", "backpressure_rejected", 0.0), ("metrics_backpressure_timeout", "backpressure_timeout", 0.0), ("metrics_circuit_open", "circuit_open", 0.0), ("metrics_circuit_trip", "circuit_open", f64::from(DEFAULT_CIRCUIT_FAILURE_THRESHOLD)), ] { assert_eq!( counter_value(&snapshot, METRIC_OPERATIONS_TOTAL, &[("operation", operation), ("outcome", outcome)]), 1 ); assert_eq!( histogram_values(&snapshot, METRIC_OPERATION_ATTEMPTS, &[("operation", operation), ("outcome", outcome)]), vec![attempts] ); } } #[test] fn successful_and_fatal_half_open_probes_clear_the_open_gauge() { for failure in [None, Some(ErrorClass::Fatal)] { let (snapshot, runtime) = record_metrics(move || { Box::pin(async move { let runtime = test_runtime(1, 1); trip_breaker(&runtime); tokio::time::advance(DEFAULT_CIRCUIT_OPEN_DURATION).await; let probe = runtime.admit(Instant::now()).expect("state").expect("half-open probe"); assert!(!runtime.complete(probe, failure)); runtime }) }); assert!(!runtime.rejects(Instant::now()).expect("closed breaker state")); assert_eq!( gauge_value(&snapshot, METRIC_CIRCUIT_OPEN, &[("backend", "test-backend"), ("scope", "test-scope")]), Some(0.0) ); } } #[test] fn cancelling_half_open_attempt_releases_capacity_and_reopens_breaker() { let (snapshot, policy) = record_metrics(|| { Box::pin(async { let policy = policy_with_limit(1, 1); trip_breaker(&policy.runtime); tokio::time::advance(DEFAULT_CIRCUIT_OPEN_DURATION).await; let cancel = CancellationToken::new(); let canceller = cancel.clone(); tokio::spawn(async move { tokio::time::sleep(Duration::from_millis(1)).await; canceller.cancel(); }); let result: Result<()> = execute_with_jitter("cancel_half_open_probe", OpClass::ReadIdempotent, &policy, &cancel, full_jitter, || { std::future::pending::>() }) .await; assert!(matches!(result, Err(KmsError::OperationCancelled { .. }))); assert_eq!(policy.runtime.capacity.total.available_permits(), 1); assert_eq!(policy.runtime.queued.available_permits(), 1); assert!(policy.runtime.rejects(Instant::now()).expect("reopened breaker state")); tokio::time::advance(DEFAULT_CIRCUIT_OPEN_DURATION).await; let probe = policy .runtime .admit(Instant::now()) .expect("state") .expect("replacement half-open probe"); assert!(probe.probe); policy }) }); let labels = [("backend", "test-backend"), ("scope", "test-scope")]; assert_eq!(gauge_value(&snapshot, METRIC_IN_FLIGHT, &labels), Some(0.0)); assert_eq!(gauge_value(&snapshot, METRIC_CIRCUIT_OPEN, &labels), Some(1.0)); drop(policy); } #[test] fn runtime_metrics_track_in_flight_and_open_lifecycle() { let in_flight_recorder = DebuggingRecorder::new(); let in_flight_snapshotter = in_flight_recorder.snapshotter(); let in_flight = metrics::with_local_recorder(&in_flight_recorder, || { let runtime = test_runtime(1, 1); let active = Arc::clone(&runtime.capacity.total).try_acquire_owned().expect("capacity"); let admission = runtime.admit(Instant::now()).expect("state").expect("closed"); let _permit = RuntimePermit::new(Arc::clone(&runtime), admission, active); in_flight_snapshotter.snapshot().into_vec() }); let dropped_recorder = DebuggingRecorder::new(); let dropped_snapshotter = dropped_recorder.snapshotter(); let dropped = metrics::with_local_recorder(&dropped_recorder, || { let runtime = test_runtime(1, 1); trip_breaker(&runtime); drop(runtime); dropped_snapshotter.snapshot().into_vec() }); let labels = [("backend", "test-backend"), ("scope", "test-scope")]; assert_eq!(gauge_value(&in_flight, METRIC_IN_FLIGHT, &labels), Some(1.0)); assert_eq!(gauge_value(&dropped, METRIC_CIRCUIT_OPEN, &labels), Some(0.0)); } }