fix(eventhubs): serialize claims-based-security authorizations

sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md

@@ -14,6 +14,7 @@
 
 ### Bugs Fixed
 
+- Claims-based-security authorizations for one connection now run in sequence. The service permits one `$cbs` link for each connection, so a client that attached more than one link at once could fail with `NotAllowed`.
 - Increased `DEFAULT_PARTITION_EXPIRATION_DURATION` from 10 seconds to 60 seconds. The previous default was shorter than `DEFAULT_UPDATE_INTERVAL` (30 seconds), so ownership records expired between load-balancing cycles. The load balancer perpetually saw `current=0` for every consumer and continuously re-claimed partitions, causing widespread duplicate event processing. `EventProcessorBuilder::build` now rejects configurations where `partition_expiration_duration <= update_interval`. ([#3851](https://github.com/Azure/azure-sdk-for-rust/issues/3851))
 - 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.

sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/claims_based_security.rs

@@ -106,13 +106,19 @@ impl AmqpClaimsBasedSecurityApis for RecoverableClaimsBasedSecurity {
                 let secret = secret.clone();
 
                 async move {
-                    let claims_based_security_client = self
+                    let connection = self
                         .recoverable_connection
                         .upgrade()
-                        .ok_or_else(|| AmqpError::with_message("Missing Connection"))?
-                        .ensure_amqp_cbs()
-                        .await
-                        .map_err(|e| {
+                        .ok_or_else(|| AmqpError::with_message("Missing Connection"))?;
+
+                    // The service permits one `$cbs` link for each connection, and
+                    // `ensure_amqp_cbs` attaches a new link for each authorization.
+                    // Hold the lock for the full round trip, so two authorizations
+                    // that start at the same time do not collide with `NotAllowed`.
+                    let _cbs_guard = connection.lock_claims_based_security().await;
+
+                    let claims_based_security_client =
+                        connection.ensure_amqp_cbs().await.map_err(|e| {
                             AmqpError::from(azure_core::Error::with_error(
                                 AzureErrorKind::Other,
                                 e,

sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs

@@ -18,7 +18,7 @@ use crate::{
     producer::DEFAULT_EVENTHUBS_APPLICATION,
     RetryOptions,
 };
-use async_lock::{Mutex as AsyncMutex, OnceCell, RwLock};
+use async_lock::{Mutex as AsyncMutex, MutexGuard as AsyncMutexGuard, OnceCell, RwLock};
 use azure_core::{credentials::TokenCredential, http::Url, time::Duration, Uuid};
 use azure_core_amqp::{
     error::{AmqpErrorCondition, AmqpErrorKind},
@@ -93,6 +93,11 @@ pub(crate) struct RecoverableConnection {
     session_instances: RwLock<HashMap<Url, Arc<OnceCell<Arc<AmqpSession>>>>>,
     receiver_instances: RwLock<HashMap<Url, Arc<OnceCell<Arc<AmqpReceiver>>>>>,
     pub(super) authorizer: Arc<Authorizer>,
+    // The service permits one `$cbs` link for each connection. Every
+    // authorization attaches a link, uses it, and then drops it, so two
+    // authorizations that overlap make the service reject the second one with
+    // `NotAllowed`. This lock keeps them in sequence. See `lock_claims_based_security`.
+    cbs_lock: AsyncMutex<()>,
     connections: AsyncMutex<Option<Arc<AmqpConnection>>>,
     connection_name: String,
     pub(super) retry_options: RetryOptions,
@@ -207,6 +212,7 @@ impl RecoverableConnection {
                 connection_name,
                 custom_endpoint,
                 retry_options,
+                cbs_lock: AsyncMutex::new(()),
                 connections: AsyncMutex::new(None),
                 session_instances: RwLock::new(HashMap::new()),
                 sender_instances: RwLock::new(HashMap::new()),
@@ -644,6 +650,23 @@ impl RecoverableConnection {
         Ok(management_client.clone())
     }
 
+    /// Takes the lock that keeps the claims-based-security round trips of this
+    /// connection in sequence.
+    ///
+    /// The service permits one `$cbs` link for each connection, and it rejects a
+    /// second attach with `NotAllowed`. [`Self::ensure_amqp_cbs`] attaches a new
+    /// link for each authorization, so the caller must hold this lock for the
+    /// full round trip.
+    ///
+    /// Without this lock, the authorizations for different paths overlap when a
+    /// client sets up more than one link at once, for example a buffered
+    /// producer that starts one sender for each partition. The lock covers only
+    /// the authorization. The link attach that follows and the session begin
+    /// stay concurrent.
+    pub(super) async fn lock_claims_based_security(&self) -> AsyncMutexGuard<'_, ()> {
+        self.cbs_lock.lock().await
+    }
+
     /// Ensures that the AMQP Claims-Based Security (CBS) client is created and attached.
     #[instrument(
         level = "debug",
@@ -1703,4 +1726,78 @@ mod tests {
              forever."
         );
     }
+
+    fn cbs_lock_test_connection() -> Arc<RecoverableConnection> {
+        let url = Url::parse("amqps://example.com").unwrap();
+        RecoverableConnection::new(
+            url,
+            None,
+            None,
+            Arc::new(MockCredential),
+            Default::default(),
+            None,
+        )
+    }
+
+    // The service permits one `$cbs` link for each connection, so an
+    // authorization must not start while another one holds the link. A second
+    // caller must wait until the first guard drops. This test needs no network,
+    // because it exercises the lock that `authorize_path` takes.
+    #[tokio::test]
+    async fn cbs_lock_blocks_a_second_caller_until_the_guard_drops() {
+        let connection = cbs_lock_test_connection();
+
+        let guard = connection.lock_claims_based_security().await;
+        assert!(
+            connection.cbs_lock.try_lock().is_none(),
+            "a second caller must not take the lock while the first one holds it"
+        );
+
+        drop(guard);
+        assert!(
+            connection.cbs_lock.try_lock().is_some(),
+            "the lock must be free after the guard drops"
+        );
+    }
+
+    // Regression guard for the `NotAllowed` failure: if a later change moves or
+    // narrows the guard in `authorize_path`, two round trips can overlap again.
+    // Count the callers that hold the lock at the same time, and make sure the
+    // count never goes above one.
+    #[tokio::test]
+    async fn cbs_lock_never_lets_two_callers_overlap() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+
+        let connection = cbs_lock_test_connection();
+        let in_flight = Arc::new(AtomicUsize::new(0));
+        let most_seen = Arc::new(AtomicUsize::new(0));
+
+        let mut tasks = Vec::new();
+        for _ in 0..8 {
+            let connection = connection.clone();
+            let in_flight = in_flight.clone();
+            let most_seen = most_seen.clone();
+            tasks.push(tokio::spawn(async move {
+                let _guard = connection.lock_claims_based_security().await;
+                let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
+                most_seen.fetch_max(now, Ordering::SeqCst);
+                // Give the other tasks a chance to run while this one holds the
+                // lock, which is what a real round trip does at its await points.
+                for _ in 0..4 {
+                    tokio::task::yield_now().await;
+                }
+                in_flight.fetch_sub(1, Ordering::SeqCst);
+            }));
+        }
+
+        for task in tasks {
+            task.await.unwrap();
+        }
+
+        assert_eq!(
+            most_seen.load(Ordering::SeqCst),
+            1,
+            "the claims-based-security round trips of one connection must not overlap"
+        );
+    }
 }

sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer.rs

@@ -4,8 +4,8 @@
 use azure_core::http::StatusCode;
 use azure_core_amqp::{message::AmqpMessageProperties, AmqpError, AmqpList, AmqpSimpleValue};
 use azure_core_test::{recorded, TestContext};
-use azure_messaging_eventhubs::{EventDataBatchOptions, ProducerClient};
-use std::{env, error::Error};
+use azure_messaging_eventhubs::{EventDataBatchOptions, ProducerClient, SendEventOptions};
+use std::{env, error::Error, sync::Arc};
 use tracing::{info, trace};
 
 #[recorded::test(live)]
@@ -490,3 +490,72 @@ async fn send_eventdata_with_connection_string(_ctx: TestContext) -> Result<(),
 
     Ok(())
 }
+
+/// Sends to every partition at the same time from one client.
+///
+/// Each send attaches a sender, and each attach needs a claims-based-security
+/// authorization. The service permits one `$cbs` link for each connection, so
+/// the authorizations must run in sequence. Before that fix, the service
+/// answered the overlapping authorizations with `NotAllowed` ("A link to
+/// connection ... $cbs node has already been opened"), which the client
+/// classifies as not retryable, and the sends failed.
+#[recorded::test(live)]
+async fn send_to_every_partition_at_once(ctx: TestContext) -> Result<(), Box<dyn Error>> {
+    let recording = ctx.recording();
+
+    let host = env::var("EVENTHUBS_HOST")?;
+    let eventhub = env::var("EVENTHUB_NAME")?;
+
+    let client = Arc::new(
+        ProducerClient::builder()
+            .with_application_id("send_to_every_partition_at_once".to_string())
+            .open(host.as_str(), eventhub.as_str(), recording.credential())
+            .await?,
+    );
+
+    // Read the partitions first, so the connection is open and only the sender
+    // attaches overlap.
+    let partitions = client.get_eventhub_properties().await?.partition_ids;
+    assert!(partitions.len() > 1, "the test needs many partitions");
+    info!("Send to {} partitions at the same time.", partitions.len());
+
+    let mut tasks = Vec::new();
+    for partition in partitions.iter() {
+        let client = client.clone();
+        let partition = partition.clone();
+        tasks.push(tokio::spawn(async move {
+            let result = client
+                .send_event(
+                    format!("Hello, partition {partition}!"),
+                    Some(SendEventOptions {
+                        partition_id: Some(partition.clone()),
+                    }),
+                )
+                .await;
+            (partition, result)
+        }));
+    }
+
+    let mut failures = Vec::new();
+    for task in tasks {
+        let (partition, result) = task.await?;
+        if let Err(e) = result {
+            info!("Partition {partition} failed. {e:?}");
+            failures.push(partition);
+        }
+    }
+
+    Arc::try_unwrap(client)
+        .map_err(|_| "A task still holds the client.")?
+        .close()
+        .await?;
+
+    assert!(
+        failures.is_empty(),
+        "{} of {} sends failed: {failures:?}",
+        failures.len(),
+        partitions.len()
+    );
+
+    Ok(())
+}