mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-13 00:26:53 +00:00
Zhengchao An
2109 lines
82 KiB
Rust
2109 lines
82 KiB
Rust
// 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 {
|
|
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 {
|
|
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 `operation`, `op_class`, and `outcome`.
|
|
const METRIC_OPERATIONS_TOTAL: &str = "rustfs_kms_backend_operations_total";
|
|
/// Counter: failed attempts, by `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 `operation` and `outcome`.
|
|
const METRIC_OPERATION_DURATION_SECONDS: &str = "rustfs_kms_backend_operation_duration_seconds";
|
|
/// Histogram: attempts one operation used before completing, by `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 operation, operation class, and outcome"
|
|
);
|
|
metrics::describe_counter!(
|
|
METRIC_ATTEMPT_FAILURES_TOTAL,
|
|
"Total failed KMS backend attempts, by operation and retry classification"
|
|
);
|
|
metrics::describe_histogram!(
|
|
METRIC_OPERATION_DURATION_SECONDS,
|
|
"Wall-clock duration of KMS backend operations including retries and backoff, in seconds"
|
|
);
|
|
metrics::describe_histogram!(
|
|
METRIC_OPERATION_ATTEMPTS,
|
|
"Number of attempts a KMS backend operation used before completing"
|
|
);
|
|
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(operation: &'static str, error_class: &'static str) {
|
|
metrics::counter!(
|
|
METRIC_ATTEMPT_FAILURES_TOTAL,
|
|
"operation" => operation,
|
|
"error_class" => error_class
|
|
)
|
|
.increment(1);
|
|
}
|
|
|
|
/// Record the terminal outcome of one policy execution.
|
|
fn record_operation(operation: &'static str, class: OpClass, outcome: Outcome, attempts: u32, elapsed: Duration) {
|
|
metrics::counter!(
|
|
METRIC_OPERATIONS_TOTAL,
|
|
"operation" => operation,
|
|
"op_class" => class.as_label(),
|
|
"outcome" => outcome.as_label()
|
|
)
|
|
.increment(1);
|
|
metrics::histogram!(
|
|
METRIC_OPERATION_DURATION_SECONDS,
|
|
"operation" => operation,
|
|
"outcome" => outcome.as_label()
|
|
)
|
|
.record(elapsed.as_secs_f64());
|
|
metrics::histogram!(
|
|
METRIC_OPERATION_ATTEMPTS,
|
|
"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(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(operation, failure.class.as_label());
|
|
failure
|
|
}
|
|
Err(_) => {
|
|
record_attempt_failure(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> {
|
|
let capacity = Arc::new(BackendCapacity {
|
|
total: Arc::new(Semaphore::new(max_concurrent)),
|
|
operations: Arc::new(Semaphore::new(max_concurrent)),
|
|
});
|
|
let mut runtime = BackendRuntime::new("test-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:?}");
|
|
}
|
|
|
|
#[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));
|
|
}
|
|
}
|