Cleaned up eventhubs disconnect operations (Azure#3038)

Partial fix for Azure#2944

The errors generated come from dropping a connection without first
closing the connection (and sessions, etc on the connection).

This PR makes the connection close logic much more deterministic and
close operations are now much more likely to succeed.

Author: LarryOsterman

sdk/core/azure_core_amqp/src/fe2o3/management.rs

@@ -86,6 +86,9 @@ impl AmqpManagementApis for Fe2o3AmqpManagement {
             .ok_or_else(Self::amqp_management_not_attached)?;
         let management = management.into_inner();
         management.close().await.map_err(AmqpError::from)?;
+
+        let mut session = self.session.lock().await;
+        session.end().await.map_err(AmqpError::from)?;
         Ok(())
     }
 

sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md

@@ -6,10 +6,14 @@
 
 ### Breaking Changes
 
+- `EventProcessor` now consumes its `ConsumerClient` parameter rather than accepting a clone of an `Arc`.
+
 ### Bugs Fixed
 
 ### Other Changes
 
+- Internal refactoring to ensure that the `close()` method on various clients works as expected.
+
 ## 0.7.0 (2025-09-16)
 
 ### Features Added

sdk/eventhubs/azure_messaging_eventhubs/examples/eventhubs_processor_client.rs

@@ -30,7 +30,7 @@ impl BackgroundProcessor {
         // Start the background task to process events
         let mut task = self.background_task.lock().await;
         //        let mut task = task.as_mut();
-        let processor = Arc::clone(&self.processor);
+        let processor = self.processor.clone();
         task.replace(tokio::spawn(async move { processor.run().await }));
         Ok(())
     }
@@ -58,21 +58,19 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let eventhub_namespace = std::env::var("EVENTHUBS_HOST")?;
     let eventhub_name = std::env::var("EVENTHUB_NAME")?;
 
-    let consumer = Arc::new(
-        ConsumerClient::builder()
-            .open(
-                eventhub_namespace.as_str(),
-                eventhub_name,
-                DeveloperToolsCredential::new(None)?.clone(),
-            )
-            .await?,
-    );
+    let consumer = ConsumerClient::builder()
+        .open(
+            eventhub_namespace.as_str(),
+            eventhub_name,
+            DeveloperToolsCredential::new(None)?.clone(),
+        )
+        .await?;
 
     println!("Opened consumer client");
 
     let checkpoint_store = Arc::new(InMemoryCheckpointStore::new());
     let processor = EventProcessor::builder()
-        .build(consumer.clone(), checkpoint_store)
+        .build(consumer, checkpoint_store)
         .await?;
     let background_processor = BackgroundProcessor::new(processor.clone()).await?;
 
@@ -116,5 +114,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     background_processor.stop().await?;
     println!("Stopped background processor");
 
+    // Close the processor
+    if let Ok(processor) = Arc::try_unwrap(processor) {
+        processor.close().await?;
+        println!("Closed processor");
+    } else {
+        println!("Processor still has references, not closing.");
+    }
+
     Ok(())
 }

sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs

