diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs index ed9e6ad31c4..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, @@ -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 @@ -215,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); } @@ -639,6 +668,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)] 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..ad60a19118e 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/recoverable/connection.rs @@ -424,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. @@ -1498,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. // @@ -1540,6 +1550,151 @@ 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>, + ) -> 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(); + 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>, + ) -> azure_core::Result { + 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, + 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. 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..d99c40ccf08 --- /dev/null +++ b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_producer_memory.rs @@ -0,0 +1,171 @@ +// Copyright (c) Microsoft Corporation. All Rights reserved +// Licensed under the MIT license. + +//! 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 +//! 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; +use std::{ + alloc::{GlobalAlloc, Layout, System}, + env, + error::Error, + sync::atomic::{AtomicUsize, Ordering}, +}; + +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) }; + 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> { + // 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; + 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 { + // 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(), + ) + .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); + } + } + + // 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; + 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; + } + + // Fit an ordinary least-squares line to the block medians. Its slope is the + // 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) + .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::(); + + // 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; + + 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(()) +}