fix(eventhubs): stop authorization refresh on close

sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md

@@ -18,6 +18,7 @@
 
 ### Bugs Fixed
 
+- `ProducerClient::close` now stops the authorization refresh task, so repeated producer life cycles release task-held memory. ([#4595](https://github.com/Azure/azure-sdk-for-rust/issues/4595))
 - `ConsumerClient::close` and `ProducerClient::close` now close the connection when another object still holds it, most often an `EventReceiver` that the caller has not dropped. Both methods used to report an error and leave the connection open. ([#4931](https://github.com/Azure/azure-sdk-for-rust/issues/4931))
 - A handle that outlives the client it came from now reports that the client is closed on its next call. Such a handle opened a second connection to the service before. ([#4931](https://github.com/Azure/azure-sdk-for-rust/issues/4931))
 - `EventProcessor::close` now continues past a partition client that the application still holds. It used to stop there, which left the partition clients behind it open and skipped the close of the consumer client. ([#4931](https://github.com/Azure/azure-sdk-for-rust/issues/4931))

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

@@ -18,7 +18,7 @@ use azure_core_amqp::{AmqpClaimsBasedSecurityApis as _, AmqpError};
 use rand::{rng, RngExt};
 use std::{
     collections::{HashMap, HashSet},
-    sync::{Arc, Mutex as SyncMutex, OnceLock, Weak},
+    sync::{Arc, Mutex as SyncMutex, Weak},
 };
 use tracing::{debug, error, info, trace, warn};
 
@@ -71,9 +71,15 @@ enum RefreshPass {
     Stop,
 }
 
+enum AuthorizationRefresherState {
+    NotStarted,
+    Running(SpawnedTask),
+    Stopped,
+}
+
 pub(crate) struct Authorizer {
     authorization_scopes: RwLock<HashMap<Url, AccessToken>>,
-    authorization_refresher: OnceLock<SpawnedTask>,
+    authorization_refresher: SyncMutex<AuthorizationRefresherState>,
     /// Bias to apply to token refresh time. This determines how much time we will refresh the token before it expires.
     token_refresh_bias: SyncMutex<TokenRefreshTimes>,
     credential: Arc<dyn TokenCredential>,
@@ -86,9 +92,6 @@ pub(crate) struct Authorizer {
     disable_authorization: SyncMutex<bool>,
 }
 
-unsafe impl Send for Authorizer {}
-unsafe impl Sync for Authorizer {}
-
 impl Authorizer {
     /// Creates an authorizer. `cbs_token_type` is `None` for JWT/Entra
     /// credentials and `Some("servicebus.windows.net:sastoken")` for SAS
@@ -99,7 +102,7 @@ impl Authorizer {
         cbs_token_type: Option<&'static str>,
     ) -> Self {
         Self {
-            authorization_refresher: OnceLock::new(),
+            authorization_refresher: SyncMutex::new(AuthorizationRefresherState::NotStarted),
             authorization_scopes: RwLock::new(HashMap::new()),
             token_refresh_bias: SyncMutex::new(TokenRefreshTimes::default()),
             credential,
@@ -116,8 +119,28 @@ impl Authorizer {
         scopes.clear();
     }
 
+    pub(crate) async fn stop_refresh_task(&self) {
+        let task = {
+            let mut state = self
+                .authorization_refresher
+                .lock()
+                .unwrap_or_else(|poisoned| poisoned.into_inner());
+            match std::mem::replace(&mut *state, AuthorizationRefresherState::Stopped) {
+                AuthorizationRefresherState::Running(task) => Some(task),
+                AuthorizationRefresherState::NotStarted | AuthorizationRefresherState::Stopped => {
+                    None
+                }
+            }
+        };
+
+        if let Some(task) = task {
+            task.abort();
+            let _ = task.await;
+        }
+    }
+
     #[cfg(test)]
-    fn disable_authorization(&self) -> Result<()> {
+    pub(crate) fn disable_authorization(&self) -> Result<()> {
         use crate::EventHubsError;
 
         let mut disable_authorization = self
@@ -226,12 +249,18 @@ impl Authorizer {
                 continue;
             };
 
-            self.authorization_refresher.get_or_init(|| {
+            let mut state = self
+                .authorization_refresher
+                .lock()
+                .unwrap_or_else(|poisoned| poisoned.into_inner());
+            if matches!(&*state, AuthorizationRefresherState::NotStarted) {
                 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()))
-            });
+                *state = AuthorizationRefresherState::Running(
+                    async_runtime.spawn(Box::pin(self_clone.refresh_tokens_task())),
+                );
+            }
 
             return Ok(stored);
         }
@@ -650,6 +679,17 @@ impl Authorizer {
         *token_refresh_bias = refresh_times;
         Ok(())
     }
+
+    #[cfg(test)]
+    pub(crate) fn set_token_refresh_bias_for_test(&self, bias: Duration) -> Result<()> {
+        self.set_token_refresh_times(TokenRefreshTimes {
+            before_expiration_refresh_time: bias,
+            jitter_min: Duration::milliseconds(0),
+            // The upper bound is exclusive. This range produces zero jitter
+            // while avoiding an empty random range in the refresh loop.
+            jitter_max: Duration::milliseconds(1),
+        })
+    }
 }
 
 #[cfg(test)]

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

@@ -427,6 +427,8 @@ impl RecoverableConnection {
         // the client.
         self.closed.store(true, Ordering::Release);
 
+        self.authorizer.stop_refresh_task().await;
+
         // Swap the cell out under the write lock, then detach without holding
         // it. The guard is a separate binding so the lock scope is visible and
         // a debugger can read it.
@@ -1508,9 +1510,17 @@ impl Drop for RecoverableConnection {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use azure_core::http::Url;
+    use azure_core::{
+        credentials::{AccessToken, TokenCredential, TokenRequestOptions},
+        http::Url,
+        time::{Duration, OffsetDateTime},
+    };
     use azure_core_test::credentials::MockCredential;
-    use std::sync::Arc;
+    use std::sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    };
+    use tokio::sync::Notify;
 
     // A close does not need exclusive ownership of the connection.
     //
@@ -1551,6 +1561,153 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn close_stops_owned_authorization_refresh_task() {
+        #[derive(Debug)]
+        struct GatedCredential {
+            requests: AtomicUsize,
+            entered_refresh: Notify,
+            release_refresh: Notify,
+        }
+
+        #[async_trait::async_trait]
+        impl TokenCredential for GatedCredential {
+            async fn get_token(
+                &self,
+                _scopes: &[&str],
+                _options: Option<TokenRequestOptions<'_>>,
+            ) -> azure_core::Result<AccessToken> {
+                match self.requests.fetch_add(1, Ordering::SeqCst) {
+                    0 => Ok(AccessToken::new(
+                        azure_core::credentials::Secret::new("initial_token"),
+                        OffsetDateTime::now_utc() + Duration::hours(1),
+                    )),
+                    1 => {
+                        self.entered_refresh.notify_one();
+                        self.release_refresh.notified().await;
+                        Ok(AccessToken::new(
+                            azure_core::credentials::Secret::new("refreshed_token"),
+                            OffsetDateTime::now_utc() + Duration::hours(1),
+                        ))
+                    }
+                    request => unreachable!("unexpected token request {request}"),
+                }
+            }
+        }
+
+        let credential = Arc::new(GatedCredential {
+            requests: AtomicUsize::new(0),
+            entered_refresh: Notify::new(),
+            release_refresh: Notify::new(),
+        });
+        let connection = RecoverableConnection::new(
+            Url::parse("amqps://example.com").unwrap(),
+            None,
+            None,
+            AmqpTransport::default(),
+            credential.clone(),
+            Default::default(),
+            None,
+        );
+        let authorizer = connection.authorizer.clone();
+        authorizer.disable_authorization().unwrap();
+        authorizer
+            .set_token_refresh_bias_for_test(Duration::hours(2))
+            .unwrap();
+
+        let path = Url::parse("amqps://example.com/close_refresh_task").unwrap();
+        authorizer.authorize_path(&connection, &path).await.unwrap();
+
+        // The second request proves that the refresher holds an Arc to the authorizer.
+        credential.entered_refresh.notified().await;
+
+        connection.close_connection().await.unwrap();
+        assert_eq!(
+            Arc::strong_count(&authorizer),
+            2,
+            "the connection and test authorizer references must remain after close"
+        );
+        drop(connection);
+        assert_eq!(
+            Arc::strong_count(&authorizer),
+            1,
+            "authorization refresh task remained alive after close; strong_count={}",
+            Arc::strong_count(&authorizer)
+        );
+    }
+
+    #[tokio::test]
+    async fn close_racing_first_authorization_does_not_start_refresher() {
+        #[derive(Debug)]
+        struct GatedCredential {
+            requests: AtomicUsize,
+            entered: Notify,
+            release: Notify,
+        }
+
+        #[async_trait::async_trait]
+        impl TokenCredential for GatedCredential {
+            async fn get_token(
+                &self,
+                _scopes: &[&str],
+                _options: Option<TokenRequestOptions<'_>>,
+            ) -> azure_core::Result<AccessToken> {
+                if self.requests.fetch_add(1, Ordering::SeqCst) == 0 {
+                    self.entered.notify_one();
+                    self.release.notified().await;
+                }
+                Ok(AccessToken::new(
+                    azure_core::credentials::Secret::new("initial_token"),
+                    OffsetDateTime::now_utc() + Duration::hours(1),
+                ))
+            }
+        }
+
+        let credential = Arc::new(GatedCredential {
+            requests: AtomicUsize::new(0),
+            entered: Notify::new(),
+            release: Notify::new(),
+        });
+        let connection = RecoverableConnection::new(
+            Url::parse("amqps://example.com").unwrap(),
+            None,
+            None,
+            AmqpTransport::default(),
+            credential.clone(),
+            Default::default(),
+            None,
+        );
+        let authorizer = connection.authorizer.clone();
+        authorizer.disable_authorization().unwrap();
+
+        let path = Url::parse("amqps://example.com/close_first_authorization").unwrap();
+        let authorization = {
+            let authorizer = authorizer.clone();
+            let connection = connection.clone();
+            tokio::spawn(async move { authorizer.authorize_path(&connection, &path).await })
+        };
+
+        credential.entered.notified().await;
+        connection.close_connection().await.unwrap();
+        credential.release.notify_one();
+        authorization
+            .await
+            .expect("authorize_path task panicked")
+            .expect("authorize_path returned an error");
+
+        assert_eq!(
+            credential.requests.load(Ordering::SeqCst),
+            1,
+            "the first authorization must make one token request"
+        );
+        drop(connection);
+        assert_eq!(
+            Arc::strong_count(&authorizer),
+            1,
+            "authorization refresh task must not start after close"
+        );
+    }
+
     // The RecoverableConnection implementation uses a UUID to identify connections unless an application ID is provided.
     // This test verifies that a new recoverable connection uses a UUID for its connection ID when no application ID is specified.
     // It also verifies that the connections aren't initialized during construction - they're created on-demand.