@@ -566,9 +566,9 @@ mod tests {
         let final_count = mock_credential.get_token_get_count();
         debug!("After sleeping, token count: {final_count}");
 
-        assert_eq!(
-            final_count, 2,
-            "Expected token get count to be 2, but got {final_count}"
+        assert!(
+            final_count >= 2,
+            "Expected token get count to be greater or equal to 2, but got {final_count}"
         );
         info!("Final token get count: {final_count}");
         Ok(())
@@ -650,7 +650,7 @@ mod tests {
         debug!("After sleeping the first time, token count: {final_count}");
         assert!(
             final_count >= 3,
-            "Expected first get token count to be 3, but got {final_count}"
+            "Expected first get token count to be at least 3, but got {final_count}"
         );
 
         info!("First token expiration get count: {}", final_count);

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

@@ -22,8 +22,9 @@ use azure_core::{credentials::TokenCredential, http::Url, time::Duration, Result
 use azure_core_amqp::{
     error::{AmqpErrorCondition, AmqpErrorKind},
     AmqpClaimsBasedSecurity, AmqpConnection, AmqpConnectionApis, AmqpConnectionOptions, AmqpError,
-    AmqpManagement, AmqpReceiver, AmqpReceiverApis, AmqpReceiverOptions, AmqpSender,
-    AmqpSenderApis, AmqpSession, AmqpSessionApis, AmqpSessionOptions, AmqpSource, AmqpSymbol,
+    AmqpManagement, AmqpManagementApis, AmqpReceiver, AmqpReceiverApis, AmqpReceiverOptions,
+    AmqpSender, AmqpSenderApis, AmqpSession, AmqpSessionApis, AmqpSessionOptions, AmqpSource,
+    AmqpSymbol,
 };
 #[cfg(test)]
 use std::sync::Mutex;
@@ -73,12 +74,12 @@ pub(crate) struct RecoverableConnection {
     pub(super) url: Url,
     application_id: Option<String>,
     custom_endpoint: Option<Url>,
-    connections: AsyncMutex<Option<Arc<AmqpConnection>>>,
     mgmt_client: AsyncMutex<Option<Arc<AmqpManagement>>>,
-    session_instances: AsyncMutex<HashMap<Url, Arc<AmqpSession>>>,
     receiver_instances: AsyncMutex<HashMap<Url, Arc<AmqpReceiver>>>,
     sender_instances: AsyncMutex<HashMap<Url, Arc<AmqpSender>>>,
+    session_instances: AsyncMutex<HashMap<Url, Arc<AmqpSession>>>,
     pub(super) authorizer: Arc<Authorizer>,
+    connections: AsyncMutex<Option<Arc<AmqpConnection>>>,
     connection_name: String,
     pub(super) retry_options: RetryOptions,
 
@@ -161,10 +162,80 @@ impl RecoverableConnection {
     /// This method will close the underlying AMQP connection, if it exists. It will also cause all outstanding sends and receives
     /// to complete with an error.
     ///
-    pub(crate) async fn close_connection(&self) -> Result<()> {
-        let connection = self.ensure_connection().await?;
+    pub(crate) async fn close_connection(self) -> Result<()> {
+        trace!("Closing recoverable connection for {}.", self.url);
 
-        connection.close().await
+        let mut management_client = self.mgmt_client.lock().await;
+        if let Some(management_client) = management_client.take() {
+            trace!("Closing management client for {}.", self.url);
+            if let Ok(management_client) = Arc::try_unwrap(management_client) {
+                trace!("Detaching management client for {}.", self.url);
+                management_client.detach().await?;
+            } else {
+                trace!(
+                    "Failed to detach management client for {}, references exist.",
+                    self.url
+                );
+            }
+        }
+
+        let mut sender_instances = self.sender_instances.lock().await;
+        for (path, sender) in sender_instances.drain() {
+            trace!("Detaching sender for path {}.", path);
+            if let Ok(sender) = Arc::try_unwrap(sender) {
+                trace!("Detaching sender for path {}.", path);
+                sender.detach().await?;
+            } else {
+                trace!(
+                    "Failed to detach sender for path {}, references exist.",
+                    path
+                );
+            }
+        }
+
+        let mut receiver_instances = self.receiver_instances.lock().await;
+        for (source_url, receiver) in receiver_instances.drain() {
+            trace!("Detaching receiver for source URL {}.", source_url);
+            if let Ok(receiver) = Arc::try_unwrap(receiver) {
+                trace!("Detaching receiver for source URL {}.", source_url);
+                receiver.detach().await?;
+            } else {
+                trace!(
+                    "Failed to detach receiver for source URL {}, references exist.",
+                    source_url
+                );
+            }
+        }
+
+        let mut session_instances = self.session_instances.lock().await;
+        for (session_id, session) in session_instances.drain() {
+            trace!("Detaching session for ID {}.", session_id);
+            if let Ok(session) = Arc::try_unwrap(session) {
+                session.end().await?;
+            } else {
+                trace!(
+                    "Failed to detach session for ID {}, references exist.",
+                    session_id
+                );
+            }
+        }
+
+        if let Some(connection) = self.connections.lock().await.take() {
+            trace!("Closing connection for {}.", self.url);
+            if let Ok(connection) = Arc::try_unwrap(connection) {
+                trace!(
+                    "No references, actually closing connection for {}.",
+                    self.url
+                );
+                connection.close().await?;
+            } else {
+                trace!(
+                    "Failed to close connection for {}, references exist.",
+                    self.url
+                );
+            }
+        }
+        Ok(())
     }
 
     /// Ensures that the connection to the Event Hubs service is established.
@@ -482,8 +553,7 @@ impl RecoverableConnection {
                 }
             }
             AmqpErrorKind::ConnectionClosedByRemote(_)
-            | AmqpErrorKind::ConnectionDetachedByRemote(_)
-            | AmqpErrorKind::ConnectionDropped(_) => {
+            | AmqpErrorKind::ConnectionDetachedByRemote(_) => {
                 debug!("Connection dropped error: {}", amqp_error);
                 ErrorRecoveryAction::ReconnectConnection
             }
@@ -535,6 +605,12 @@ impl RecoverableConnection {
     }
 }
 
+impl Drop for RecoverableConnection {
+    fn drop(&mut self) {
+        trace!("Dropping RecoverableConnection for {}", self.url);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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

@@ -65,6 +65,12 @@ impl RecoverableReceiver {
     }
 }
 
+impl Drop for RecoverableReceiver {
+    fn drop(&mut self) {
+        debug!("Dropping RecoverableReceiver for {}", self.source_url);
+    }
+}
+
 #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
 #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
 impl AmqpReceiverApis for RecoverableReceiver {
@@ -93,24 +99,34 @@ impl AmqpReceiverApis for RecoverableReceiver {
     }
 
     async fn receive_delivery(&self) -> azure_core::Result<azure_core_amqp::AmqpDelivery> {
+        let retry_options = {
+            self.recoverable_connection
+                .upgrade()
+                .ok_or_else(|| EventHubsError::from(ErrorKind::MissingConnection))?
+                .retry_options
+                .clone()
+        };
         let delivery = recover_azure_operation(
             || async move {
-                let connection = self
-                    .recoverable_connection
-                    .upgrade()
-                    .ok_or_else(|| EventHubsError::from(ErrorKind::MissingConnection))?;
-
-                // Check for forced error.
-                #[cfg(test)]
-                connection.get_forced_error()?;
-
-                let receiver = connection
-                    .ensure_receiver(
-                        &self.source_url,
-                        &self.message_source,
-                        &self.receiver_options,
-                    )
-                    .await?;
+                debug!("Starting receive_delivery operation");
+                let receiver = {
+                    let connection = self
+                        .recoverable_connection
+                        .upgrade()
+                        .ok_or_else(|| EventHubsError::from(ErrorKind::MissingConnection))?;
+
+                    // Check for forced error.
+                    #[cfg(test)]
+                    connection.get_forced_error()?;
+
+                    connection
+                        .ensure_receiver(
+                            &self.source_url,
+                            &self.message_source,
+                            &self.receiver_options,
+                        )
+                        .await?
+                };
                 if let Some(delivery_timeout) = self.timeout {
                     select! {
                         delivery = receiver.receive_delivery().fuse() => Ok(delivery),
@@ -124,11 +140,7 @@ impl AmqpReceiverApis for RecoverableReceiver {
                     receiver.receive_delivery().await
                 }
             },
-            &self
-                .recoverable_connection
-                .upgrade()
-                .ok_or_else(|| EventHubsError::from(ErrorKind::MissingConnection))?
-                .retry_options,
+            &retry_options,
             Self::should_retry_receive_operation,
             Some(move |connection: Weak<RecoverableConnection>, reason| {
                 let connection = connection.clone();

sdk/eventhubs/azure_messaging_eventhubs/src/consumer/event_receiver.rs

@@ -144,3 +144,9 @@ impl EventReceiver {
         self.connection.close_receiver(&self.source_url).await
     }
 }
+
+impl Drop for EventReceiver {
+    fn drop(&mut self) {
+        trace!("Dropping EventReceiver for partition {}", self.partition_id);
+    }
+}

sdk/eventhubs/azure_messaging_eventhubs/src/consumer/mod.rs

@@ -148,7 +148,19 @@ impl ConsumerClient {
     /// ```
     pub async fn close(self) -> Result<()> {
         trace!("Closing consumer client for {}.", self.endpoint);
-        self.recoverable_connection.close_connection().await
+        let recoverable_connection =
+            Arc::try_unwrap(self.recoverable_connection).map_err(|_| {
+                Error::with_message(
+                    AzureErrorKind::Other,
+                    "Could not close consumer recoverable connection, multiple references exist",
+                )
+            })?;
+        trace!(
+            "No references to connection, closing connection for {}.",
+            self.endpoint
+        );
+        recoverable_connection.close_connection().await?;
+        Ok(())
     }
 
     /// Forces an error on the connection.
@@ -890,6 +902,12 @@ pub(crate) mod tests {
         )
         .await?;
 
+        if let Ok(consumer) = Arc::try_unwrap(consumer) {
+            consumer.close().await?;
+        } else {
+            panic!("Consumer client has unresolved references.");
+        }
+
         Ok(())
     }
 
@@ -935,6 +953,12 @@ pub(crate) mod tests {
         )
         .await?;
 
+        if let Ok(consumer) = Arc::try_unwrap(consumer) {
+            consumer.close().await?;
+        } else {
+            panic!("Consumer client has unresolved references.");
+        }
+
         Ok(())
     }
     #[recorded::test(live)]