From 152d27942ab2d4bdc16eeec78426fa6678448a4b Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Mon, 24 Aug 2026 13:14:17 -0400 Subject: [PATCH 1/9] test: add producer lifecycle memory trend test Track live allocation bytes across sequential producer lifecycle tests. Use warmups, block medians, and a noise-adjusted regression bound. --- .../tests/eventhubs_producer_memory.rs | 156 ++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs diff --git a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs new file mode 100644 index 00000000000..2f978f6410c --- /dev/null +++ b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs @@ -0,0 +1,156 @@ +// Copyright (c) Microsoft Corporation. All Rights reserved +// Licensed under the MIT license. + +use azure_core_test::recorded; +use azure_identity::DeveloperToolsCredential; +use azure_messaging_eventhubs::ProducerClient; +use std::{ + alloc::{GlobalAlloc, Layout, System}, + env, + error::Error, + sync::atomic::{AtomicUsize, Ordering}, +}; + +struct LiveByteAllocator; + +static LIVE_BYTES: AtomicUsize = AtomicUsize::new(0); + +#[global_allocator] +static ALLOCATOR: LiveByteAllocator = LiveByteAllocator; + +unsafe impl GlobalAlloc for LiveByteAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let pointer = unsafe { System.alloc(layout) }; + if !pointer.is_null() { + LIVE_BYTES.fetch_add(layout.size(), Ordering::Relaxed); + } + pointer + } + + unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) { + unsafe { System.dealloc(pointer, layout) }; + LIVE_BYTES.fetch_sub(layout.size(), Ordering::Relaxed); + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let pointer = unsafe { System.alloc_zeroed(layout) }; + if !pointer.is_null() { + LIVE_BYTES.fetch_add(layout.size(), Ordering::Relaxed); + } + pointer + } + + unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_pointer = unsafe { System.realloc(pointer, layout, new_size) }; + if !new_pointer.is_null() { + match new_size.cmp(&layout.size()) { + std::cmp::Ordering::Greater => { + LIVE_BYTES.fetch_add(new_size - layout.size(), Ordering::Relaxed); + } + std::cmp::Ordering::Less => { + LIVE_BYTES.fetch_sub(layout.size() - new_size, Ordering::Relaxed); + } + std::cmp::Ordering::Equal => {} + } + } + new_pointer + } +} + +#[recorded::test(live)] +async fn producer_lifecycle_send_close_has_no_sustained_heap_trend() -> Result<(), Box> { + let host = match env::var("EVENTHUBS_HOST") { + Ok(host) if !host.is_empty() => host, + _ => return Ok(()), + }; + let eventhub = match env::var("EVENT_HUB_NAME") { + Ok(eventhub) if !eventhub.is_empty() => eventhub, + _ => match env::var("EVENTHUB_NAME") { + Ok(eventhub) if !eventhub.is_empty() => eventhub, + _ => return Ok(()), + }, + }; + let credential = match DeveloperToolsCredential::new(None) { + Ok(credential) => credential, + Err(_) => return Ok(()), + }; + + const WARMUP_CYCLES: usize = 5; + const MEASURED_CYCLES: usize = 95; + const BLOCK_SIZE: usize = 5; + const BLOCK_COUNT: usize = MEASURED_CYCLES / BLOCK_SIZE; + const TOTAL_CYCLES: usize = WARMUP_CYCLES + MEASURED_CYCLES; + + let mut samples = [0usize; MEASURED_CYCLES]; + for cycle in 0..TOTAL_CYCLES { + let client = match ProducerClient::builder() + .with_application_id( + "producer_lifecycle_send_close_has_no_sustained_heap_trend".to_string(), + ) + .open(host.as_str(), eventhub.as_str(), credential.clone()) + .await + { + Ok(client) => client, + Err(error) => panic!("producer open failed during cycle {cycle}: {error}"), + }; + + let send_result = client.send_event("lifecycle memory test", None).await; + assert!( + send_result.is_ok(), + "producer send failed during cycle {cycle}: {send_result:?}" + ); + + let close_result = client.close().await; + assert!( + close_result.is_ok(), + "producer close failed during cycle {cycle}: {close_result:?}" + ); + + if cycle >= WARMUP_CYCLES { + samples[cycle - WARMUP_CYCLES] = LIVE_BYTES.load(Ordering::Relaxed); + } + } + + let mut medians = [0.0; BLOCK_COUNT]; + for (block_index, median) in medians.iter_mut().enumerate() { + let start = block_index * BLOCK_SIZE; + let mut block = [0usize; BLOCK_SIZE]; + block.copy_from_slice(&samples[start..start + BLOCK_SIZE]); + block.sort_unstable(); + *median = block[BLOCK_SIZE / 2] as f64; + } + + let x_mean = (BLOCK_COUNT - 1) as f64 / 2.0; + let y_mean = medians.iter().sum::() / BLOCK_COUNT as f64; + let sum_squared_x_deviations = (0..BLOCK_COUNT) + .map(|index| { + let deviation = index as f64 - x_mean; + deviation * deviation + }) + .sum::(); + let slope = medians + .iter() + .enumerate() + .map(|(index, median)| (index as f64 - x_mean) * (median - y_mean)) + .sum::() + / sum_squared_x_deviations; + let intercept = y_mean - slope * x_mean; + let residual_sum_of_squares = medians + .iter() + .enumerate() + .map(|(index, median)| { + let residual = median - (intercept + slope * index as f64); + residual * residual + }) + .sum::(); + let noise = (residual_sum_of_squares / (BLOCK_COUNT - 2) as f64).sqrt(); + let slope_standard_error = noise / sum_squared_x_deviations.sqrt(); + let lower_bound = slope - 3.0 * slope_standard_error; + + assert!( + lower_bound <= 0.0, + "lower_bound <= 0: slope={slope:.3} bytes/block, noise={noise:.3}, lower_bound={lower_bound:.3}, block medians={medians:?}" + ); + + Ok(()) +} From b1ae1e5217444206dbb3356c6a0d9093932762c4 Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Mon, 24 Aug 2026 13:25:48 -0400 Subject: [PATCH 2/9] test: cover closing authorization refresh tasks Gate a refresh credential and verify close releases the authorizer task. --- .../src/common/authorizer.rs | 90 +++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs index ed9e6ad31c4..83951ef18da 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs @@ -1428,4 +1428,94 @@ mod tests { "a token refreshed during recovery must be discarded, not written back" ); } + + #[tokio::test] + async fn close_stops_authorization_refresh_task() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::sync::Notify; + + #[derive(Debug)] + struct GatedTokenCredential { + requests: AtomicUsize, + entered_refresh: Notify, + release_refresh: Notify, + } + + #[async_trait::async_trait] + impl TokenCredential for GatedTokenCredential { + async fn get_token( + &self, + _scopes: &[&str], + _options: Option>, + ) -> Result { + 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(GatedTokenCredential { + requests: AtomicUsize::new(0), + entered_refresh: Notify::new(), + release_refresh: Notify::new(), + }); + let connection = RecoverableConnection::new( + Url::parse("amqps://example.com").unwrap(), + None, + None, + credential.clone(), + Default::default(), + None, + ); + + let authorizer = Arc::new(Authorizer::new( + Arc::downgrade(&connection), + credential.clone(), + None, + )); + authorizer.disable_authorization().unwrap(); + authorizer + .set_token_refresh_times(TokenRefreshTimes { + before_expiration_refresh_time: Duration::hours(2), + jitter_min: Duration::milliseconds(0), + jitter_max: Duration::milliseconds(1), + }) + .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(); + drop(connection); + + let stopped = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if Arc::strong_count(&authorizer) == 1 { + return; + } + tokio::task::yield_now().await; + } + }) + .await; + assert!( + stopped.is_ok(), + "authorization refresh task remained alive after close; strong_count={}", + Arc::strong_count(&authorizer) + ); + } } From a935823fd3f628d7be6a1f91a2490fba762c463d Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Mon, 24 Aug 2026 13:38:50 -0400 Subject: [PATCH 3/9] test: move refresh close test to connection Exercise the connection-owned authorizer with a gated refresh credential. Expose only test-gated helpers for deterministic refresh timing and setup. --- .../src/common/authorizer.rs | 103 ++---------------- .../src/common/recoverable/connection.rs | 85 +++++++++++++++ 2 files changed, 97 insertions(+), 91 deletions(-) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs index 83951ef18da..73e3108827b 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs @@ -117,7 +117,7 @@ impl Authorizer { } #[cfg(test)] - fn disable_authorization(&self) -> Result<()> { + pub(crate) fn disable_authorization(&self) -> Result<()> { use crate::EventHubsError; let mut disable_authorization = self @@ -639,6 +639,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)] @@ -1428,94 +1439,4 @@ mod tests { "a token refreshed during recovery must be discarded, not written back" ); } - - #[tokio::test] - async fn close_stops_authorization_refresh_task() { - use std::sync::atomic::{AtomicUsize, Ordering}; - use tokio::sync::Notify; - - #[derive(Debug)] - struct GatedTokenCredential { - requests: AtomicUsize, - entered_refresh: Notify, - release_refresh: Notify, - } - - #[async_trait::async_trait] - impl TokenCredential for GatedTokenCredential { - async fn get_token( - &self, - _scopes: &[&str], - _options: Option>, - ) -> Result { - 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(GatedTokenCredential { - requests: AtomicUsize::new(0), - entered_refresh: Notify::new(), - release_refresh: Notify::new(), - }); - let connection = RecoverableConnection::new( - Url::parse("amqps://example.com").unwrap(), - None, - None, - credential.clone(), - Default::default(), - None, - ); - - let authorizer = Arc::new(Authorizer::new( - Arc::downgrade(&connection), - credential.clone(), - None, - )); - authorizer.disable_authorization().unwrap(); - authorizer - .set_token_refresh_times(TokenRefreshTimes { - before_expiration_refresh_time: Duration::hours(2), - jitter_min: Duration::milliseconds(0), - jitter_max: Duration::milliseconds(1), - }) - .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(); - drop(connection); - - let stopped = tokio::time::timeout(std::time::Duration::from_secs(1), async { - loop { - if Arc::strong_count(&authorizer) == 1 { - return; - } - tokio::task::yield_now().await; - } - }) - .await; - assert!( - stopped.is_ok(), - "authorization refresh task remained alive after close; strong_count={}", - Arc::strong_count(&authorizer) - ); - } } 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 678516feeb9..f256a2abb81 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs @@ -1540,6 +1540,91 @@ mod tests { ); } + #[tokio::test] + async fn close_stops_owned_authorization_refresh_task() { + use azure_core::{ + credentials::{AccessToken, TokenCredential, TokenRequestOptions}, + time::OffsetDateTime, + }; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::sync::Notify; + + #[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>, + ) -> azure_core::Result { + 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, + 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(); + drop(connection); + + let stopped = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if Arc::strong_count(&authorizer) == 1 { + return; + } + tokio::task::yield_now().await; + } + }) + .await; + assert!( + stopped.is_ok(), + "authorization refresh task remained alive after close; strong_count={}", + Arc::strong_count(&authorizer) + ); + } + // 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. From 6a5ca29e6214e150746d907862307b81cdcd5e9b Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Mon, 24 Aug 2026 13:42:29 -0400 Subject: [PATCH 4/9] fix(eventhubs): stop authorization refresh on close --- .../azure_messaging_eventhubs/src/common/authorizer.rs | 7 +++++++ .../src/common/recoverable/connection.rs | 2 ++ 2 files changed, 9 insertions(+) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs index 73e3108827b..256dae05d9a 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs @@ -116,6 +116,13 @@ impl Authorizer { scopes.clear(); } + pub(crate) async fn stop_refresh_task(&self) { + if let Some(task) = self.authorization_refresher.get() { + task.abort(); + } + get_async_runtime().yield_now().await; + } + #[cfg(test)] pub(crate) fn disable_authorization(&self) -> Result<()> { use crate::EventHubsError; 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 f256a2abb81..70ed520f3b9 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs @@ -417,6 +417,8 @@ impl RecoverableConnection { "Closing recoverable connection." ); + self.authorizer.stop_refresh_task().await; + // Record the close before the teardown starts. A handle that outlives // the client, for example an `EventReceiver` that the caller still // holds, shares this object, and `ensure_connection` would otherwise From 55dc1cb89d038fef6d76dd075fabb2bffb59e0ee Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Mon, 24 Aug 2026 13:52:16 -0400 Subject: [PATCH 5/9] test: move connection test imports --- .../src/common/recoverable/connection.rs | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) 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 70ed520f3b9..e5202e83036 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs @@ -1500,9 +1500,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. // @@ -1544,13 +1552,6 @@ mod tests { #[tokio::test] async fn close_stops_owned_authorization_refresh_task() { - use azure_core::{ - credentials::{AccessToken, TokenCredential, TokenRequestOptions}, - time::OffsetDateTime, - }; - use std::sync::atomic::{AtomicUsize, Ordering}; - use tokio::sync::Notify; - #[derive(Debug)] struct GatedCredential { requests: AtomicUsize, From c0ae637574f47e2f9afdd2dbcf182a7f1ceca419 Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Thu, 27 Aug 2026 11:06:11 -0400 Subject: [PATCH 6/9] fix(eventhubs): stop authorization refresher on close Use a terminal refresher state to synchronize startup with shutdown. Await task cancellation before close returns so the authorizer ownership cycle is released. --- .../src/common/authorizer.rs | 44 ++++++--- .../src/common/recoverable/connection.rs | 93 ++++++++++++++++--- 2 files changed, 113 insertions(+), 24 deletions(-) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs index 256dae05d9a..423b06b8fcb 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs +++ b/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>, - authorization_refresher: OnceLock, + authorization_refresher: SyncMutex, /// Bias to apply to token refresh time. This determines how much time we will refresh the token before it expires. token_refresh_bias: SyncMutex, credential: Arc, @@ -86,9 +92,6 @@ pub(crate) struct Authorizer { disable_authorization: SyncMutex, } -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, @@ -117,10 +120,23 @@ impl Authorizer { } pub(crate) async fn stop_refresh_task(&self) { - if let Some(task) = self.authorization_refresher.get() { + 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; } - get_async_runtime().yield_now().await; } #[cfg(test)] @@ -222,12 +238,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); } 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 e5202e83036..ad60a19118e 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs @@ -417,8 +417,6 @@ impl RecoverableConnection { "Closing recoverable connection." ); - self.authorizer.stop_refresh_task().await; - // Record the close before the teardown starts. A handle that outlives // the client, for example an `EventReceiver` that the caller still // holds, shares this object, and `ensure_connection` would otherwise @@ -426,6 +424,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. @@ -1610,21 +1610,88 @@ mod tests { 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) + ); + } - let stopped = tokio::time::timeout(std::time::Duration::from_secs(1), async { - loop { - if Arc::strong_count(&authorizer) == 1 { - return; + #[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>, + ) -> azure_core::Result { + if self.requests.fetch_add(1, Ordering::SeqCst) == 0 { + self.entered.notify_one(); + self.release.notified().await; } - tokio::task::yield_now().await; + Ok(AccessToken::new( + azure_core::credentials::Secret::new("initial_token"), + OffsetDateTime::now_utc() + Duration::hours(1), + )) } - }) - .await; - assert!( - stopped.is_ok(), - "authorization refresh task remained alive after close; strong_count={}", - Arc::strong_count(&authorizer) + } + + 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, + 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" ); } From 0e4e7b4729aabae23d538776642d7024896432f0 Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Thu, 27 Aug 2026 11:09:53 -0400 Subject: [PATCH 7/9] test(eventhubs): propagate live test configuration errors --- .../tests/eventhubs_producer_memory.rs | 18 +++--------------- 1 file changed, 3 insertions(+), 15 deletions(-) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs index 2f978f6410c..a7190264c54 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs @@ -59,21 +59,9 @@ unsafe impl GlobalAlloc for LiveByteAllocator { #[recorded::test(live)] async fn producer_lifecycle_send_close_has_no_sustained_heap_trend() -> Result<(), Box> { - let host = match env::var("EVENTHUBS_HOST") { - Ok(host) if !host.is_empty() => host, - _ => return Ok(()), - }; - let eventhub = match env::var("EVENT_HUB_NAME") { - Ok(eventhub) if !eventhub.is_empty() => eventhub, - _ => match env::var("EVENTHUB_NAME") { - Ok(eventhub) if !eventhub.is_empty() => eventhub, - _ => return Ok(()), - }, - }; - let credential = match DeveloperToolsCredential::new(None) { - Ok(credential) => credential, - Err(_) => return Ok(()), - }; + let host = env::var("EVENTHUBS_HOST")?; + let eventhub = env::var("EVENT_HUB_NAME").or_else(|_| env::var("EVENTHUB_NAME"))?; + let credential = DeveloperToolsCredential::new(None)?; const WARMUP_CYCLES: usize = 5; const MEASURED_CYCLES: usize = 95; From 761e67b503987e497a1479b003bb45ffabe83645 Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Thu, 27 Aug 2026 12:48:56 -0400 Subject: [PATCH 8/9] test(eventhubs): explain producer memory regression --- .../tests/eventhubs_producer_memory.rs | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs index a7190264c54..1f25798299c 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs @@ -1,6 +1,14 @@ // Copyright (c) Microsoft Corporation. All Rights reserved // Licensed under the MIT license. +//! Live regression test for memory retained across repeated producer lifecycles. +//! +//! The custom global allocator tracks bytes that the Rust allocator currently +//! owns. Each sample is taken after a producer has opened, sent an event, and +//! closed. Other allocations in the test process add noise, so the test groups +//! samples into blocks, uses each block's median, and checks the trend across +//! those medians instead of comparing individual samples. + use azure_core_test::recorded; use azure_identity::DeveloperToolsCredential; use azure_messaging_eventhubs::ProducerClient; @@ -13,11 +21,16 @@ use std::{ struct LiveByteAllocator; +// This counter is process-wide because Rust permits only one global allocator. +// Relaxed ordering is sufficient because the test needs approximate snapshots, +// not synchronization with the threads that allocate or release memory. static LIVE_BYTES: AtomicUsize = AtomicUsize::new(0); #[global_allocator] static ALLOCATOR: LiveByteAllocator = LiveByteAllocator; +// Forward each operation to the system allocator and mirror successful size +// changes in LIVE_BYTES. A failed allocation leaves the count intact. unsafe impl GlobalAlloc for LiveByteAllocator { unsafe fn alloc(&self, layout: Layout) -> *mut u8 { let pointer = unsafe { System.alloc(layout) }; @@ -59,10 +72,14 @@ unsafe impl GlobalAlloc for LiveByteAllocator { #[recorded::test(live)] async fn producer_lifecycle_send_close_has_no_sustained_heap_trend() -> Result<(), Box> { + // A requested live run must fail when its configuration is incomplete. It + // must not report success without exercising a producer lifecycle. let host = env::var("EVENTHUBS_HOST")?; let eventhub = env::var("EVENT_HUB_NAME").or_else(|_| env::var("EVENTHUB_NAME"))?; let credential = DeveloperToolsCredential::new(None)?; + // Warmup absorbs one-time client and runtime initialization. The remaining + // 95 cycles form 19 blocks of 5 samples for the trend calculation. const WARMUP_CYCLES: usize = 5; const MEASURED_CYCLES: usize = 95; const BLOCK_SIZE: usize = 5; @@ -71,6 +88,8 @@ async fn producer_lifecycle_send_close_has_no_sustained_heap_trend() -> Result<( let mut samples = [0usize; MEASURED_CYCLES]; for cycle in 0..TOTAL_CYCLES { + // Recreate the full producer lifecycle on every cycle. Sampling after + // close exposes memory that the client failed to release during close. let client = match ProducerClient::builder() .with_application_id( "producer_lifecycle_send_close_has_no_sustained_heap_trend".to_string(), @@ -99,6 +118,8 @@ async fn producer_lifecycle_send_close_has_no_sustained_heap_trend() -> Result<( } } + // Use the median of each block to reduce the effect of temporary runtime, + // transport, and credential allocations on any one sample. let mut medians = [0.0; BLOCK_COUNT]; for (block_index, median) in medians.iter_mut().enumerate() { let start = block_index * BLOCK_SIZE; @@ -108,6 +129,8 @@ async fn producer_lifecycle_send_close_has_no_sustained_heap_trend() -> Result<( *median = block[BLOCK_SIZE / 2] as f64; } + // Fit an ordinary least-squares line to the block medians. Its slope is the + // estimated retained-byte change per block of producer lifecycles. let x_mean = (BLOCK_COUNT - 1) as f64 / 2.0; let y_mean = medians.iter().sum::() / BLOCK_COUNT as f64; let sum_squared_x_deviations = (0..BLOCK_COUNT) @@ -131,6 +154,10 @@ async fn producer_lifecycle_send_close_has_no_sustained_heap_trend() -> Result<( residual * residual }) .sum::(); + + // Estimate slope uncertainty from the regression residuals. The test fails + // only when the three-standard-error lower bound is positive. This identifies + // a sustained upward trend while tolerating normal allocator noise. let noise = (residual_sum_of_squares / (BLOCK_COUNT - 2) as f64).sqrt(); let slope_standard_error = noise / sum_squared_x_deviations.sqrt(); let lower_bound = slope - 3.0 * slope_standard_error; From 035b9b7466e299d4d32da3b25d1c235e61df38db Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Thu, 27 Aug 2026 13:21:59 -0400 Subject: [PATCH 9/9] fix(eventhubs): satisfy spell check Use the accepted two-word spelling in the producer memory regression test so the pull request spell-check job can complete. --- .../tests/eventhubs_producer_memory.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs index 1f25798299c..d99c40ccf08 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All Rights reserved // Licensed under the MIT license. -//! Live regression test for memory retained across repeated producer lifecycles. +//! Live regression test for memory retained across repeated producer life cycles. //! //! The custom global allocator tracks bytes that the Rust allocator currently //! owns. Each sample is taken after a producer has opened, sent an event, and @@ -130,7 +130,7 @@ async fn producer_lifecycle_send_close_has_no_sustained_heap_trend() -> Result<( } // Fit an ordinary least-squares line to the block medians. Its slope is the - // estimated retained-byte change per block of producer lifecycles. + // estimated retained-byte change per block of producer life cycles. let x_mean = (BLOCK_COUNT - 1) as f64 / 2.0; let y_mean = medians.iter().sum::() / BLOCK_COUNT as f64; let sum_squared_x_deviations = (0..BLOCK_COUNT)