diff --git a/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md b/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md index 1640dcabc0d..425269dade4 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md +++ b/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md @@ -19,6 +19,7 @@ - A partition stolen by a higher-or-equal-epoch attacher now surfaces as `ErrorKind::ConsumerDisconnected` when the broker reports `amqp:link:stolen` on a re-attach, not only on an in-flight receive. Other attach failures inside the receive loop now classify by their own kind. The wrapper reported all of them as a message error, which the retry decider treated as non-retryable. - The `EventProcessor`'s load-balancer reconciliation now closes the underlying AMQP receiver for any partition that has been reassigned to another consumer, so the consumer's `stream_events()` resolves and the loop can terminate. Previously a stolen partition's client could continue to attempt receives until the broker tore down the link. - Fixed a deadlock when a CBS failure during management-client creation started connection recovery. ([#4728](https://github.com/Azure/azure-sdk-for-rust/issues/4728)) +- Closed a stale-resource window in connection recovery. A `ReconnectConnection` recovery that fired while a slow-path attach (authorize, session begin, or sender/receiver link attach) was in flight could cache a resource bound to the just-dropped connection; the next operation on that resource failed (unauthorized / detached / closed) and triggered a second, redundant recovery cycle. A recovery generation counter now tags each cached resource, and a slow path that completes across a recovery discards its result and re-attaches against the new connection instead of caching the stale one. The authorizer's token cache is mutable (a background task refreshes tokens) so it cannot use the same one-shot cell as the connection caches; both of its writers, `authorize_path` and the refresh task, instead re-check the generation under the same lock that recovery's clear takes, and a recovery brackets its invalidation with a generation bump on each side, which leaves the counter odd for as long as the recovery runs, so a slow path that overlaps a recovery at either end also discards rather than caching a resource bound to the connection that recovery is dropping. A token refresh pass that a recovery discards now applies the same backoff floor as a failed pass, so a recovery storm cannot turn the refresh loop into an uncapped stream of credential and CBS calls. The per-path / per-partition concurrency is preserved. ([#4454](https://github.com/Azure/azure-sdk-for-rust/issues/4454)) ### Other Changes diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs index a4844c10a72..ed9e6ad31c4 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs @@ -3,7 +3,10 @@ // cspell:ignore sastoken refreshable -use crate::{common::recoverable::RecoverableConnection, error::Result}; +use crate::{ + common::recoverable::{RecoverableConnection, MAX_GENERATION_RETRIES}, + error::Result, +}; use async_lock::RwLock; use azure_core::{ async_runtime::{get_async_runtime, SpawnedTask}, @@ -24,11 +27,13 @@ const TOKEN_REFRESH_BIAS: Duration = Duration::minutes(6); // By default, we ref const TOKEN_REFRESH_JITTER_MIN: Duration = Duration::seconds(-5); // Minimum jitter (added from the bias, so a negative number means we refresh before the bias) const TOKEN_REFRESH_JITTER_MAX: Duration = Duration::seconds(5); // Maximum jitter (added to the bias) -// Floor delay applied after a pass in which one or more token refreshes failed. -// A failed refresh leaves the token's `expires_on` unchanged, so its computed -// refresh_time stays in the past and the next pass would not sleep. Without this -// floor a persistent failure (credential outage, CBS down) busy-spins, hammering -// get_token / perform_authorization with no backoff. See PR #4593 review. +// Floor delay applied after a pass that did not write any fresh token back: +// either one or more refreshes failed, or a recovery made the pass discard its +// results. Both outcomes leave the cached token's `expires_on` unchanged, so its +// computed refresh_time stays in the past and the next pass would not sleep. +// Without this floor a persistent failure (credential outage, CBS down) or a +// recovery storm busy-spins, hammering get_token / perform_authorization with no +// backoff. See PR #4593 review and the `Discarded` arm below (#4454). const TOKEN_REFRESH_RETRY_BACKOFF: Duration = Duration::seconds(30); const EVENTHUBS_AUTHORIZATION_SCOPE: &str = "https://eventhubs.azure.net/.default"; @@ -50,6 +55,22 @@ impl Default for TokenRefreshTimes { } } +/// What one [`Authorizer::refresh_due_tokens`] pass concluded. +#[derive(Debug, PartialEq, Eq)] +enum RefreshPass { + /// The pass ran to the end. `failed` is true when at least one path did not + /// refresh, so the caller must back off before the next pass. + Completed { failed: bool }, + /// A recovery advanced the generation while the pass was in flight, so the + /// pass dropped the tokens it had refreshed (#4454). This is a distinct + /// outcome from `Completed { failed: true }`, because nothing failed, but the + /// caller must still back off: the cache keeps the old tokens, so they stay + /// due and the next pass would not sleep. + Discarded, + /// The recoverable connection is gone, so the refresh task must stop. + Stop, +} + pub(crate) struct Authorizer { authorization_scopes: RwLock>, authorization_refresher: OnceLock, @@ -123,49 +144,95 @@ impl Authorizer { ) -> azure_core_amqp::Result { debug!("Authorizing path: {path}"); - // Fast path: cached token under a brief lock. - if let Some(token) = self.authorization_scopes.read().await.get(path).cloned() { - debug!("Token already exists for path: {path}"); - return Ok(token); - } - - // Slow path: fetch the credential and perform the CBS attach *without* - // holding the scope cache lock. Holding it across `perform_authorization` - // would block `clear()` (called from `recover_from_error`) for as long as - // the CBS attach is in flight; if that CBS attach is itself the operation - // that triggers recovery, the result is a self-deadlock. Matches the - // pattern used by `ensure_sender` / `ensure_receiver` / `get_session` in - // `RecoverableConnection`. - debug!("Creating new authorization scope for path: {path}"); - - debug!("Get Token."); - let token = self - .credential - .get_token(&[EVENTHUBS_AUTHORIZATION_SCOPE], None) - .await - .map_err(AmqpError::from)?; - - debug!("Token for path {path} expires at {}", token.expires_on); + // #4454 stale-token guard. The token cache is mutable (the + // refresh task rewrites entries), so unlike the connection caches it can't + // use a `OnceCell`; the generation check is applied here directly. We + // capture the connection's recovery generation before the lock-free CBS + // attach and re-check it after: if a recovery cleared this cache and bumped + // the generation mid-attach, the token we just authorized is bound to the + // torn-down connection's CBS link, so we discard it and re-authorize + // against the new connection instead of caching a stale entry (which the + // next operation would otherwise use and fail on, costing a second recovery + // cycle). `MAX_GENERATION_RETRIES` bounds the loop, so a storm of + // back-to-back recoveries surfaces an error rather than spinning forever. + for _ in 0..MAX_GENERATION_RETRIES { + // Fast path: cached token under a brief lock. + if let Some(token) = self.authorization_scopes.read().await.get(path).cloned() { + debug!("Token already exists for path: {path}"); + return Ok(token); + } - self.perform_authorization(connection, path, &token).await?; - debug!("Token verified."); + let generation = connection.generation(); + + // Slow path: fetch the credential and perform the CBS attach *without* + // holding the scope cache lock. Holding it across `perform_authorization` + // would block `clear()` (called from `recover_from_error`) for as long as + // the CBS attach is in flight; if that CBS attach is itself the operation + // that triggers recovery, the result is a self-deadlock. Matches the + // pattern used by `ensure_sender` / `ensure_receiver` / `get_session` in + // `RecoverableConnection`. + debug!("Creating new authorization scope for path: {path}"); + + debug!("Get Token."); + let token = self + .credential + .get_token(&[EVENTHUBS_AUTHORIZATION_SCOPE], None) + .await + .map_err(AmqpError::from)?; + + debug!("Token for path {path} expires at {}", token.expires_on); + + self.perform_authorization(connection, path, &token).await?; + debug!("Token verified."); + + // Insert under the write lock, but re-check the recovery generation + // *inside* that lock before inserting. `clear()` (from + // `recover_from_error`) takes this same lock and runs after the + // generation bump, so re-reading the generation here, rather than before + // acquiring the lock, closes the window where a recovery lands between + // the check and the insert: we either observe the bump and discard, or + // we hold the lock across the insert so no recovery can interleave. If a + // recovery raced the lock-free attach above, the CBS link we authorized + // against is gone, so we drop this token and retry against the new + // generation rather than repopulating the just-cleared cache with a + // stale entry (which the next operation would use and fail on, costing a + // second recovery cycle). See #4454. + let stored = { + let mut scopes = self.authorization_scopes.write().await; + if !connection.generation_is_current(generation) { + None + } else { + // If another task won the race, return its cached token and drop + // ours. Both CBS auths succeeded against the same link, so either + // credential is acceptable to the broker. + Some(scopes.entry(path.clone()).or_insert(token).clone()) + } + }; + let Some(stored) = stored else { + debug!( + "Discarding token authorized during recovery (#4454) for path: {path}; re-authorizing." + ); + continue; + }; - // Insert; if another task won the race, return its cached token and drop - // ours. Both CBS auths succeeded against the same link, so either - // credential is acceptable to the broker. - let stored = { - let mut scopes = self.authorization_scopes.write().await; - scopes.entry(path.clone()).or_insert(token).clone() - }; + self.authorization_refresher.get_or_init(|| { + debug!("Starting authorization refresh task."); + let self_clone = self.clone(); + let async_runtime = get_async_runtime(); + async_runtime.spawn(Box::pin(self_clone.refresh_tokens_task())) + }); - self.authorization_refresher.get_or_init(|| { - debug!("Starting authorization refresh task."); - let self_clone = self.clone(); - let async_runtime = get_async_runtime(); - async_runtime.spawn(Box::pin(self_clone.refresh_tokens_task())) - }); + return Ok(stored); + } - Ok(stored) + // Intentionally a plain `AmqpError::with_message`: the connection's + // `should_retry_amqp_error` classifies this unrecognized kind as + // `ReturnError`, so exhausting the budget surfaces to the caller instead of + // looping. Do not "fix" this into a retryable kind, that would let a recovery + // storm spin here forever (#4454). + Err(AmqpError::with_message(format!( + "Exceeded retry budget ({MAX_GENERATION_RETRIES}) authorizing path '{path}' across recoveries" + ))) } /// Actually perform an authorization against the Event Hubs service. @@ -344,35 +411,128 @@ impl Authorizer { debug!("Not sleeping because refresh time ({refresh_time}) is in the past (now = {now})."); } - // Refresh the tokens. - // First, collect the tokens that need refreshing while holding the lock briefly - let tokens_to_refresh = { - let scopes = self.authorization_scopes.read().await; - let mut to_refresh = Vec::new(); - for (url, token) in scopes.iter() { - if non_refreshable.contains(url) { - continue; - } - if token.expires_on >= now + (token_refresh_bias) { - debug!( - "Token not expired for {url}: ExpiresOn: {}, Now: {now}, Bias: {token_refresh_bias:?}", - token.expires_on - ); - continue; - } + // Refresh every token that is due as of `now`, then write the fresh + // tokens back, guarded against a racing recovery (#4454). + match self + .refresh_due_tokens(now, token_refresh_bias, &mut non_refreshable) + .await? + { + RefreshPass::Stop => return Ok(()), + RefreshPass::Completed { failed: false } => {} + RefreshPass::Completed { failed: true } => { + // The failed path keeps its old `expires_on`, so its refresh + // time is still in the past and the top-of-loop sleep would be + // skipped. Back off for a bounded interval before the next pass + // so a persistent failure does not busy-spin on get_token / + // perform_authorization. + warn!( + backoff = ?TOKEN_REFRESH_RETRY_BACKOFF, + "One or more token refreshes failed this pass; backing off before retrying." + ); + azure_core::sleep::sleep(TOKEN_REFRESH_RETRY_BACKOFF).await; + } + RefreshPass::Discarded => { + // A recovery advanced the generation mid-pass, so the pass + // dropped its refreshed tokens. Nothing failed, but the cache + // keeps the old tokens, so they stay due and the top-of-loop + // sleep would be skipped exactly as it is after a failed pass. + // Apply the same floor. `ReconnectSession` and `ReconnectLink` + // advance the generation and leave the token cache populated, so + // without this floor a recovery storm turns this loop into an + // uncapped stream of get_token and CBS calls (#4454). + warn!( + backoff = ?TOKEN_REFRESH_RETRY_BACKOFF, + "A recovery discarded the tokens refreshed this pass; backing off before retrying." + ); + azure_core::sleep::sleep(TOKEN_REFRESH_RETRY_BACKOFF).await; + } + } + } + } + /// One refresh pass: re-authorize every cached token that is within + /// `token_refresh_bias` of expiring as of `now`, then write the fresh tokens + /// back, guarded against a racing recovery. + /// + /// #4454: the refresh task is the token cache's second writer (alongside + /// `authorize_path`), so it needs the same generation guard. We capture the + /// recovery generation before the lock-free CBS re-authorizations below and + /// re-check it under the write lock before writing the refreshed tokens back. + /// If a recovery clears the token cache and bumps the generation mid-refresh, + /// these tokens are bound to the torn-down connection; writing them back would + /// repopulate the just-cleared cache with stale entries that the next operation + /// would use and fail on. On a mismatch we drop them and let the next + /// `authorize_path` re-establish fresh tokens against the new connection. A + /// mismatch returns [`RefreshPass::Discarded`], so the caller applies the same + /// backoff floor it applies to a failed pass; the cache keeps the old tokens, + /// so they stay due and the next pass would otherwise start with no delay. + /// + /// `non_refreshable` is owned by the caller's loop and carries across passes: + /// a path that lands in it is skipped by every later pass. + /// + /// Extracted from the `refresh_tokens` loop (which owns the expiry scheduling + /// and sleeping) so the generation guard can be exercised deterministically in + /// tests. + async fn refresh_due_tokens( + self: &Arc, + now: OffsetDateTime, + token_refresh_bias: Duration, + non_refreshable: &mut HashSet, + ) -> Result { + // First, collect the tokens that need refreshing while holding the lock briefly + let tokens_to_refresh = { + let scopes = self.authorization_scopes.read().await; + let mut to_refresh = Vec::new(); + for (url, token) in scopes.iter() { + if non_refreshable.contains(url) { + continue; + } + if token.expires_on >= now + (token_refresh_bias) { debug!( - "Token about to be expired for {url}: ExpiresOn: {}, Now: {now}, Bias: {token_refresh_bias:?}", + "Token not expired for {url}: ExpiresOn: {}, Now: {now}, Bias: {token_refresh_bias:?}", token.expires_on ); - // Carry the current expiry so a refresh that does not extend - // it can be detected as non-refreshable below. - to_refresh.push((url.clone(), token.expires_on)); + continue; } - to_refresh + + debug!( + "Token about to be expired for {url}: ExpiresOn: {}, Now: {now}, Bias: {token_refresh_bias:?}", + token.expires_on + ); + // Carry the current expiry so a refresh that does not extend + // it can be detected as non-refreshable below. + to_refresh.push((url.clone(), token.expires_on)); + } + to_refresh + }; + + // Nothing due: skip the connection upgrade and lock dance entirely. Scoping + // the work inside this branch keeps the connection and the recovery + // generation as plain values that exist only where they are valid, so they + // cannot drift out of sync (no hand-maintained `Option` invariant, no + // `expect()` that a future edit could turn into a panic in this background + // task and silently stop all token refresh). + // + // `refresh_failed` records whether any path failed this pass, so the caller + // can back off instead of retrying with no delay. + let mut refresh_failed = false; + if !tokens_to_refresh.is_empty() { + // A failed upgrade is terminal, not retryable: once the last strong + // reference to the RecoverableConnection is dropped the Weak can never + // upgrade again, so continuing would loop forever with nothing left to + // refresh against. Stop the task. + let Some(connection) = self.recoverable_connection.upgrade() else { + info!( + operation = "upgrade_connection", + "Recoverable connection has been dropped; stopping token refresher." + ); + return Ok(RefreshPass::Stop); }; + // Capture the recovery generation before the lock-free re-authorizations + // below, so a recovery that races them is detected before write-back (#4454). + let captured = connection.generation(); - // Now refresh tokens without holding the lock to avoid deadlocks. + // Refresh tokens without holding the scopes lock to avoid deadlocks. // // A failure to refresh a single path (transient get_token failure, // authorization failure, etc.) must NOT tear down the refresh task: @@ -381,11 +541,6 @@ impl Authorizer { // and continue to the next path; the unrefreshed token will be // retried on the next pass. let mut updated_tokens = HashMap::new(); - // Tracks whether any path failed to refresh this pass. A failure - // leaves that token's `expires_on` unchanged, so refresh_time stays - // in the past and the next pass would not sleep; we back off below to - // avoid a no-delay retry storm. - let mut refresh_failed = false; for (url, previous_expiry) in tokens_to_refresh { let new_token = match self .credential @@ -418,22 +573,6 @@ impl Authorizer { continue; } - // Create an ephemeral connection to host the authentication. - // - // A failed upgrade is terminal, not retryable: once the last - // strong reference to the RecoverableConnection is dropped the - // Weak can never upgrade again, so continuing here would loop - // forever with nothing left to refresh against. Stop the task. - let connection = match self.recoverable_connection.upgrade() { - Some(connection) => connection, - None => { - info!( - operation = "upgrade_connection", - "Recoverable connection has been dropped; stopping token refresher." - ); - return Ok(()); - } - }; if let Err(e) = self .perform_authorization(&connection, &url, &new_token) .await @@ -455,27 +594,41 @@ impl Authorizer { updated_tokens.insert(url.clone(), new_token); } - // Finally, update the scopes map with the new tokens + // Finally, update the scopes map with the new tokens, unless a recovery + // raced us. Re-check the generation under the same write lock `clear()` + // takes (#4454) before writing anything back. if !updated_tokens.is_empty() { let mut scopes = self.authorization_scopes.write().await; + if !connection.generation_is_current(captured) { + debug!( + "Discarding tokens refreshed during recovery (#4454); a recovery overlapped the refresh." + ); + // Report the discard to the caller so it applies the backoff + // floor. The cache keeps the old tokens, so they are still due + // and the next pass would start with no delay (#4454). + return Ok(RefreshPass::Discarded); + } for (url, token) in updated_tokens.into_iter() { scopes.insert(url.clone(), token); } debug!("Updated tokens."); } - - // If any path failed to refresh, its refresh_time is still in the - // past, so the top-of-loop sleep would be skipped. Back off for a - // bounded interval before retrying so a persistent failure does not - // busy-spin on get_token / perform_authorization. - if refresh_failed { - warn!( - backoff = ?TOKEN_REFRESH_RETRY_BACKOFF, - "One or more token refreshes failed this pass; backing off before retrying." - ); - azure_core::sleep::sleep(TOKEN_REFRESH_RETRY_BACKOFF).await; - } } + + Ok(RefreshPass::Completed { + failed: refresh_failed, + }) + } + + /// Test hook: hold the token cache's write lock. A recovery that clears the + /// authorizer blocks inside [`Authorizer::clear`] until the guard is dropped, + /// which gives a test a deterministic point part way through + /// `apply_recovery_plan`. See `recovery_generation_differs_for_a_mid_recovery_capture`. + #[cfg(test)] + pub(crate) async fn lock_scopes_for_test( + &self, + ) -> async_lock::RwLockWriteGuard<'_, HashMap> { + self.authorization_scopes.write().await } #[cfg(test)] @@ -992,4 +1145,287 @@ mod tests { .expect("authorize_path task panicked") .expect("authorize_path returned an error"); } + + // #4454: when a recovery races an in-flight `authorize_path` slow path, the + // token authorized against the now-dead CBS link must be discarded and the + // path re-authorized against the new connection, rather than cached and handed + // out stale. + // + // The token cache is mutable (the refresh task rewrites it), so it cannot use + // an `OnceCell` like the connection caches; `authorize_path` guards itself with + // the connection's recovery generation instead. This test drives that guard + // deterministically: a gated credential blocks the first `get_token` so the + // test can fire a simulated reconnect (which bumps the generation) precisely + // during the lock-free authorization window. The first attempt's token must be + // thrown away and a second authorization performed; the cached token must be + // the second one, stamped at the post-recovery generation. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn authorize_path_discards_token_authorized_during_recovery() { + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + #[derive(Debug)] + struct CountingGatedCredential { + calls: AtomicUsize, + first_call_entered: AtomicBool, + release_first_call: AtomicBool, + } + + #[async_trait::async_trait] + impl TokenCredential for CountingGatedCredential { + async fn get_token( + &self, + _scopes: &[&str], + _options: Option>, + ) -> azure_core::Result { + let call = self.calls.fetch_add(1, Ordering::SeqCst); + // Gate only the first call so the test can interleave a recovery + // while the slow path is mid-authorization. Later calls (the + // re-authorization) proceed immediately. + if call == 0 { + self.first_call_entered.store(true, Ordering::SeqCst); + while !self.release_first_call.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + } + // Far-future expiry so the spawned refresh task sleeps on its first + // pass instead of immediately re-fetching (which would add a third, + // racy `get_token` call and make the exact-count assert flaky). + Ok(AccessToken::new( + azure_core::credentials::Secret::new("mock_token"), + OffsetDateTime::now_utc() + Duration::hours(1), + )) + } + } + + let credential = Arc::new(CountingGatedCredential { + calls: AtomicUsize::new(0), + first_call_entered: AtomicBool::new(false), + release_first_call: AtomicBool::new(false), + }); + + let url = Url::parse("amqps://example.com").unwrap(); + let connection = RecoverableConnection::new( + url.clone(), + None, + None, + credential.clone(), + Default::default(), + None, + ); + + let authorizer = Arc::new(Authorizer::new( + Arc::downgrade(&connection), + credential.clone(), + None, + )); + // Skip the real CBS attach; we are exercising the generation guard, not the + // broker handshake. + authorizer.disable_authorization().unwrap(); + connection.disable_connection().await.unwrap(); + + let path = Url::parse("amqps://example.com/test").unwrap(); + + let auth_task = { + let authorizer = authorizer.clone(); + let connection = connection.clone(); + let path = path.clone(); + tokio::spawn(async move { authorizer.authorize_path(&connection, &path).await }) + }; + + // Wait until the slow path is inside the first (gated) get_token: it has + // captured the pre-recovery generation and is mid-authorization. + while !credential.first_call_entered.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + assert_eq!(connection.generation(), 0); + + // Fire a recovery now, in the lock-free window. This advances the generation + // past one whole recovery, exactly the race #4454 describes. We advance the + // generation directly rather than running the full `simulate_reconnect`, + // because clearing the caches is irrelevant to the token guard and keeps the + // test focused. + connection.bump_generation_for_test(); + assert_eq!(connection.generation(), 2); + + // Release the first authorization; its token is now stale and must be + // discarded, triggering a re-authorization against the new generation. + credential.release_first_call.store(true, Ordering::SeqCst); + + let token = auth_task + .await + .expect("authorize_path task panicked") + .expect("authorize_path returned an error"); + + // Two get_token calls: the discarded first attempt and the clean retry. + assert_eq!( + credential.calls.load(Ordering::SeqCst), + 2, + "expected exactly one discard-and-retry" + ); + + // The token returned is cached and stamped at the stable post-recovery + // generation; the next lookup is a clean fast-path hit, with no further + // recovery needed. + let cached = authorizer + .authorization_scopes + .read() + .await + .get(&path) + .cloned(); + assert!( + cached.is_some(), + "a fresh token must be cached after the discard-and-retry" + ); + assert_eq!(cached.unwrap().token.secret(), token.token.secret()); + assert_eq!( + connection.generation(), + 2, + "no second recovery cycle should have been needed" + ); + } + + // #4454: the background refresh task is the token cache's *second* writer + // (alongside `authorize_path`), and it needs the same generation guard. When a + // recovery races a refresh, the token re-authorized against the now-dead CBS + // link must be discarded, not written back over the cache the recovery just + // cleared, otherwise the next operation serves a stale token and forces a + // second recovery cycle. + // + // This drives the guard deterministically by calling the extracted single-pass + // `refresh_due_tokens`: the cache is seeded with a token that is already due, + // a gated credential blocks the refresh's `get_token` inside the lock-free + // window, the test fires a simulated reconnect (bumping the generation) there, + // then releases. The refreshed token must be thrown away and the original left + // untouched. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn refresh_discards_tokens_refreshed_during_recovery() { + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + #[derive(Debug)] + struct GatedRefreshCredential { + calls: AtomicUsize, + entered: AtomicBool, + release: AtomicBool, + } + + #[async_trait::async_trait] + impl TokenCredential for GatedRefreshCredential { + async fn get_token( + &self, + _scopes: &[&str], + _options: Option>, + ) -> azure_core::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + self.entered.store(true, Ordering::SeqCst); + // Block inside the lock-free refresh window so the test can fire a + // recovery before the refreshed token is written back. + while !self.release.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + Ok(AccessToken::new( + azure_core::credentials::Secret::new("refreshed_token"), + OffsetDateTime::now_utc() + Duration::hours(1), + )) + } + } + + let credential = Arc::new(GatedRefreshCredential { + calls: AtomicUsize::new(0), + entered: AtomicBool::new(false), + release: AtomicBool::new(false), + }); + + let url = Url::parse("amqps://example.com").unwrap(); + let connection = RecoverableConnection::new( + url.clone(), + None, + None, + credential.clone(), + Default::default(), + None, + ); + let authorizer = Arc::new(Authorizer::new( + Arc::downgrade(&connection), + credential.clone(), + None, + )); + // Exercise the generation guard, not the broker handshake. + authorizer.disable_authorization().unwrap(); + connection.disable_connection().await.unwrap(); + + // Seed the cache with an "original" token that is already due for refresh. + let path = Url::parse("amqps://example.com/test").unwrap(); + let now = OffsetDateTime::now_utc(); + let original = AccessToken::new( + azure_core::credentials::Secret::new("original_token"), + now + Duration::seconds(1), + ); + authorizer + .authorization_scopes + .write() + .await + .insert(path.clone(), original); + + // A 10s bias makes the 1s-from-now token due, so the pass refreshes it. + let bias = Duration::seconds(10); + let refresh_task = { + let authorizer = authorizer.clone(); + tokio::spawn(async move { + let mut non_refreshable = HashSet::new(); + authorizer + .refresh_due_tokens(now, bias, &mut non_refreshable) + .await + }) + }; + + // Wait until the refresh is inside the gated `get_token`: it has captured + // the pre-recovery generation and is mid re-authorization. + while !credential.entered.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + assert_eq!(connection.generation(), 0); + + // Fire a recovery in the lock-free window: the #4454 race, refresh edition. + connection.bump_generation_for_test(); + assert_eq!(connection.generation(), 2); + + // Release the gated refresh; its token is now stale and must be discarded. + credential.release.store(true, Ordering::SeqCst); + let outcome = refresh_task + .await + .expect("refresh task panicked") + .expect("refresh_due_tokens returned an error"); + + // The pass must report the discard, not a clean completion. The caller + // backs off on `Discarded`, and it must: the cache still holds the due + // original token, so a `Completed { failed: false }` here would send the + // refresh loop straight back around with no sleep, once for every + // generation bump a recovery storm produces (#4454). + assert_eq!( + outcome, + RefreshPass::Discarded, + "a pass whose tokens a recovery discarded must ask the caller to back off" + ); + + // Exactly one refresh attempt was made, and the cache still holds the + // original token: the token refreshed against the torn-down connection was + // dropped at the guarded write-back rather than overwriting the cache. + assert_eq!( + credential.calls.load(Ordering::SeqCst), + 1, + "exactly one refresh attempt" + ); + let cached = authorizer + .authorization_scopes + .read() + .await + .get(&path) + .cloned() + .expect("the original token must remain cached"); + assert_eq!( + cached.token.secret(), + "original_token", + "a token refreshed during recovery must be discarded, not written back" + ); + } } diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs index 7208714449c..694f9d63935 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs @@ -5,7 +5,7 @@ use super::{ claims_based_security::RecoverableClaimsBasedSecurity, management::RecoverableManagementClient, - receiver::RecoverableReceiver, sender::RecoverableSender, + receiver::RecoverableReceiver, sender::RecoverableSender, MAX_GENERATION_RETRIES, }; use crate::{ common::{ @@ -31,7 +31,11 @@ use azure_core_amqp::{ use std::sync::Mutex; use std::{ collections::HashMap, - sync::{Arc, Weak}, + future::Future, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Weak, + }, }; use tracing::{debug, info, instrument, trace, warn}; @@ -89,9 +93,15 @@ pub(crate) struct RecoverableConnection { // *different* partitions never serialize on a shared lock, and the expensive // attach (authorize + session begin + link attach) happens without holding // the map-wide lock. See issues #2243 and #4563. - sender_instances: RwLock>>>>, - session_instances: RwLock>>>>, - receiver_instances: RwLock>>>>, + // + // Each cell is tagged with the recovery `generation` it was created under (see + // the `generation` field). A slow-path attach that races a recovery completes + // against a now-dead connection; comparing the cell's generation against the + // current one after the attach lets that path discard its stale result instead + // of caching and handing out a resource bound to the old connection. See #4454. + sender_instances: RwLock>>, + session_instances: RwLock>>, + receiver_instances: RwLock>>, pub(super) authorizer: Arc, // The service permits one `$cbs` link for each connection. Every // authorization attaches a link, uses it, and then drops it, so two @@ -102,6 +112,28 @@ pub(crate) struct RecoverableConnection { connection_name: String, pub(super) retry_options: RetryOptions, + // Recovery generation counter (#4454), used as a sequence lock. + // `apply_recovery_plan` bumps it once before it invalidates anything and once + // after, so the value is odd for exactly as long as a recovery is tearing state + // down and even at rest. Two properties follow, and `generation_is_current` + // tests both: + // + // * The value changes across every recovery, so a capture taken before one does + // not match afterwards. + // * The value is odd for a capture taken during one, so such a capture is + // rejected even when the recovery has not finished by the time it is tested. + // + // Invariant: a cached resource is only valid if the generation it was created + // under is even and still equals the current generation. The four slow paths + // (authorize_path, get_session, ensure_sender, ensure_receiver) do their AMQP IO + // with no map lock held; a recovery that overlaps that window leaves the + // captured generation odd, changed, or both, so the slow path discards its + // result rather than caching a resource bound to the dead connection. A single + // counter is used for all resource types: session-level recovery is rare, so the + // occasional extra re-init of an unaffected type after a narrower recovery is + // cheaper than the bookkeeping of per-type counters. + generation: AtomicU64, + #[cfg(test)] forced_error: Mutex>, @@ -111,29 +143,76 @@ pub(crate) struct RecoverableConnection { // how the operation wrappers behave. #[cfg(test)] forced_attach_error: Mutex>, + + // Test seam for the caller side of the #4454 supersession race. When armed, + // `run_peer_supersession_hook` fires once on the next generational init. It + // plays a peer task that drove a recovery in the window between a caller's + // generation capture and its cell resolution. See the hook for details. + #[cfg(test)] + peer_supersession_pending: std::sync::atomic::AtomicBool, +} + +/// A per-path cache cell tagged with the recovery [`RecoverableConnection::generation`] +/// it was created under. See #4454. +struct GenerationalCell { + generation: u64, + cell: Arc>>, +} + +impl Clone for GenerationalCell { + fn clone(&self) -> Self { + Self { + generation: self.generation, + cell: self.cell.clone(), + } + } } unsafe impl Send for RecoverableConnection {} unsafe impl Sync for RecoverableConnection {} -/// Returns the per-path `OnceCell` for `key`, inserting an uninitialized one if -/// absent. The read path is taken first so steady-state lookups share a read -/// lock; only the first insert for a key takes the write lock. The attach then -/// runs inside the returned `OnceCell`, so the map lock is never held across it -/// and different paths set up concurrently. Shared by the sender, session, and -/// receiver caches so all three keep identical concurrency semantics. +/// Returns the per-path cell for `key` valid at `generation`, inserting an +/// uninitialized one if absent. The read path is taken first so steady-state +/// lookups share a read lock; only the first insert for a key (or replacing a cell +/// left over from a previous generation) takes the write lock. The attach then +/// runs inside the returned `OnceCell`, so the map lock is never held across it and +/// different paths set up concurrently. Shared by the sender, session, and receiver +/// caches so all three keep identical concurrency semantics. +/// +/// A cached cell whose generation predates `generation` is stale: a recovery +/// cleared and re-stamped the connection after it was created. Such a cell is +/// replaced with a fresh one so the caller re-attaches against the live +/// connection. A cell at a *newer* generation than `generation` is returned +/// as-is rather than overwritten: a recovery already advanced past the +/// generation the caller captured, and a peer task may have attached a valid +/// resource into that newer cell. Clobbering it with a fresh cell stamped at the +/// older `generation` would discard that peer's work and force a redundant +/// re-attach, the exact wasted recovery cycle #4454 set out to remove. The +/// caller's post-`init` generation check sorts out the captured-then-superseded +/// case instead (see [`RecoverableConnection::get_or_init_generational`]). Only +/// a strictly-older or absent entry is replaced. See #4454. async fn or_init_cell( - map: &RwLock>>>>, + map: &RwLock>>, key: &Url, -) -> Arc>> { - if let Some(cell) = map.read().await.get(key) { - return cell.clone(); + generation: u64, +) -> GenerationalCell { + if let Some(entry) = map.read().await.get(key) { + if entry.generation >= generation { + return entry.clone(); + } + } + let mut guard = map.write().await; + match guard.get(key) { + Some(entry) if entry.generation >= generation => entry.clone(), + _ => { + let fresh = GenerationalCell { + generation, + cell: Arc::new(OnceCell::new()), + }; + guard.insert(key.clone(), fresh.clone()); + fresh + } } - map.write() - .await - .entry(key.clone()) - .or_insert_with(|| Arc::new(OnceCell::new())) - .clone() } /// Describes which per-connection caches an [`ErrorRecoveryAction`] must invalidate. @@ -219,10 +298,13 @@ impl RecoverableConnection { receiver_instances: RwLock::new(HashMap::new()), mgmt_client: RwLock::new(Arc::new(OnceCell::new())), authorizer, + generation: AtomicU64::new(0), #[cfg(test)] forced_error: Mutex::new(None), #[cfg(test)] forced_attach_error: Mutex::new(None), + #[cfg(test)] + peer_supersession_pending: std::sync::atomic::AtomicBool::new(false), } }) } @@ -337,7 +419,7 @@ impl RecoverableConnection { } let mut sender_instances = self.sender_instances.write().await; - for (path, cell) in sender_instances.drain() { + for (path, GenerationalCell { cell, .. }) in sender_instances.drain() { trace!("Detaching sender for path {}.", path); let Some(sender) = Arc::try_unwrap(cell).ok().and_then(OnceCell::into_inner) else { trace!( @@ -358,7 +440,7 @@ impl RecoverableConnection { } let mut receiver_instances = self.receiver_instances.write().await; - for (source_url, cell) in receiver_instances.drain() { + for (source_url, GenerationalCell { cell, .. }) in receiver_instances.drain() { trace!("Detaching receiver for source URL {}.", source_url); let Some(receiver) = Arc::try_unwrap(cell).ok().and_then(OnceCell::into_inner) else { trace!( @@ -379,7 +461,7 @@ impl RecoverableConnection { } let mut session_instances = self.session_instances.write().await; - for (session_id, cell) in session_instances.drain() { + for (session_id, GenerationalCell { cell, .. }) in session_instances.drain() { trace!("Detaching session for ID {}.", session_id); let Some(session) = Arc::try_unwrap(cell).ok().and_then(OnceCell::into_inner) else { trace!( @@ -494,7 +576,9 @@ impl RecoverableConnection { pub(crate) async fn close_receiver(self: &Arc, source_url: &Url) -> Result<()> { // Drop the map's write lock as soon as the cell is removed so the detach // (network I/O) doesn't hold it. - let Some(cell) = self.receiver_instances.write().await.remove(source_url) else { + let Some(GenerationalCell { cell, .. }) = + self.receiver_instances.write().await.remove(source_url) + else { // No entry for this path; nothing to detach. return Ok(()); }; @@ -535,6 +619,105 @@ impl RecoverableConnection { Ok(()) } + /// The recovery generation the caches are currently stamped at. See the + /// `generation` field and #4454. + fn current_generation(&self) -> u64 { + self.generation.load(Ordering::Acquire) + } + + /// Whether a resource created under `captured` is still bound to live state + /// (#4454). Every slow-path guard tests this, and it holds only when both + /// halves of the sequence-lock rule do: + /// + /// * `captured` is even, so no recovery was in flight when it was taken. An odd + /// capture came from inside `apply_recovery_plan`, where the connection may + /// already be gone (or about to be taken) and the caches are being cleared. + /// * `captured` still equals the current generation, so no recovery has started + /// since. + /// + /// Testing the parity matters on its own. `apply_recovery_plan` releases every + /// lock it takes, so it can stall between its two bumps under contention; a slow + /// path that captured an odd generation there has time to finish a whole attach + /// and test its capture while the value is still unchanged. Equality alone would + /// accept that attach. + pub(crate) fn generation_is_current(&self, captured: u64) -> bool { + captured.is_multiple_of(2) && self.current_generation() == captured + } + + /// Resolves the per-path cell for `key`, runs `init` to attach the resource + /// without holding the map lock, and guards the result against a racing + /// recovery via the generation counter (#4454). + /// + /// The attach (`init`) does its AMQP IO with no map lock held, so a recovery + /// can clear the caches and bump the generation while it is in flight. After + /// `init` completes this re-reads the generation: if it still matches the one + /// the cell was created under, the result is fresh and is returned. If it + /// changed, the just-attached resource is bound to a now-dead connection; + /// rather than caching and handing out that stale resource (the old behavior + /// that cost an extra recovery cycle on the next operation), the stale cell is + /// evicted and the whole attach retries against the new generation. + /// + /// `init` is therefore an `FnMut`: it may run more than once if recovery keeps + /// racing. [`MAX_GENERATION_RETRIES`] bounds the loop, so a pathological storm + /// of back-to-back recoveries surfaces an error instead of spinning forever. + async fn get_or_init_generational( + &self, + map: &RwLock>>, + key: &Url, + mut init: F, + ) -> azure_core_amqp::Result> + where + F: FnMut() -> Fut, + Fut: Future>>, + { + for _ in 0..MAX_GENERATION_RETRIES { + let generation = self.current_generation(); + + // Test seam (#4454): reproduce a peer task that drives a recovery and + // installs a newer cell in the window between this capture and the + // resolve below. In production this compiles away. + #[cfg(test)] + self.run_peer_supersession_hook(map, key).await; + + let entry = or_init_cell(map, key, generation).await; + let value = entry.cell.get_or_try_init(&mut init).await?; + + // If no recovery raced the attach above, the cell is valid; return it. + // Test the cell's *own* generation, not the value captured at the top of + // the loop: `or_init_cell` may have handed back a newer cell that a + // racing task installed, and such a cell is valid as long as its + // generation is still current. Using the captured `generation` here + // would wrongly discard (and evict) that peer's freshly-attached + // resource. See #4454. + if self.generation_is_current(entry.generation) { + return Ok(value.clone()); + } + + // A recovery cleared the caches mid-attach. The value we just produced + // (or read from a cell another racing task initialized) is bound to the + // old connection. Evict this cell if it is still the one mapped for + // `key` so the next pass re-inits against the new generation, then loop. + debug!( + %key, + "Discarding stale resource produced during recovery (#4454); re-initializing." + ); + let mut guard = map.write().await; + if let Some(current) = guard.get(key) { + if Arc::ptr_eq(¤t.cell, &entry.cell) { + guard.remove(key); + } + } + } + + // Intentionally a plain `AmqpError::with_message`: `should_retry_amqp_error` + // classifies this unrecognized kind as `ReturnError`, so exhausting the + // budget surfaces to the caller instead of looping. Do not "fix" this into a + // retryable kind, that would let a recovery storm spin here forever (#4454). + Err(AmqpError::with_message(format!( + "Exceeded retry budget ({MAX_GENERATION_RETRIES}) re-initializing resource '{key}' across recoveries" + ))) + } + #[instrument( level = "debug", skip_all, @@ -550,10 +733,11 @@ impl RecoverableConnection { ) -> azure_core_amqp::Result> { // Resolve the per-path cell while holding the map lock only briefly, then // initialize (which may begin a new AMQP session) without holding it, so - // that sessions for other partitions can be created concurrently. - let cell = self.session_cell(source_url).await; - let session = cell - .get_or_try_init(|| async { + // that sessions for other partitions can be created concurrently. The + // generation guard discards a session begun against a connection that a + // racing recovery has since replaced (#4454). + let session = self + .get_or_init_generational(&self.session_instances, source_url, || async { debug!(source_url = %source_url, "Creating session for partition."); let connection = self.ensure_connection().await?; @@ -568,13 +752,21 @@ impl RecoverableConnection { }) .await?; debug!(source_url = %source_url, "Cloning session for partition."); - Ok(session.clone()) + Ok(session) } - /// Returns the `OnceCell` that owns the session for `source_url`, inserting an - /// uninitialized one if absent. See [`or_init_cell`] for the locking strategy. + /// Returns the `OnceCell` that owns the session for `source_url` at the current + /// generation, inserting an uninitialized one if absent. See [`or_init_cell`] + /// for the locking strategy. Used in tests to assert cell identity. + #[cfg(test)] async fn session_cell(&self, source_url: &Url) -> Arc>> { - or_init_cell(&self.session_instances, source_url).await + or_init_cell( + &self.session_instances, + source_url, + self.current_generation(), + ) + .await + .cell } #[instrument( @@ -705,9 +897,8 @@ impl RecoverableConnection { // without holding it, so receivers for other partitions can be created // concurrently and steady-state receives never serialize on a shared // lock. See issues #2243 and #4563. - let cell = self.receiver_cell(source_url).await; - let receiver = cell - .get_or_try_init(|| async { + let receiver = self + .get_or_init_generational(&self.receiver_instances, source_url, || async { // Test seam: fail the attach with an injected error before // any network activity. The error leaves this closure on the // same path a rejected `receiver.attach` below takes. @@ -746,13 +937,21 @@ impl RecoverableConnection { }) .await?; - Ok(receiver.clone()) + Ok(receiver) } - /// Returns the `OnceCell` that owns the receiver for `source_url`, inserting - /// an uninitialized one if absent. See [`or_init_cell`] for the locking strategy. + /// Returns the `OnceCell` that owns the receiver for `source_url` at the current + /// generation, inserting an uninitialized one if absent. See [`or_init_cell`] + /// for the locking strategy. Used in tests to assert cell identity. + #[cfg(test)] async fn receiver_cell(&self, source_url: &Url) -> Arc>> { - or_init_cell(&self.receiver_instances, source_url).await + or_init_cell( + &self.receiver_instances, + source_url, + self.current_generation(), + ) + .await + .cell } #[instrument( @@ -771,9 +970,8 @@ impl RecoverableConnection { // attach (authorize + session begin + link attach) without holding it, so // that senders for other partitions can be created concurrently and // steady-state sends never serialize on a shared lock. See issue #2243. - let cell = self.sender_cell(path).await; - let sender = cell - .get_or_try_init(|| async { + let sender = self + .get_or_init_generational(&self.sender_instances, path, || async { // Ensure that we are authorized to access the senders path. self.authorizer.authorize_path(self, path).await?; @@ -812,13 +1010,17 @@ impl RecoverableConnection { }) .await?; - Ok(sender.clone()) + Ok(sender) } - /// Returns the `OnceCell` that owns the sender for `path`, inserting an - /// uninitialized one if absent. See [`or_init_cell`] for the locking strategy. + /// Returns the `OnceCell` that owns the sender for `path` at the current + /// generation, inserting an uninitialized one if absent. See [`or_init_cell`] + /// for the locking strategy. Used in tests to assert cell identity. + #[cfg(test)] async fn sender_cell(&self, path: &Url) -> Arc>> { - or_init_cell(&self.sender_instances, path).await + or_init_cell(&self.sender_instances, path, self.current_generation()) + .await + .cell } #[instrument( @@ -875,23 +1077,65 @@ impl RecoverableConnection { /// Side-effecting half of `recover_from_error`: takes the locks and clears /// whichever caches the [`RecoveryPlan`] flagged. + /// + /// #4454 stale-resource window. Any plan that invalidates something brackets + /// that invalidation with a bump of the recovery `generation`: one before it + /// and one after it, which leaves the counter odd for the whole span. A slow + /// path (authorize_path / get_session / ensure_sender / ensure_receiver) that is + /// mid-attach captured a generation from inside or before that bracket, so + /// `generation_is_current` rejects it on completion and the slow path discards + /// its result instead of caching a resource bound to the connection this + /// recovery just tore down. The body explains why one bump on either side alone + /// is not enough. async fn apply_recovery_plan(&self, plan: RecoveryPlan) { let connection_id = self.get_connection_id(); + + // A plan that invalidates anything brackets the invalidation with a + // generation bump: one before it touches the connection or any cache, and + // one after the last of them. The generation is therefore odd for exactly + // the span in which this connection's state is inconsistent, which is the + // sequence-lock rule `generation_is_current` tests (#4454). + // + // Both bumps are needed, and so is the parity test: + // + // * Without the closing bump, a slow path that captured the old generation + // can clone the connection, attach, and test its capture before the single + // bump lands. The generation still matches, so it caches and returns a + // resource bound to the connection this recovery drops a moment later. + // * Without the opening bump, a slow path can capture the new generation and + // *then* clone the old connection, which `connections` still holds. Its + // post-init test matches too, so the same stale resource reaches the + // caller. The token cache has the same shape: a reader that runs after the + // bump and before `clear()` gets a token that was authorized on the CBS + // link of the connection being dropped. + // * Without the parity test, a slow path that captured a generation between + // the two bumps is accepted for as long as this function has not reached + // the closing one. Every lock below is released before the next is taken, + // so a contended recovery can stall here long enough for that slow path to + // finish a whole attach against the connection being dropped. + // + // A task that captures the final, even generation started after the last + // invalidation, so it finds an empty cache and builds against the new + // connection. + let invalidates = plan.drop_connection + || plan.clear_authorizer + || plan.clear_sessions + || plan.clear_senders + || plan.clear_receivers; + + if invalidates { + self.generation.fetch_add(1, Ordering::AcqRel); + } + if plan.drop_connection { self.connections.lock().await.take(); debug!(connection_id = %connection_id, "Recovery: dropped AMQP connection."); } + if plan.clear_authorizer { self.authorizer.clear().await; debug!(connection_id = %connection_id, "Recovery: cleared authorizer tokens."); } - // Clearing a cache drops its per-path `OnceCell` entries, so the next - // `ensure_*` for a path re-inits a fresh cell against the new connection. - // A task already mid-`get_or_try_init` (or holding a value it just read) - // keeps using the old cell until it next re-enters `ensure_*`, so there's - // a brief window where stale, pre-recovery links can still be handed out. - // Closing that window (invalidating in-flight work immediately via a - // recovery generation counter) is tracked in #4454. if plan.clear_sessions { let mut sessions = self.session_instances.write().await; let count = sessions.len(); @@ -918,6 +1162,95 @@ impl RecoverableConnection { *self.mgmt_client.write().await = Arc::new(OnceCell::new()); debug!(connection_id = %connection_id, "Recovery: dropped management client."); } + + // Closing bump. See the comment above the opening one. + if invalidates { + self.generation.fetch_add(1, Ordering::AcqRel); + } + } + + /// The recovery generation, exposed for the authorizer's slow-path guard + /// (#4454) and for tests. See `current_generation`. + pub(crate) fn generation(&self) -> u64 { + self.current_generation() + } + + /// Test hook: simulate the cache-clearing half of a `ReconnectConnection` + /// recovery (bump the generation and clear the per-path caches) without needing + /// a live broker connection. Used to drive the #4454 stale-resource race + /// deterministically. + #[cfg(test)] + pub(crate) async fn simulate_reconnect(&self) { + self.apply_recovery_plan( + RecoveryPlan::for_action(&ErrorRecoveryAction::ReconnectConnection) + .expect("ReconnectConnection has a recovery plan"), + ) + .await; + } + + /// Test hook: advance the recovery generation past one whole recovery without + /// taking the cache locks. Used to exercise the #4454 generation guard in + /// isolation. The step is two, the same as a completed `apply_recovery_plan`, + /// so the counter is left even and a capture taken after this hook is current. + #[cfg(test)] + pub(crate) fn bump_generation_for_test(&self) { + self.generation.fetch_add(2, Ordering::AcqRel); + } + + /// Test hook: park the generation mid-recovery, as `apply_recovery_plan` does + /// between its two bumps, and leave it there. Used to assert that a capture + /// taken during a recovery is rejected even when the recovery has not finished. + #[cfg(test)] + pub(crate) fn enter_recovery_generation_for_test(&self) { + self.generation.fetch_add(1, Ordering::AcqRel); + } + + /// Arms `run_peer_supersession_hook` to fire on the next generational init. + /// See that hook and the `get_or_init_generational_returns_superseding_peer_cell` + /// test. + #[cfg(test)] + pub(crate) fn arm_peer_supersession_for_test(&self) { + self.peer_supersession_pending + .store(true, Ordering::Release); + } + + /// Test seam for the caller side of the #4454 supersession property. When + /// armed by `arm_peer_supersession_for_test`, this fires once, in the window + /// between a caller capturing its generation and resolving its cell. It plays + /// a peer task that drove a recovery in that window: it bumps the generation, + /// so the caller's captured value is now stale, and installs a fresh, empty + /// cell for `key` at the new generation. `or_init_cell` then hands that newer + /// cell back to the caller. + /// + /// The correct guard in `get_or_init_generational` compares the cell's own + /// generation, so it returns the resource the caller attaches into that cell: + /// the cell is current, the resource is valid, and no eviction happens. A + /// guard that compared the captured local generation instead would see a + /// mismatch, evict the valid cell, and re-attach, the wasted recovery cycle + /// #4454 removes. The empty cell makes that difference observable: a correct + /// return keeps the caller's single `init`, a wrong eviction forces a second. + /// This compiles away in production; the field is `cfg(test)` only. + #[cfg(test)] + async fn run_peer_supersession_hook( + &self, + map: &RwLock>>, + key: &Url, + ) { + if self.peer_supersession_pending.swap(false, Ordering::AcqRel) { + // Drive the peer's recovery to completion, so the caller's captured + // generation is now behind and the new one is settled (even). + self.generation.fetch_add(2, Ordering::AcqRel); + // Install the peer's fresh, higher-generation cell. It is empty on + // purpose, so the caller's own `init` fills it. + let generation = self.current_generation(); + map.write().await.insert( + key.clone(), + GenerationalCell { + generation, + cell: Arc::new(OnceCell::new()), + }, + ); + } } /// Classifies an [`AmqpError`] into the recovery action the retry loop should take. @@ -1241,6 +1574,337 @@ mod tests { assert!(connection.receiver_cell(&path_a).await.get().is_none()); } + // #4454: a recovery that clears the per-path caches must bump the recovery + // generation so racing slow-path attaches can detect it. A simulated + // ReconnectConnection must advance `generation()`. + // + // The step is two, not one: `apply_recovery_plan` brackets its invalidation + // with a bump on each side, so a task that captures the generation part way + // through the recovery also sees a mismatch when it completes. The exact value + // is asserted here to pin that bracketing; nothing else compares generations by + // anything other than equality. + #[tokio::test] + async fn simulate_reconnect_bumps_generation() { + let url = Url::parse("amqps://example.com").unwrap(); + let connection = RecoverableConnection::new( + url, + None, + None, + Arc::new(MockCredential), + Default::default(), + None, + ); + + assert_eq!(connection.generation(), 0); + connection.simulate_reconnect().await; + assert_eq!(connection.generation(), 2); + connection.simulate_reconnect().await; + assert_eq!(connection.generation(), 4); + } + + // #4454: a generation captured *part way through* a recovery must also end up + // stale. This is the edge a single leading bump leaves open: a slow path that + // starts after the bump can still clone the connection that `apply_recovery_plan` + // has not taken yet, or read a token that `clear()` has not removed yet, and its + // post-init check would then match and hand the caller a resource bound to the + // connection this recovery is dropping. + // + // The token cache's write lock is the seam. Holding it stops the recovery inside + // `authorizer.clear()`, which is after the opening bump and before the closing + // one, so the test can capture the generation a racing slow path would see. + #[tokio::test] + async fn recovery_generation_differs_for_a_mid_recovery_capture() { + let url = Url::parse("amqps://example.com").unwrap(); + let connection = RecoverableConnection::new( + url, + None, + None, + Arc::new(MockCredential), + Default::default(), + None, + ); + connection.disable_connection().await.unwrap(); + + let scopes = connection.authorizer.lock_scopes_for_test().await; + + let recovery = { + let connection = connection.clone(); + tokio::spawn(async move { connection.simulate_reconnect().await }) + }; + + // Wait for the opening bump. The recovery then blocks on the guard above. + while connection.generation() == 0 { + tokio::task::yield_now().await; + } + let captured_mid_recovery = connection.generation(); + + drop(scopes); + recovery.await.expect("recovery task panicked"); + + assert_ne!( + connection.generation(), + captured_mid_recovery, + "a generation captured during a recovery must not survive it, or a slow \ + path that started mid-recovery would pass its post-init check and cache \ + a resource bound to the dropped connection" + ); + } + + // #4454: the core of the fix. A cell resolved under generation N must be + // replaced by a fresh, distinct cell once the generation advances to N+1, + // because the recovery that bumped the generation tore down the connection the + // old cell's resource was attached to. Resolving at the same generation must + // keep returning the same cell (so we don't lose single-init within a + // generation). + #[tokio::test] + async fn stale_generation_cell_is_replaced() { + let url = Url::parse("amqps://example.com").unwrap(); + let connection = RecoverableConnection::new( + url, + None, + None, + Arc::new(MockCredential), + Default::default(), + None, + ); + let path = Url::parse("amqps://example.com/eh/Partitions/0").unwrap(); + + // Two resolutions at the same generation share a cell. + let cell_gen0 = connection.sender_cell(&path).await; + assert!(Arc::ptr_eq( + &cell_gen0, + &connection.sender_cell(&path).await + )); + + // After a recovery, the next resolution returns a brand-new cell. + connection.simulate_reconnect().await; + let cell_gen1 = connection.sender_cell(&path).await; + assert!( + !Arc::ptr_eq(&cell_gen0, &cell_gen1), + "cell from the previous generation must be discarded after recovery" + ); + + // And that new cell is itself stable within its generation. + assert!(Arc::ptr_eq( + &cell_gen1, + &connection.sender_cell(&path).await + )); + } + + // #4454: an attach that both starts and finishes inside a recovery must be + // discarded too. The two bumps make the generation odd for the span in which + // `apply_recovery_plan` is invalidating state, so a capture taken there is + // rejected on parity alone, without waiting for the recovery to end. + // + // Equality against the captured value cannot catch this case: the generation + // has not moved since the capture. `apply_recovery_plan` releases each lock + // before it takes the next, so a contended recovery can stall between its bumps + // long enough for a slow path to finish attaching to the connection it is + // dropping. The test parks the generation mid-recovery to hold that state open. + #[tokio::test] + async fn generation_captured_mid_recovery_is_never_current() { + let url = Url::parse("amqps://example.com").unwrap(); + let connection = RecoverableConnection::new( + url, + None, + None, + Arc::new(MockCredential), + Default::default(), + None, + ); + let path = Url::parse("amqps://example.com/eh/Partitions/0").unwrap(); + + assert!(connection.generation_is_current(connection.generation())); + + // Park the counter where `apply_recovery_plan` holds it between its bumps. + connection.enter_recovery_generation_for_test(); + let captured_mid_recovery = connection.generation(); + assert!( + !connection.generation_is_current(captured_mid_recovery), + "a generation captured during a recovery must never be current, even \ + while the recovery is still in flight and the value is unchanged" + ); + + // An attach that runs entirely inside the recovery is therefore never + // cached. It retries to the budget and surfaces an error instead. + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let map: RwLock>> = RwLock::new(HashMap::new()); + let result = + connection + .get_or_init_generational(&map, &path, || { + let calls = calls.clone(); + async move { + Ok::<_, AmqpError>(Arc::new(calls.fetch_add(1, Ordering::SeqCst) as u64)) + } + }) + .await; + + assert!( + result.is_err(), + "a resource attached during a recovery must not be handed to the caller" + ); + assert_eq!(calls.load(Ordering::SeqCst), MAX_GENERATION_RETRIES); + assert!( + map.read().await.get(&path).is_none(), + "no cell attached during a recovery may stay cached" + ); + } + + // #4454: `get_or_init_generational` must discard a value produced during a + // racing recovery and re-init against the new generation. Here the init closure + // fires a simulated reconnect on its first call (the in-flight-slow-path + // window), so the first attempt's value is stale and must be thrown away; the + // second attempt runs at a stable generation and its value is the one returned + // and cached. + #[tokio::test] + async fn get_or_init_generational_discards_value_produced_during_recovery() { + let url = Url::parse("amqps://example.com").unwrap(); + let connection = RecoverableConnection::new( + url, + None, + None, + Arc::new(MockCredential), + Default::default(), + None, + ); + let path = Url::parse("amqps://example.com/eh/Partitions/0").unwrap(); + + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let map: RwLock>> = RwLock::new(HashMap::new()); + + let result = connection + .get_or_init_generational(&map, &path, || { + let calls = calls.clone(); + let connection = &connection; + async move { + let attempt = calls.fetch_add(1, Ordering::SeqCst); + // On the first attempt only, simulate a recovery firing during + // the lock-free init window. This bumps the generation, so the + // value produced here is stale and must be discarded. + if attempt == 0 { + connection.simulate_reconnect().await; + } + Ok::<_, AmqpError>(Arc::new(attempt as u64)) + } + }) + .await + .expect("init should succeed on the second, stable-generation attempt"); + + // The closure ran twice: once racing the recovery (discarded), once clean. + assert_eq!(calls.load(Ordering::SeqCst), 2); + // The returned value is the second attempt's (index 1), not the stale first. + assert_eq!(*result, 1); + // The cached cell holds the fresh value, stamped at the post-recovery generation. + let cached = map.read().await.get(&path).cloned().unwrap(); + assert_eq!(cached.generation, connection.generation()); + assert_eq!(**cached.cell.get().unwrap(), 1); + } + + // #4454 regression: `or_init_cell` must never overwrite a cell at a *newer* + // generation than the one the caller captured. A slow task that captured + // generation N can reach the lookup only after a recovery advanced to N+1 and a + // peer task already cached a valid resource there; clobbering it with a fresh + // gen-N cell would discard the peer's freshly-attached resource and force a + // redundant re-attach, the exact wasted recovery cycle the fix removes. A + // strictly-older cached cell must still be replaced so the caller re-attaches + // against the live connection. + #[tokio::test] + async fn or_init_cell_reuses_newer_cell_and_replaces_older() { + let path = Url::parse("amqps://example.com/eh/Partitions/0").unwrap(); + let map: RwLock>> = RwLock::new(HashMap::new()); + + // A peer task at generation 1 attached a resource and cached it. + let newer = or_init_cell(&map, &path, 1).await; + newer.cell.set(Arc::new(42)).await.unwrap(); + assert_eq!(newer.generation, 1); + + // A slow task that captured the stale generation 0 resolves the same key. It + // must get the gen-1 cell back, value intact, not a fresh gen-0 cell that + // throws the peer's work away. + let stale = or_init_cell(&map, &path, 0).await; + assert!( + Arc::ptr_eq(&stale.cell, &newer.cell), + "a cell newer than the captured generation must be reused, not clobbered" + ); + assert_eq!(stale.generation, 1); + assert_eq!(**stale.cell.get().unwrap(), 42); + + // Resolving at a generation strictly newer than the cached cell replaces it + // with a fresh, empty cell so the caller re-attaches against the live + // connection. + let replaced = or_init_cell(&map, &path, 2).await; + assert!( + !Arc::ptr_eq(&replaced.cell, &newer.cell), + "a cell older than the captured generation must be replaced" + ); + assert_eq!(replaced.generation, 2); + assert!(replaced.cell.get().is_none()); + } + + // #4454 regression, the caller side of the supersession property. `or_init_cell` + // owns one half: it never overwrites a newer cell (see + // `or_init_cell_reuses_newer_cell_and_replaces_older`). `get_or_init_generational` + // owns the other half, exercised here: when `or_init_cell` hands back a cell that + // is newer than the caller's captured generation but still current, the caller + // must attach into it and return the result, not evict it against the stale + // captured generation and re-attach. + // + // The scenario is a slow caller that captures generation N, then a peer task + // drives a recovery to N+1 and installs a fresh cell there before the caller + // resolves its own cell. The `run_peer_supersession_hook` seam reproduces that + // peer exactly, in the capture-to-resolve window, so the race is deterministic. + // + // With the correct guard (compare the cell's own generation) the init closure + // runs once and its value is returned. If the guard wrongly compared the captured + // local generation, the caller would evict the valid cell and run init a second + // time; this test then fails on the call count and the returned value. That is the + // mutation the earlier proof left uncaught, so this test closes the gap. + #[tokio::test] + async fn get_or_init_generational_returns_superseding_peer_cell() { + let url = Url::parse("amqps://example.com").unwrap(); + let connection = RecoverableConnection::new( + url, + None, + None, + Arc::new(MockCredential), + Default::default(), + None, + ); + let path = Url::parse("amqps://example.com/eh/Partitions/0").unwrap(); + + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let map: RwLock>> = RwLock::new(HashMap::new()); + + // Arm the peer: on the next generational init it bumps the generation and + // installs a fresh, higher-generation cell between the capture and the resolve. + connection.arm_peer_supersession_for_test(); + + let result = connection + .get_or_init_generational(&map, &path, || { + let calls = calls.clone(); + async move { + let attempt = calls.fetch_add(1, Ordering::SeqCst); + Ok::<_, AmqpError>(Arc::new(attempt as u64)) + } + }) + .await + .expect("init should succeed against the peer's current-generation cell"); + + // The init closure ran exactly once: the caller attached into the peer's + // newer-but-current cell and returned it, with no eviction and no retry. + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "a newer-but-current cell must be returned, not evicted and re-initialized" + ); + // The returned value is that single init's value (attempt index 0). + assert_eq!(*result, 0); + // The cached cell is stamped at the post-recovery generation and holds the value. + let cached = map.read().await.get(&path).cloned().unwrap(); + assert_eq!(cached.generation, connection.generation()); + assert_eq!(**cached.cell.get().unwrap(), 0); + } + // The RecoverableConnection supports using a custom endpoint for connecting to Event Hubs proxies. // This test verifies that the custom endpoint is properly stored in the RecoverableConnection. #[test] @@ -1375,6 +2039,20 @@ mod tests { RecoverableConnection::should_retry_amqp_error(&err), ErrorRecoveryAction::ReconnectLink ); + + // Test SimpleMessage (the kind `AmqpError::with_message` produces) -> + // ReturnError. The retry-budget-exhausted errors in `get_or_init_generational` + // and `authorize_path` are `with_message` errors that intentionally rely on + // this classification to surface instead of spinning across a recovery storm + // (#4454). This pins that contract: adding an explicit `SimpleMessage` arm, or + // flipping the `_` default to a retryable action, must fail here and force a + // deliberate decision rather than silently turning those backstops into an + // infinite retry loop. + let err = AmqpError::with_message("retry budget exhausted"); + assert_eq!( + RecoverableConnection::should_retry_amqp_error(&err), + ErrorRecoveryAction::ReturnError + ); } #[test] diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/mod.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/mod.rs index 8ed863a0cba..d5a73802bd2 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/mod.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/mod.rs @@ -9,3 +9,25 @@ mod sender; pub(crate) use connection::RecoverableConnection; pub(crate) use sender::RecoverableSender; + +/// How many times a generation-guarded cache fill retries when a recovery races +/// it (#4454). +/// +/// A resource that is attached while a recovery runs is bound to the connection +/// that the recovery tore down, so the cache discards it and attaches again +/// against the new generation. This bounds that loop: a storm of back-to-back +/// recoveries surfaces an error instead of spinning forever. +/// +/// The bound is generous because recovery is rare, and each pass makes forward +/// progress against a newer connection. To reach it, the connection must go down +/// faster than it can come up. +/// +/// One recovery costs one pass in the usual case. It costs two for a task that +/// captures its generation inside the recovery, because such a capture is odd and +/// is rejected on parity, so the task attaches once more before it reaches a +/// settled generation. The bound therefore covers at least four back-to-back +/// recoveries in the worst case. +/// +/// Every generation-guarded cache in this crate uses this one value, so the +/// policy stays the same for connections, senders, receivers, and tokens. +pub(crate) const MAX_GENERATION_RETRIES: usize = 8; diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/mod.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/mod.rs index da868359421..67cb50c8e31 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/mod.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/mod.rs @@ -936,6 +936,100 @@ mod tests { Ok(()) } + // Send to a single partition in a tight loop; any error (including a + // post-reconnect unauthorized / detached error from a stale token) panics the + // loop and fails the test. + async fn send_to_partition(producer: Arc, partition: &str) { + loop { + let batch = producer + .create_batch(Some(EventDataBatchOptions { + partition_id: Some(partition.to_string()), + ..Default::default() + })) + .await + .unwrap(); + assert!(batch + .try_add_event_data( + EventData::builder().with_body(b"Hello, World!").build(), + None, + ) + .unwrap()); + producer.send_batch(batch, None).await.unwrap(); + } + } + + // #4454: after a connection-level reconnect the per-path authorization tokens + // must be re-established cleanly on the new connection. Sending to several + // partitions concurrently keeps multiple `authorize_path` re-authorizations in + // flight across the forced `ConnectionClosedByRemote`, so a token cached against + // the torn-down connection (the stale-resource race this issue targets) would + // surface here as an unauthorized / detached error and panic a send loop's + // `unwrap`. A clean 30s run means every partition re-authorized against the new + // connection without a second recovery cycle. + #[recorded::test(live)] + async fn force_errors_concurrent_authorize_send_reconnect(ctx: TestContext) -> Result<()> { + const TEST_NAME: &str = "force_errors_concurrent_authorize_send_reconnect"; + let recording = ctx.recording(); + let host = recording.var("EVENTHUBS_HOST", None); + let eventhub = recording.var("EVENTHUB_NAME", None); + let credential = recording.credential(); + let producer = Arc::new( + ProducerClient::builder() + .with_application_id(TEST_NAME.to_string()) + .open(host.as_str(), eventhub.as_str(), credential.clone()) + .await?, + ); + + // Derive the partition IDs from the Event Hub rather than hard-coding + // "0".."3", which would panic on a hub configured with fewer than four + // partitions. The race this test targets only needs several + // `authorize_path` re-authorizations in flight at once, so send to up to + // four of whatever partitions the hub actually exposes. + let partition_ids = producer.get_eventhub_properties().await?.partition_ids; + assert!( + partition_ids.len() >= 2, + "this test needs at least 2 partitions to keep concurrent authorizations \ + in flight across the reconnect, but the configured Event Hub has {}", + partition_ids.len() + ); + let partition_ids: Vec = partition_ids.into_iter().take(4).collect(); + + force_errors( + producer.clone(), + move |producer: Arc| { + let partition_ids = partition_ids.clone(); + async move { + // Run the send loops via `join_all` (not `tokio::spawn`) so they + // are cancelled with the test future when `force_errors`'s + // timeout arm fires. + futures::future::join_all( + partition_ids + .iter() + .map(|partition| send_to_partition(producer.clone(), partition)), + ) + .await; + } + }, + |producer| { + producer + .force_error(azure_core_amqp::AmqpError::from( + AmqpErrorKind::ConnectionClosedByRemote(Box::new( + azure_core::error::Error::new( + azure_core::error::ErrorKind::Other, + "Forced error", + ), + )), + )) + .unwrap(); + }, + Duration::seconds(10), // Seconds until forcing the error. + Duration::seconds(30), // Seconds until test timeout. + ) + .await?; + + Ok(()) + } + #[recorded::test(live)] async fn force_errors_producer_properties_connection(ctx: TestContext) -> Result<()> { const TEST_NAME: &str = "force_errors_producer_properties_connection";