fix(eventhubs): stop authorization refresh on close

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

@@ -116,8 +116,15 @@ 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)]
-    fn disable_authorization(&self) -> Result<()> {
+    pub(crate) fn disable_authorization(&self) -> Result<()> {
         use crate::EventHubsError;
 
         let mut disable_authorization = self
@@ -639,6 +646,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

@@ -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
@@ -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,84 @@ 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,
+            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.

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<dyn Error>> {
+    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::<f64>() / BLOCK_COUNT as f64;
+    let sum_squared_x_deviations = (0..BLOCK_COUNT)
+        .map(|index| {
+            let deviation = index as f64 - x_mean;
+            deviation * deviation
+        })
+        .sum::<f64>();
+    let slope = medians
+        .iter()
+        .enumerate()
+        .map(|(index, median)| (index as f64 - x_mean) * (median - y_mean))
+        .sum::<f64>()
+        / 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::<f64>();
+    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(())
+}