From 169a5aea2c05fae36642c56d3c7f36ec5dfc852b Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Thu, 20 Aug 2026 14:21:43 -0400 Subject: [PATCH 1/2] test(eventhubs): pin a lost ownership claim as a normal outcome InMemoryCheckpointStore::claim_ownership returns an error when a competing instance rotates the ETag first. That error travels up through the load balancer to EventProcessor::run, which stops every partition. BlobCheckpointStore already treats the same race as a normal outcome and returns an empty result. Add four tests for the target contract: - a lost claim returns Ok with no ownership and does not change the winner's record - a lost claim in the middle of a batch does not cancel the partitions behind it, and the winners' rotated ETags reach the caller - a validation failure still returns an error, which stops a fix that swallows every error in the loop - the load balancer survives a lost claim and owns no partition Three of the four fail against the current store. The validation test passes today and guards the behavior that must not change. --- .../src/event_processor/load_balancer.rs | 90 ++++++++++- .../tests/eventhubs_checkpoint_store.rs | 143 ++++++++++++++++++ 2 files changed, 231 insertions(+), 2 deletions(-) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/event_processor/load_balancer.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/event_processor/load_balancer.rs index fb72c0c42da..4ff0ac68805 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/event_processor/load_balancer.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/event_processor/load_balancer.rs @@ -392,8 +392,10 @@ impl LoadBalancer { pub(crate) mod tests { use super::*; use crate::{ - event_processor::Ownership, in_memory_checkpoint_store::InMemoryCheckpointStore, - models::ConsumerClientDetails, CheckpointStore, + event_processor::Ownership, + in_memory_checkpoint_store::InMemoryCheckpointStore, + models::{Checkpoint, ConsumerClientDetails}, + CheckpointStore, }; use azure_core::{time::OffsetDateTime, Result}; use azure_core_test::{recorded, TestContext}; @@ -580,6 +582,90 @@ pub(crate) mod tests { } } + /// A checkpoint store that lets a competitor claim the partition first. + /// A sequential test cannot reach the race on its own, because + /// `get_available_partitions` and `load_balance` have no await point + /// between them. The wrapper only orders the two calls. + struct LostRaceCheckpointStore { + inner: Arc, + competitor: Mutex>>, + } + + #[async_trait::async_trait] + impl CheckpointStore for LostRaceCheckpointStore { + async fn claim_ownership(&self, ownerships: &[Ownership]) -> Result> { + // Take the competitor in its own statement, so the guard drops before the await. + let competitor = self.competitor.lock().unwrap().take(); + if let Some(competitor) = competitor { + self.inner.claim_ownership(&competitor).await?; + } + self.inner.claim_ownership(ownerships).await + } + + async fn list_checkpoints( + &self, + namespace: &str, + event_hub_name: &str, + consumer_group: &str, + ) -> Result> { + self.inner + .list_checkpoints(namespace, event_hub_name, consumer_group) + .await + } + + async fn list_ownerships( + &self, + namespace: &str, + event_hub_name: &str, + consumer_group: &str, + ) -> Result> { + self.inner + .list_ownerships(namespace, event_hub_name, consumer_group) + .await + } + + async fn update_checkpoint(&self, checkpoint: Checkpoint) -> Result<()> { + self.inner.update_checkpoint(checkpoint).await + } + } + + /// A partition that another instance claimed first must not stop the load + /// balancer. The instance that lost the race reports no partition. + #[tokio::test] + async fn load_balance_survives_a_lost_claim() { + let inner = Arc::new(InMemoryCheckpointStore::new()); + let store = Arc::new(LostRaceCheckpointStore { + inner: inner.clone(), + competitor: Mutex::new(Some(vec![new_test_ownership("0", "winning-client")])), + }); + + let load_balancer = LoadBalancer::new( + store, + new_test_consumer_client_details("losing-client"), + ProcessorStrategy::Balanced, + Duration::seconds(3600), + None, + ); + + let result = load_balancer.load_balance(&["0"]).await; + assert!( + result.is_ok(), + "a lost claim must not stop the load balancer, got: {:?}", + result.as_ref().err() + ); + assert!( + result.unwrap().is_empty(), + "the losing instance must own no partition" + ); + + let ownerships = inner + .list_ownerships(TEST_EVENTHUB_FQDN, TEST_EVENTHUB_NAME, TEST_CONSUMER_GROUP) + .await + .unwrap(); + assert_eq!(ownerships.len(), 1); + assert_eq!(ownerships[0].owner_id, Some("winning-client".to_string())); + } + fn find_common(a: Vec, b: Vec) -> Vec { let mut common = vec![]; for item in a.into_iter() { diff --git a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_checkpoint_store.rs b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_checkpoint_store.rs index 63c414183a5..4fc8c00f252 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_checkpoint_store.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_checkpoint_store.rs @@ -101,6 +101,149 @@ async fn test_claim_ownership_renewal_rotates_etag_and_timestamp() { assert_eq!(*stale.unwrap_err().kind(), AzureErrorKind::Other); } +/// A claim that lost the race is a normal outcome. The store must report the +/// loss with an empty result and keep the winner's record. +#[tokio::test] +async fn test_claim_ownership_lost_claim_is_not_an_error() { + common::setup(); + let store = InMemoryCheckpointStore::new(); + let ownership = Ownership { + fully_qualified_namespace: "ns.servicebus.windows.net".to_string(), + event_hub_name: "event_hub".to_string(), + consumer_group: "consumer_group".to_string(), + partition_id: "0".to_string(), + owner_id: Some("owner-a".to_string()), + ..Default::default() + }; + + let first = store.claim_ownership(&[ownership]).await.unwrap(); + assert_eq!(first.len(), 1); + + // A competing renewal rotates the ETag, which makes the first record stale. + let second = store.claim_ownership(&[first[0].clone()]).await.unwrap(); + assert_eq!(second.len(), 1); + + let lost = store.claim_ownership(&[first[0].clone()]).await; + assert!( + lost.is_ok(), + "a lost claim must not be an error, got: {:?}", + lost.as_ref().err() + ); + assert!( + lost.unwrap().is_empty(), + "the losing claim must return no ownership" + ); + + let ownerships = store + .list_ownerships("ns.servicebus.windows.net", "event_hub", "consumer_group") + .await + .unwrap(); + assert_eq!(ownerships.len(), 1); + assert_eq!( + ownerships[0].etag, second[0].etag, + "a lost claim must not mutate the record" + ); +} + +/// One lost claim must not cancel the partitions behind it in the batch. The +/// stale partition sits in the middle of the batch, so a store that stops at +/// the first conflict fails this test. +#[tokio::test] +async fn test_claim_ownership_continues_past_a_lost_claim() { + common::setup(); + let store = InMemoryCheckpointStore::new(); + let new_ownership = |partition_id: &str, owner_id: &str| Ownership { + fully_qualified_namespace: "ns.servicebus.windows.net".to_string(), + event_hub_name: "event_hub".to_string(), + consumer_group: "consumer_group".to_string(), + partition_id: partition_id.to_string(), + owner_id: Some(owner_id.to_string()), + ..Default::default() + }; + + let claimed = store + .claim_ownership(&[ + new_ownership("0", "owner-a"), + new_ownership("1", "owner-a"), + new_ownership("2", "owner-a"), + ]) + .await + .unwrap(); + assert_eq!(claimed.len(), 3); + + // A second instance takes partition 1 behind the caller's back. + let mut rotated = claimed[1].clone(); + rotated.owner_id = Some("owner-b".to_string()); + let winner_b = store.claim_ownership(&[rotated]).await.unwrap(); + assert_eq!(winner_b.len(), 1); + + let result = store + .claim_ownership(&[claimed[0].clone(), claimed[1].clone(), claimed[2].clone()]) + .await; + assert!( + result.is_ok(), + "one lost claim must not cancel the batch, got: {:?}", + result.as_ref().err() + ); + + let kept = result.unwrap(); + let mut partition_ids = kept + .iter() + .map(|o| o.partition_id.clone()) + .collect::>(); + partition_ids.sort(); + assert_eq!(partition_ids, vec!["0".to_string(), "2".to_string()]); + + let kept_zero = kept + .iter() + .find(|o| o.partition_id == "0") + .expect("partition 0 stays with the caller"); + assert_ne!( + kept_zero.etag, claimed[0].etag, + "the winner's rotated ETag must reach the caller" + ); + let kept_two = kept + .iter() + .find(|o| o.partition_id == "2") + .expect("partition 2 stays with the caller"); + assert_ne!( + kept_two.etag, claimed[2].etag, + "the winner's rotated ETag must reach the caller" + ); + + let ownerships = store + .list_ownerships("ns.servicebus.windows.net", "event_hub", "consumer_group") + .await + .unwrap(); + let lost_partition = ownerships + .iter() + .find(|o| o.partition_id == "1") + .expect("partition 1 stays in the store"); + assert_eq!( + lost_partition.etag, winner_b[0].etag, + "the loser must not overwrite the winner" + ); + assert_eq!(lost_partition.owner_id, Some("owner-b".to_string())); +} + +/// A validation failure is a different outcome from a lost claim, so +/// `claim_ownership` must still return an error for a record it cannot use. +#[tokio::test] +async fn test_claim_ownership_invalid_ownership_still_errors() { + common::setup(); + let store = InMemoryCheckpointStore::new(); + let ownership = Ownership { + fully_qualified_namespace: "fqdn.servicebus.windows.net".to_string(), + partition_id: "partition_id".to_string(), + owner_id: Some("owner_id".to_string()), + etag: Some("etag".into()), + ..Default::default() + }; + let result = store.claim_ownership(&[ownership]).await; + assert!(result.is_err(), "a validation failure is not a lost claim"); + assert_eq!(*result.unwrap_err().kind(), AzureErrorKind::Other); +} + #[tokio::test] async fn test_update_checkpoint() { common::setup(); From 32d68607cfd6fa0edc5b4eabd28276117ac2244b Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Thu, 20 Aug 2026 14:36:31 -0400 Subject: [PATCH 2/2] fix(eventhubs): treat a lost ownership claim as a normal outcome `InMemoryCheckpointStore::claim_ownership` returned an error when the caller presented a stale ETag. One lost partition therefore cancelled the whole batch: the error reached `EventProcessor::run`, which stopped every partition the instance owned and discarded the new ETag for each partition that the same batch had already renewed. A private `try_update_ownership` helper now reports a stale ETag as `Ok(None)` and logs the conflict at the debug level with the `claim-conflict` event field that `BlobCheckpointStore` uses. `claim_ownership` skips that partition and keeps the claims it made. A store failure still leaves the batch as an error. The public `update_ownership` keeps its error contract for a stale ETag, and the `CheckpointStore` trait doc now states that a lost claim is not an error. --- .../azure_messaging_eventhubs/CHANGELOG.md | 1 + .../src/event_processor/mod.rs | 9 +++- .../src/in_memory_checkpoint_store.rs | 52 ++++++++++++++----- 3 files changed, 46 insertions(+), 16 deletions(-) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md b/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md index 8575377c254..31664f0a4f4 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md +++ b/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md @@ -26,6 +26,7 @@ - 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)) - `InMemoryCheckpointStore` now rotates the ETag and refreshes `last_modified_time` when an existing ownership is renewed, matching the create path and the production `BlobCheckpointStore`. Previously the renewal path reinserted the caller's record verbatim, leaving a stale ETag and timestamp; that divergence from the real store could mask bugs in code that relies on ETag rotation for optimistic concurrency. ([#4594](https://github.com/Azure/azure-sdk-for-rust/issues/4594)) +- `InMemoryCheckpointStore::claim_ownership` now treats a stale ETag as a lost claim instead of an error, the same way `BlobCheckpointStore` does. It skips that partition, returns the partitions it did claim, and logs the conflict at the debug level. A single lost claim used to return an error that reached `EventProcessor::run`, which stopped every partition the instance owned and discarded the new ETag for each partition that the same batch had already renewed. ([#5095](https://github.com/Azure/azure-sdk-for-rust/issues/5095)) ### Other Changes diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/event_processor/mod.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/event_processor/mod.rs index 659c36d3643..388e87e54f0 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/event_processor/mod.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/event_processor/mod.rs @@ -25,10 +25,15 @@ pub trait CheckpointStore: Send + Sync { /// * `ownerships` - A vector of `Ownership` objects representing the partitions to claim. /// /// # Returns - /// A vector of claimed `Ownership` objects. + /// A vector of the `Ownership` objects that the caller claimed. Each record + /// carries the new ETag for the next claim. The vector holds fewer records + /// than the `ownerships` argument when another owner holds a partition. /// /// # Errors - /// Returns an error if the ownership claim fails. + /// Returns an error only when the store fails, for example an + /// authentication failure or a missing container. A lost claim is not an + /// error. The implementation must omit that partition from the returned + /// vector and continue with the other partitions. /// async fn claim_ownership(&self, ownerships: &[Ownership]) -> Result>; diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/in_memory_checkpoint_store.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/in_memory_checkpoint_store.rs index 73bd1133f15..d5c97bfb91a 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/in_memory_checkpoint_store.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/in_memory_checkpoint_store.rs @@ -7,7 +7,7 @@ use azure_core::{ }; use std::collections::HashMap; use std::sync::{Arc, Mutex}; -use tracing::{error, trace}; +use tracing::{debug, error, trace}; /// An in-memory checkpoint store for Event Hubs. /// This store is used to manage checkpoints and ownerships in memory. @@ -51,7 +51,35 @@ impl InMemoryCheckpointStore { /// `last_modified_time`, for a renewal and for a first claim. A renewal /// makes the caller's ETag stale, so the caller must keep the returned /// record for its next claim. + /// + /// A stale ETag returns an error. [`CheckpointStore::claim_ownership`] + /// reports the same condition as a lost claim instead. pub fn update_ownership(&self, ownership: &Ownership) -> Result { + if let Some(updated_ownership) = self.try_update_ownership(ownership)? { + return Ok(updated_ownership); + } + + let key = Ownership::get_ownership_name( + &ownership.fully_qualified_namespace, + &ownership.event_hub_name, + &ownership.consumer_group, + &ownership.partition_id, + )?; + error!( + partition_id = %ownership.partition_id, + expected_etag = ?ownership.etag, + "ETag mismatch claiming ownership for key {}", + key + ); + Err(Error::with_message( + AzureErrorKind::Other, + format!("ETag mismatch for partition {key}"), + )) + } + + /// Updates the ownership for a specific partition, and reports a lost + /// claim as `Ok(None)`. An `Err` is a store failure, not a lost claim. + fn try_update_ownership(&self, ownership: &Ownership) -> Result> { trace!("Update ownership for partition {}", ownership.partition_id); check_non_empty_parameter!(ownership.fully_qualified_namespace); @@ -74,20 +102,15 @@ impl InMemoryCheckpointStore { Some(existing) => { let actual_etag = existing.etag.clone(); if ownership.etag != actual_etag { - // The call returns `Err` from here, so this logs at the - // error level, the same as the other failure path in this - // file. - error!( + debug!( + event = "claim-conflict", partition_id = %ownership.partition_id, expected_etag = ?ownership.etag, actual_etag = ?actual_etag, - "ETag mismatch claiming ownership for key {}", + "Lost ownership claim: ETag mismatch for key {}", key ); - return Err(Error::with_message( - AzureErrorKind::Other, - format!("ETag mismatch for partition {key}"), - )); + return Ok(None); } true } @@ -107,7 +130,7 @@ impl InMemoryCheckpointStore { } else { trace!("Inserted new ownership for key {}", key); } - Ok(updated_ownership) + Ok(Some(updated_ownership)) } } @@ -145,9 +168,10 @@ impl CheckpointStore for InMemoryCheckpointStore { trace!("Claim ownership for {} partitions", ownerships.len()); let mut claimed_ownerships = Vec::new(); for ownership in ownerships { - let ownership = self.update_ownership(ownership)?; - if ownership.etag.is_some() { - claimed_ownerships.push(ownership); + // A lost claim is not a failure. Skip that partition and keep the + // claims this batch already made. + if let Some(claimed) = self.try_update_ownership(ownership)? { + claimed_ownerships.push(claimed); } } Ok(claimed_ownerships)