feat(kms): bound backend concurrency and failures (#5651)

crates/kms/src/backends/aws.rs

@@ -309,11 +309,11 @@ impl AwsKmsBackend {
             loader = loader.region(aws_sdk_kms::config::Region::new(region.clone()));
         }
         let sdk_config = loader.load().await;
-        if sdk_config.region().is_none() {
+        let Some(region) = sdk_config.region() else {
             return Err(KmsError::configuration_error(
                 "AWS KMS backend could not resolve a region; set the backend region or AWS_REGION",
             ));
-        }
+        };
 
         let mut builder = aws_sdk_kms::config::Builder::from(&sdk_config)
             // `crate::policy` owns retries and timeouts; leaving the SDK's own
@@ -324,13 +324,18 @@ impl AwsKmsBackend {
             builder = builder.endpoint_url(endpoint_url);
         }
 
-        Ok(Self::with_client(aws_sdk_kms::Client::from_conf(builder.build()), &config))
+        Ok(Self::with_client(
+            aws_sdk_kms::Client::from_conf(builder.build()),
+            &config,
+            aws_backend_config.endpoint_url.as_deref().unwrap_or_default(),
+            region.as_ref(),
+        ))
     }
 
-    fn with_client(client: aws_sdk_kms::Client, config: &KmsConfig) -> Self {
+    fn with_client(client: aws_sdk_kms::Client, config: &KmsConfig, endpoint: &str, region: &str) -> Self {
         Self {
             client,
-            retry: RetryPolicy::from_config(config),
+            retry: RetryPolicy::for_backend(config, "aws", endpoint, Some(region), "operations"),
             cancel: CancellationToken::new(),
         }
     }
@@ -826,6 +831,7 @@ mod tests {
     use aws_smithy_types::body::SdkBody;
     use base64::Engine as _;
     use base64::engine::general_purpose::STANDARD as BASE64;
+    use std::sync::atomic::{AtomicU64, Ordering};
 
     /// AWS KMS speaks awsJson1_1; every request goes to `/` on the regional
     /// endpoint, so the replayed request side carries no useful assertion.
@@ -857,6 +863,15 @@ mod tests {
         )
     }
 
+    fn scripted_endpoint() -> String {
+        static NEXT_SCRIPTED_BACKEND_ID: AtomicU64 = AtomicU64::new(0);
+
+        format!(
+            "https://scripted-{}.example.invalid",
+            NEXT_SCRIPTED_BACKEND_ID.fetch_add(1, Ordering::Relaxed)
+        )
+    }
+
     fn scripted_backend(events: Vec<ReplayEvent>) -> (StaticReplayClient, AwsKmsBackend) {
         let http_client = StaticReplayClient::new(events);
         let sdk_config = aws_sdk_kms::Config::builder()
@@ -867,10 +882,20 @@ mod tests {
             .retry_config(aws_sdk_kms::config::retry::RetryConfig::disabled())
             .build();
         let kms_config = KmsConfig::aws(Some("us-east-1".to_string()));
-        let backend = AwsKmsBackend::with_client(aws_sdk_kms::Client::from_conf(sdk_config), &kms_config);
+        let endpoint = scripted_endpoint();
+        let backend = AwsKmsBackend::with_client(aws_sdk_kms::Client::from_conf(sdk_config), &kms_config, &endpoint, "us-east-1");
         (http_client, backend)
     }
 
+    fn aws_config(endpoint: &str, region: &str) -> KmsConfig {
+        let mut config = KmsConfig::aws(Some(region.to_owned()));
+        let BackendConfig::Aws(aws) = &mut config.backend_config else {
+            panic!("AWS constructor must create an AWS backend configuration");
+        };
+        aws.endpoint_url = Some(endpoint.to_owned());
+        config
+    }
+
     fn key_metadata_json(key_id: &str, state: &str) -> serde_json::Value {
         serde_json::json!({
             "KeyMetadata": {
@@ -900,6 +925,27 @@ mod tests {
             .expect("capabilities should deserialize into a flat bool map")
     }
 
+    #[tokio::test]
+    async fn aws_backend_new_identity_includes_endpoint_and_region() {
+        let endpoint = scripted_endpoint();
+        let first = AwsKmsBackend::new(aws_config(&endpoint, "us-east-1"))
+            .await
+            .expect("first AWS backend");
+        let matching = AwsKmsBackend::new(aws_config(&endpoint, "us-east-1"))
+            .await
+            .expect("matching AWS backend");
+        let other_region = AwsKmsBackend::new(aws_config(&endpoint, "us-west-2"))
+            .await
+            .expect("other-region AWS backend");
+        let other_endpoint = AwsKmsBackend::new(aws_config(&scripted_endpoint(), "us-east-1"))
+            .await
+            .expect("other-endpoint AWS backend");
+
+        assert!(first.retry.shares_active_capacity_with(&matching.retry));
+        assert!(!first.retry.shares_active_capacity_with(&other_region.retry));
+        assert!(!first.retry.shares_active_capacity_with(&other_endpoint.retry));
+    }
+
     #[tokio::test]
     async fn aws_backend_capabilities_golden() {
         let (_http, backend) = scripted_backend(Vec::new());
@@ -1023,7 +1069,7 @@ mod tests {
         let mut kms_config = KmsConfig::aws(Some("us-east-1".to_string()));
         kms_config.timeout = std::time::Duration::from_millis(5_000);
         kms_config.retry_attempts = 1;
-        let backend = AwsKmsBackend::with_client(aws_sdk_kms::Client::from_conf(sdk_config), &kms_config);
+        let backend = AwsKmsBackend::with_client(aws_sdk_kms::Client::from_conf(sdk_config), &kms_config, "", "us-east-1");
 
         let error = backend
             .describe_key(DescribeKeyRequest {

crates/kms/src/backends/vault.rs

@@ -227,7 +227,13 @@ impl VaultKmsClient {
             attempt_timeout: kms_config.effective_timeout(),
         };
         let source = token_source_for(&config.auth_method, &settings)?;
-        let policy = VaultCredentialPolicy::from_kms_config(kms_config, &config.auth_method);
+        let policy = VaultCredentialPolicy::from_kms_config(
+            kms_config,
+            &config.auth_method,
+            "vault-kv2",
+            &config.address,
+            config.namespace.as_deref(),
+        );
         let credentials = Arc::new(VaultCredentialProvider::new(settings, source, policy).await?);
 
         info!(address = %config.address, "Vault KMS backend connected");
@@ -237,7 +243,7 @@ impl VaultKmsClient {
             kv_mount: config.kv_mount.clone(),
             key_path_prefix: config.key_path_prefix.clone(),
             dek_crypto: AesDekCrypto::new(),
-            retry: RetryPolicy::from_config(kms_config),
+            retry: RetryPolicy::for_backend(kms_config, "vault-kv2", &config.address, config.namespace.as_deref(), "operations"),
             cancel: CancellationToken::new(),
         })
     }

crates/kms/src/backends/vault_credentials.rs

@@ -520,8 +520,10 @@ impl VaultConnectionSettings {
 /// Refresh and fail-closed tuning for a [`VaultCredentialProvider`].
 #[derive(Debug, Clone)]
 pub(crate) struct VaultCredentialPolicy {
-    /// Retry budget for one login/renewal cycle.
+    /// Retry budget for one login cycle.
     pub(crate) retry: RetryPolicy,
+    /// Retry budget for one token-renewal cycle.
+    pub(crate) renew_retry: RetryPolicy,
     /// Fail-closed margin: once the current token is within this window of
     /// expiry without a successful refresh, [`VaultCredentialProvider::current`]
     /// refuses to hand it out.
@@ -536,8 +538,14 @@ impl VaultCredentialPolicy {
     /// The default safety window equals the per-attempt timeout: a request
     /// issued now can stay in flight for up to one attempt timeout, so the
     /// token must outlive at least that.
-    pub(crate) fn from_kms_config(config: &KmsConfig, auth_method: &VaultAuthMethod) -> Self {
-        let retry = RetryPolicy::from_config(config);
+    pub(crate) fn from_kms_config(
+        config: &KmsConfig,
+        auth_method: &VaultAuthMethod,
+        backend: &'static str,
+        endpoint: &str,
+        namespace: Option<&str>,
+    ) -> Self {
+        let retry = RetryPolicy::for_credentials(config, backend, endpoint, namespace, "credentials-login");
         let safety_window = match auth_method {
             VaultAuthMethod::AppRole {
                 refresh_safety_window_secs: Some(secs),
@@ -551,6 +559,7 @@ impl VaultCredentialPolicy {
         };
         Self {
             retry,
+            renew_retry: RetryPolicy::for_credentials(config, backend, endpoint, namespace, "credentials-renew"),
             safety_window,
             retry_interval: DEFAULT_REFRESH_RETRY_INTERVAL,
         }
@@ -717,7 +726,7 @@ impl VaultCredentialProvider {
 
         let renewable = current.lease.map(|lease| lease.renewable).unwrap_or(false);
         let renewed = if renewable {
-            match policy::execute("vault_token_renew", OpClass::Auth, &self.policy.retry, cancel, || {
+            match policy::execute("vault_token_renew", OpClass::Auth, &self.policy.renew_retry, cancel, || {
                 self.source.renew(&current.client)
             })
             .await
@@ -868,19 +877,56 @@ mod tests {
     /// Tight retry budget so paused-clock tests stay deterministic: one
     /// attempt per cycle, failed cycles spaced by `retry_interval`.
     fn test_policy(safety_window: Duration, retry_interval: Duration) -> VaultCredentialPolicy {
+        let retry = RetryPolicy::for_test(
+            Duration::from_secs(1),
+            Duration::from_secs(1),
+            1,
+            Duration::from_millis(10),
+            Duration::from_millis(10),
+        );
+        let renew_retry = RetryPolicy::for_test(
+            Duration::from_secs(1),
+            Duration::from_secs(1),
+            1,
+            Duration::from_millis(10),
+            Duration::from_millis(10),
+        );
         VaultCredentialPolicy {
-            retry: RetryPolicy {
-                attempt_timeout: Duration::from_secs(1),
-                op_deadline: Duration::from_secs(1),
-                max_attempts: 1,
-                base_backoff: Duration::from_millis(10),
-                max_backoff: Duration::from_millis(10),
-            },
+            retry,
+            renew_retry,
             safety_window,
             retry_interval,
         }
     }
 
+    #[test]
+    fn configured_login_and_renewal_use_reserved_credential_capacity() {
+        let config = KmsConfig::default();
+        let auth_method = VaultAuthMethod::Token {
+            token: TEST_TOKEN.to_string(),
+        };
+        let credentials = VaultCredentialPolicy::from_kms_config(
+            &config,
+            &auth_method,
+            "vault-kv2",
+            "https://credential-policy.example.invalid",
+            Some("team-namespace"),
+        );
+        let operations = RetryPolicy::for_backend(
+            &config,
+            "vault-kv2",
+            "https://credential-policy.example.invalid",
+            Some("team-namespace"),
+            "operations",
+        );
+
+        assert!(credentials.retry.uses_credential_reserve());
+        assert!(credentials.renew_retry.uses_credential_reserve());
+        assert!(!operations.uses_credential_reserve());
+        assert!(credentials.retry.shares_active_capacity_with(&credentials.renew_retry));
+        assert!(credentials.retry.shares_active_capacity_with(&operations));
+    }
+
     /// Shared observable state of a [`ScriptedSource`].
     #[derive(Debug, Default)]
     struct ScriptedState {
@@ -1084,11 +1130,12 @@ mod tests {
             "expected CredentialsUnavailable, got {error:?}"
         );
 
-        // Recovery: the next retry cycle succeeds, installs a fresh
-        // generation, and the provider serves requests again.
+        // Recovery: after the failed-refresh circuit's cool-down, its single
+        // half-open probe succeeds, installs a fresh generation, and the
+        // provider serves requests again.
         state.fail_renew.store(false, Ordering::SeqCst);
         state.fail_login.store(false, Ordering::SeqCst);
-        tokio::time::sleep(Duration::from_secs(6)).await;
+        tokio::time::sleep(Duration::from_secs(31)).await;
         let handle = provider.current().expect("provider must recover after a successful refresh");
         assert!(handle.generation >= 1);
 

crates/kms/src/backends/vault_transit.rs

@@ -230,8 +230,16 @@ impl VaultTransitKmsClient {
             attempt_timeout: kms_config.effective_timeout(),
         };
         let source = token_source_for(&config.auth_method, &settings)?;
-        let policy = VaultCredentialPolicy::from_kms_config(kms_config, &config.auth_method);
+        let policy = VaultCredentialPolicy::from_kms_config(
+            kms_config,
+            &config.auth_method,
+            "vault-transit",
+            &config.address,
+            config.namespace.as_deref(),
+        );
         let credentials = Arc::new(VaultCredentialProvider::new(settings, source, policy).await?);
+        let retry =
+            RetryPolicy::for_backend(kms_config, "vault-transit", &config.address, config.namespace.as_deref(), "operations");
 
         Ok(Self {
             credentials,
@@ -242,7 +250,7 @@ impl VaultTransitKmsClient {
                 .max_capacity(METADATA_CACHE_CAPACITY)
                 .time_to_live(METADATA_CACHE_TTL)
                 .build(),
-            retry: RetryPolicy::from_config(kms_config),
+            retry,
             cancel: CancellationToken::new(),
         })
     }

crates/kms/src/backup/vault_restore.rs

@@ -452,12 +452,24 @@ impl VaultRestoreClient {
             attempt_timeout: kms_config.effective_timeout(),
         };
         let source = token_source_for(&target.auth_method, &settings)?;
-        let policy = VaultCredentialPolicy::from_kms_config(kms_config, &target.auth_method);
+        let policy = VaultCredentialPolicy::from_kms_config(
+            kms_config,
+            &target.auth_method,
+            "vault-restore",
+            &target.address,
+            target.namespace.as_deref(),
+        );
         let credentials = Arc::new(VaultCredentialProvider::new(settings, source, policy).await?);
         Ok(Self {
             credentials,
             kv_mount: target.kv_mount.clone(),
-            retry: RetryPolicy::from_config(kms_config),
+            retry: RetryPolicy::for_backend(
+                kms_config,
+                "vault-restore",
+                &target.address,
+                target.namespace.as_deref(),
+                "operations",
+            ),
             cancel: CancellationToken::new(),
         })
     }