gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-08-09 14:49:25 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
a0a8eaa0f3c8d54aee4d27a2d387b226b78e5e88
rustfs/crates/kms/src/policy.rs
T
Zhengchao An 8003912bb1 fix(kms): report unreadable keys and bound list-keys page size (#5764) 
A key that cannot be described was handled two incompatible ways. Vault KV2 swallowed every describe failure and dropped the key from the page, so a damaged or newer-format record silently disappeared from the operator's inventory and from the deletion sweep's census. Local failed the whole listing instead, so one bad record stopped every scheduled deletion on the node for as long as the damage lasted. Both force a per-key problem into a whole-page answer.

ListKeysResponse now carries unreadable_key_ids, and the backends that read local key records classify per-key failures in one place: KeyNotFound is a concurrent deletion and is skipped, a material-level error names the key on the page, and anything else fails the listing, because it says nothing about a particular key and reporting it as key damage would turn a backend outage into a false data-loss alarm. A listing that covered the entire key set and found nothing readable still fails, since an empty page there is indistinguishable from a deployment with no keys; the guard is scoped to a page with no successor so a damaged key can never strand the keys behind it. The deletion sweep destroys the expired keys it can read, counts the unreadable ones, and withholds its lifecycle gauges rather than publishing a census over a key set it did not fully see.

Vault Transit needs the same treatment and is easy to miss: its per-key metadata records live in KV2 too, so folding every non-404 failure into a backend error left its per-key classification unreachable and one metadata record written by a newer build still failed every listing on the node.

Vault KV2 record reads gain the typed errors this needs: an unparseable body is MaterialCorrupt and an absent data envelope is MaterialMissing, where both were previously indistinguishable from Vault being unreachable. Only the parse failure's category and position are reported, because serde's own message embeds the offending scalar and that message reaches a log line and an admin HTTP body.

The admin list handlers refuse a malformed limit with 400 instead of silently substituting the default page size, and every page is capped at 1000 where it is cut, so a single request can no longer fan out one metadata lookup per key without bound. The four operation-level KMS metrics gain a backend label, since operation names are shared across backends and a Transit latency regression was previously indistinguishable from an AWS one. The Static backend captures its reported creation date once instead of reading the clock on every describe and list. POST /kms/clear-cache gains a named response type with an unchanged wire shape.
2026-08-06 22:01:30 +08:00

2206 lines
86 KiB
Rust
Raw Blame History

// 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<BackendRuntime>,
}
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<Mutex<HashMap<[u8; 32], Weak<BackendCapacity>>>> = LazyLock::new(|| Mutex::new(HashMap::new()));
#[derive(Debug)]
struct BackendCapacity {
total: Arc<Semaphore>,
operations: Arc<Semaphore>,
}
#[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<BackendCapacity>,
capacity_class: CapacityClass,
queued: Arc<Semaphore>,
breaker: Mutex<CircuitState>,
in_flight: metrics::Gauge,
circuit_open: metrics::Gauge,
}
impl BackendRuntime {
fn shared(identity: [u8; 32], backend: &'static str, scope: &'static str, capacity: CapacityClass) -> Arc<Self> {
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<BackendCapacity>, 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<ActivePermits, TryAcquireError> {
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<ActivePermits, AcquireError> {
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<std::sync::MutexGuard<'_, CircuitState>> {
self.breaker
.lock()
.map_err(|_| KmsError::backend_error("KMS backend circuit state is unavailable"))
}
fn rejects(&self, now: Instant) -> Result<bool> {
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<Option<CircuitAdmission>> {
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<ErrorClass>) -> 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<Instant>,
probe_in_flight: bool,
}
#[derive(Debug, Clone, Copy)]
struct CircuitAdmission {
generation: u64,
probe: bool,
}
struct RuntimePermit {
runtime: Arc<BackendRuntime>,
admission: Option<CircuitAdmission>,
_active: ActivePermits,
}
struct ActivePermits {
_total: OwnedSemaphorePermit,
_operations: Option<OwnedSemaphorePermit>,
}
impl RuntimePermit {
#[cfg(test)]
fn new(runtime: Arc<BackendRuntime>, admission: CircuitAdmission, active: OwnedSemaphorePermit) -> Self {
Self::with_active(
runtime,
admission,
ActivePermits {
_total: active,
_operations: None,
},
)
}
fn with_active(runtime: Arc<BackendRuntime>, admission: CircuitAdmission, active: ActivePermits) -> Self {
runtime.in_flight.increment(1.0);
Self {
runtime,
admission: Some(admission),
_active: active,
}
}
fn complete(&mut self, failure: Option<ErrorClass>) -> 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<BackendRuntime>,
operation: &'static str,
cancel: &CancellationToken,
deadline: Instant,
) -> std::result::Result<RuntimePermit, (Outcome, KmsError)> {
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<T, F, Fut>(
operation: &'static str,
class: OpClass,
policy: &RetryPolicy,
cancel: &CancellationToken,
attempt: F,
) -> Result<T>
where
F: FnMut() -> Fut,
Fut: Future<Output = std::result::Result<T, AttemptError>>,
{
// 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<T, F, Fut, J>(
operation: &'static str,
class: OpClass,
policy: &RetryPolicy,
cancel: &CancellationToken,
jitter: J,
attempt: F,
) -> Result<T>
where
F: FnMut() -> Fut,
Fut: Future<Output = std::result::Result<T, AttemptError>>,
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<T, F, Fut, J>(
operation: &'static str,
class: OpClass,
policy: &RetryPolicy,
cancel: &CancellationToken,
mut jitter: J,
mut attempt: F,
attempts_made: &mut u32,
) -> (Outcome, Result<T>)
where
F: FnMut() -> Fut,
Fut: Future<Output = std::result::Result<T, AttemptError>>,
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<BackendRuntime> {
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<BackendRuntime> {
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<T> = std::result::Result<T, AttemptError>;
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::<AttemptResult<()>>()
})
.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::<Vec<_>>();
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::<AttemptResult<()>>,
)
.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::<AttemptResult<()>>()
})
.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<Duration> = (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::<serde_json::Value>("{").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<metrics::Unit>,
Option<metrics::SharedString>,
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<Out>(test: impl FnOnce() -> std::pin::Pin<Box<dyn Future<Output = Out>>>) -> (Vec<MetricEntry>, 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<f64> {
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<f64> {
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::<AttemptResult<()>>()
})
.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::<AttemptResult<()>>()
})
.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));
}
}
Reference in New Issue View Git Blame Copy Permalink