Cache Azure user delegation key to avoid GetUserDelegationKey throttling

The Azure SAS signer fetched a fresh user delegation key (a
GetUserDelegationKey network round-trip) on every signing request.
Under load Azure throttles these with HTTP 503, which surfaced as
QueryDataset Internal errors in redap.

Cache a long-lived (6h) user delegation key on AzureClient and reuse
it to mint many short-lived SAS tokens, refreshing only when it no
longer covers the requested window.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: emilk

src/azure/client.rs

@@ -46,7 +46,7 @@ use http::{
 use rand::RngExt;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use url::Url;
 
@@ -662,16 +662,54 @@ async fn parse_blob_batch_delete_body(
     Ok(results)
 }
 
+/// How long a freshly fetched user delegation key is requested to remain valid.
+///
+/// The SAS tokens we sign with it stay short-lived; this only bounds how often
+/// we call `GetUserDelegationKey`. Azure caps the key lifetime at 7 days.
+const USER_DELEGATION_KEY_VALIDITY: Duration = Duration::from_secs(6 * 60 * 60);
+
+/// Refresh margin so we rotate a cached key before it actually lapses, and to
+/// absorb clock skew between us and Azure.
+const USER_DELEGATION_KEY_REFRESH_MARGIN: Duration = Duration::from_secs(5 * 60);
+
+/// A user delegation key plus the signing key derived from it, cached so we can
+/// sign many SAS URLs without re-fetching the key on every request.
+#[derive(Debug, Clone)]
+struct CachedDelegationKey {
+    signing_key: AzureAccessKey,
+    key: UserDelegationKey,
+    valid_from: DateTime<Utc>,
+    valid_until: DateTime<Utc>,
+}
+
+impl CachedDelegationKey {
+    /// Whether this key is valid for the entire `[start, expiry]` SAS window,
+    /// keeping a refresh margin so we rotate before it lapses.
+    fn covers(&self, start: DateTime<Utc>, expiry: DateTime<Utc>) -> bool {
+        self.valid_from <= start && expiry + USER_DELEGATION_KEY_REFRESH_MARGIN <= self.valid_until
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct AzureClient {
     config: AzureConfig,
     client: HttpClient,
+    /// Cached user delegation key used to sign SAS URLs.
+    ///
+    /// Fetching a key is a network round-trip (`GetUserDelegationKey`), so we
+    /// request a long-lived key once and reuse it to mint many short-lived SAS
+    /// tokens, refreshing only when it no longer covers the requested window.
+    delegation_key_cache: Mutex<Option<CachedDelegationKey>>,
 }
 
 impl AzureClient {
     /// create a new instance of [AzureClient]
     pub(crate) fn new(config: AzureConfig, client: HttpClient) -> Self {
-        Self { config, client }
+        Self {
+            config,
+            client,
+            delegation_key_cache: Mutex::new(None),
+        }
     }
 
     /// Returns the config
@@ -1019,16 +1057,15 @@ impl AzureClient {
         let signed_expiry = signed_start + expires_in;
         match credential.as_deref() {
             Some(AzureCredential::BearerToken(_)) => {
-                let key = self
-                    .get_user_delegation_key(&signed_start, &signed_expiry)
+                let cached = self
+                    .cached_user_delegation_key(signed_start, signed_expiry)
                     .await?;
-                let signing_key = AzureAccessKey::try_new(&key.value)?;
                 Ok(AzureSigner::new(
-                    signing_key,
+                    cached.signing_key,
                     self.config.account.clone(),
                     signed_start,
                     signed_expiry,
-                    Some(key),
+                    Some(cached.key),
                 ))
             }
             Some(AzureCredential::AccessKey(key)) => Ok(AzureSigner::new(
@@ -1043,6 +1080,51 @@ impl AzureClient {
         }
     }
 
+    /// Return a user delegation key that is valid for the whole `[start, expiry]`
+    /// SAS window, reusing the cached key when possible.
+    ///
+    /// `GetUserDelegationKey` is a network round-trip and Azure throttles it
+    /// (HTTP 503) under load, so we fetch a longer-lived key once and reuse it
+    /// to sign many short-lived SAS URLs.
+    async fn cached_user_delegation_key(
+        &self,
+        start: DateTime<Utc>,
+        expiry: DateTime<Utc>,
+    ) -> Result<CachedDelegationKey> {
+        // Fast path: a cached key that still covers the requested window.
+        if let Some(cached) = self.delegation_key_cache.lock().unwrap().clone() {
+            if cached.covers(start, expiry) {
+                return Ok(cached);
+            }
+        }
+
+        // Slow path: fetch a fresh, long-lived key. We may race other callers
+        // here and fetch more than once, but that is rare and self-correcting.
+        let key_start = start;
+        let key_expiry = expiry.max(start + USER_DELEGATION_KEY_VALIDITY);
+        let key = self
+            .get_user_delegation_key(&key_start, &key_expiry)
+            .await?;
+        let signing_key = AzureAccessKey::try_new(&key.value)?;
+
+        // Trust the validity window Azure actually granted (it may clamp it).
+        let valid_from = DateTime::parse_from_rfc3339(&key.signed_start)
+            .map(|t| t.with_timezone(&Utc))
+            .unwrap_or(key_start);
+        let valid_until = DateTime::parse_from_rfc3339(&key.signed_expiry)
+            .map(|t| t.with_timezone(&Utc))
+            .unwrap_or(key_expiry);
+
+        let cached = CachedDelegationKey {
+            signing_key,
+            key,
+            valid_from,
+            valid_until,
+        };
+        *self.delegation_key_cache.lock().unwrap() = Some(cached.clone());
+        Ok(cached)
+    }
+
     #[cfg(test)]
     pub(crate) async fn get_blob_tagging(&self, path: &Path) -> Result<HttpResponse> {
         let credential = self.get_credential().await?;
@@ -1385,7 +1467,7 @@ impl BlockList {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Deserialize)]
+#[derive(Debug, Clone, Default, PartialEq, Deserialize)]
 #[serde(rename_all = "PascalCase")]
 pub(crate) struct UserDelegationKey {
     pub signed_oid: String,
@@ -1594,6 +1676,34 @@ mod tests {
             quick_xml::de::from_str(S).unwrap();
     }
 
+    #[test]
+    fn test_cached_delegation_key_covers() {
+        let cached = CachedDelegationKey {
+            signing_key: AzureAccessKey::try_new("Zm9vYmFy").unwrap(),
+            key: UserDelegationKey::default(),
+            valid_from: DateTime::parse_from_rfc3339("2026-06-25T00:00:00Z")
+                .unwrap()
+                .with_timezone(&Utc),
+            valid_until: DateTime::parse_from_rfc3339("2026-06-25T06:00:00Z")
+                .unwrap()
+                .with_timezone(&Utc),
+        };
+
+        let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
+
+        // SAS window comfortably inside the key validity: reuse.
+        assert!(cached.covers(at("2026-06-25T01:00:00Z"), at("2026-06-25T02:00:00Z")));
+
+        // SAS starts before the key is valid: refetch.
+        assert!(!cached.covers(at("2026-06-24T23:59:00Z"), at("2026-06-25T00:30:00Z")));
+
+        // SAS expiry past the key expiry: refetch.
+        assert!(!cached.covers(at("2026-06-25T05:00:00Z"), at("2026-06-25T06:30:00Z")));
+
+        // SAS expiry within the refresh margin of the key expiry: refetch.
+        assert!(!cached.covers(at("2026-06-25T05:00:00Z"), at("2026-06-25T05:58:00Z")));
+    }
+
     #[cfg(feature = "reqwest")]
     #[tokio::test]
     async fn test_build_bulk_delete_body() {