From bc832bd02ed8e7257b43be6cd19be4e368a55b9e Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Mon, 22 Jun 2026 10:58:23 -0500 Subject: [PATCH 1/3] test(eventhubs): de-flake multiple_token_refresh The test read the mock credential's token-get count at fixed wall-clock instants and asserted on it. The final assertion fired only ~1.5-2.5s after the earliest moment the second token's refresh can complete, so under CPU contention the refresh task's wakeup slipped past the deadline, the test observed the stale pre-refresh count, and it panicked. Replace the fixed sleep-then-read assertions with bounded polling (wait_for_token_count): wait up to a generous timeout for the count to reach the expected value. This widens the acceptance window from ~2s to ~10s without changing happy-path runtime, and still fails fast if a refresh never happens. Also strengthen the first assertion from the trivially-true >= 2 to >= 3 to match its documented intent. --- .../src/common/authorizer.rs | 52 ++++++++++++++----- 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs index ed9e6ad31c..596e9687cd 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs @@ -721,6 +721,31 @@ mod tests { } } + // Poll `get_token_get_count` until it reaches `target` or `timeout` elapses, + // returning the last observed count either way. + // + // The token refresh task wakes on a real timer and only refreshes once the + // sleep overshoots the scheduled instant, so the exact moment a refresh lands + // drifts with scheduler latency. Reading the count at a single fixed instant + // makes the timing tests flaky under load (the read can race ahead of a + // refresh that is merely a little late). Waiting for the expected count keeps + // the assertions deterministic without widening the race window: a refresh + // that never happens still fails fast once the generous timeout expires. + async fn wait_for_token_count( + credential: &MockTokenCredential, + target: usize, + timeout: Duration, + ) -> usize { + let deadline = OffsetDateTime::now_utc() + timeout; + loop { + let count = credential.get_token_get_count(); + if count >= target || OffsetDateTime::now_utc() >= deadline { + return count; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + } + // When a token is created, it needs to have a proper expiration time. // This test verifies that the expiration time of tokens is set correctly when // authorizing a path. It also confirms that tokens are properly stored for reuse @@ -937,30 +962,29 @@ mod tests { // Token_refresh_1 will be refreshed between 4 and 6 seconds from now. // Token_refresh_2 will be refreshed between 14 and 16 from now. - trace!("Sleeping for 7 seconds to allow token_refresh_1 to expire and be refreshed. Current token count: {current_count}"); - tokio::time::sleep(std::time::Duration::from_secs(7)).await; - - // Verify that the token get count has increased, indicating a single refresh was attempted - we refreshed token_refresh_1 but not token_refresh_2. - let final_count = mock_credential.get_token_get_count(); - trace!("After sleeping the first time, token count: {final_count}"); + // Wait for token_refresh_1 to be refreshed (count goes 2 -> 3). The + // refresh lands ~5s from now; token_refresh_2 does not refresh until + // ~15s from now, so the count should reach exactly 3 within this window. + trace!("Waiting for token_refresh_1 to expire and be refreshed. Current token count: {current_count}"); + let final_count = wait_for_token_count(&mock_credential, 3, Duration::seconds(12)).await; + trace!("After waiting the first time, token count: {final_count}"); assert!( - final_count >= 2, - "Expected first get token count to be at least 2, but got {final_count}" + final_count >= 3, + "Expected token_refresh_1 to be refreshed (count >= 3), but got {final_count}" ); trace!("First token expiration get count: {}", final_count); // Token_refresh_1 will be refreshed between 13 and 15 seconds from now. // Token_refresh_2 will be refreshed between 7 and 9 seconds from now. - // Sleep for 10 seconds to allow the second token to expire and be refreshed. - tokio::time::sleep(std::time::Duration::from_secs(10)).await; - - // Verify that the token get count has increased, indicating a single refresh was attempted - we refreshed token_refresh_2. - let final_count = mock_credential.get_token_get_count(); + // Wait for token_refresh_2 to be refreshed (count goes 3 -> 4). It + // refreshes ~7-9s from now; the timeout is generous so that scheduler + // latency under load delays the test rather than failing it. + let final_count = wait_for_token_count(&mock_credential, 4, Duration::seconds(20)).await; trace!("Getting second token count: {final_count}"); assert!( final_count >= 4, - "Expected second get token count to be 4, but got {final_count}" + "Expected token_refresh_2 to be refreshed (count >= 4), but got {final_count}" ); trace!("Second token expiration get count: {}", final_count); From d0c82238a13c8ecb592ae2e0d50a54059cad0f42 Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Mon, 3 Aug 2026 15:32:04 -0400 Subject: [PATCH 2/3] test(eventhubs): un-ignore token_refresh The sibling test `token_refresh` was ignored in #3691 for "frequent off-by-one issues in dev loop". That is the same flakiness class this branch fixes: the test read the mock credential's token-get count at a fixed instant after a 13-second sleep, and the refresh it waits for lands 8 to 12 seconds in. The margin was about one second, so scheduler latency pushed the refresh past the read. Replace the fixed sleep and read with `wait_for_token_count`, then remove the `#[ignore]`. --- .../src/common/authorizer.rs | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs index 596e9687cd..ed877dc0fe 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs @@ -820,7 +820,6 @@ mod tests { // If this feature fails in production, clients would disconnect when their tokens expire, // which could lead to data loss, application failures, or service degradation. #[recorded::test] - #[ignore = "frequent off-by-one issues in dev loop"] async fn token_refresh(_ctx: TestContext) -> Result<()> { let url = Url::parse("amqps://example.com").unwrap(); let path = Url::parse("amqps://example.com/test_token_refresh").unwrap(); @@ -873,16 +872,14 @@ mod tests { let current_count = mock_credential.get_token_get_count(); assert_eq!(current_count, 1); - trace!("Sleeping for 15 seconds to allow token to expire and be refreshed. Current token count: {current_count}"); + trace!("Waiting for the token to expire and be refreshed. Current token count: {current_count}"); - // Sleep a bit to ensure we will have refreshed the token - since the token expires in 20 seconds, - // we will refresh it between 8 and 12 seconds before the expiration time. If we wait for 13 seconds, - // we should have refreshed the token. - tokio::time::sleep(std::time::Duration::from_secs(13)).await; - - // Verify that the token get count has increased, indicating a refresh was attempted - let final_count = mock_credential.get_token_get_count(); - trace!("After sleeping, token count: {final_count}"); + // The token expires in 20 seconds and refreshes 8 to 12 seconds before + // that, so the refresh lands 8 to 12 seconds from now. Wait for the count + // to reach 2 instead of reading it at a fixed instant, because the exact + // moment the refresh lands drifts with scheduler latency. + let final_count = wait_for_token_count(&mock_credential, 2, Duration::seconds(25)).await; + trace!("After waiting, token count: {final_count}"); assert!( final_count >= 2, From 3576d7c4964d311f10076b3142c6ded8b719027b Mon Sep 17 00:00:00 2001 From: Johnathan W Date: Mon, 3 Aug 2026 16:01:45 -0400 Subject: [PATCH 3/3] test(eventhubs): wait per path and use a monotonic deadline Two points from the Copilot review. The second wait in multiple_token_refresh spans about 20 seconds, which is wide enough to cover path1's next refresh. The get_token count is shared by both paths, so path1 could raise the count to 4 on its own and satisfy the assertion even if path2 never refreshed. A probe of the real schedule measured path2 refreshing 15.1s in and path1 refreshing again 20.7s in, both inside that window. Add a test-only peek_token hook and a wait_for_token_refresh helper that watches one path's cached expiry. A refresh replaces the entry in place and only ever advances expires_on, so each assertion now names the path it tests. Also measure both timeouts with tokio::time::Instant instead of the wall clock, so a clock correction cannot end a wait early or stretch it past the stated bound. --- .../src/common/authorizer.rs | 106 +++++++++++++++--- 1 file changed, 90 insertions(+), 16 deletions(-) diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs index ed877dc0fe..06b1144dae 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs @@ -128,6 +128,17 @@ impl Authorizer { Ok(()) } + /// Test hook: read the cached token for `path` without authorizing. + /// + /// A refresh replaces the map entry in place and only ever advances + /// `expires_on`, so a test can watch one path's expiry to tell whether that + /// specific path was refreshed. The shared `get_token` call count cannot make + /// that distinction. + #[cfg(test)] + async fn peek_token(&self, path: &Url) -> Option { + self.authorization_scopes.read().await.get(path).cloned() + } + #[tracing::instrument( level = "debug", skip_all, @@ -731,21 +742,52 @@ mod tests { // refresh that is merely a little late). Waiting for the expected count keeps // the assertions deterministic without widening the race window: a refresh // that never happens still fails fast once the generous timeout expires. + // + // The deadline uses `tokio::time::Instant` rather than the wall clock, so a + // clock correction cannot cut the wait short or stretch it past the bound. async fn wait_for_token_count( credential: &MockTokenCredential, target: usize, - timeout: Duration, + timeout: std::time::Duration, ) -> usize { - let deadline = OffsetDateTime::now_utc() + timeout; + let deadline = tokio::time::Instant::now() + timeout; loop { let count = credential.get_token_get_count(); - if count >= target || OffsetDateTime::now_utc() >= deadline { + if count >= target || tokio::time::Instant::now() >= deadline { return count; } tokio::time::sleep(std::time::Duration::from_millis(100)).await; } } + // Poll the cached token for `path` until its expiry advances past + // `previous_expiry` or `timeout` elapses, returning the last observed expiry. + // + // `wait_for_token_count` watches a counter shared by every path, so it cannot + // tell which path a refresh belonged to. Once a wait window is wide enough to + // overlap another path's refresh, a count target can be reached by the wrong + // path and the assertion passes for the wrong reason. A refresh replaces the + // cache entry in place and only ever advances `expires_on` (a credential that + // returns the same expiry is marked non-refreshable instead), so watching one + // path's expiry ties each assertion to the path it names. + async fn wait_for_token_refresh( + authorizer: &Arc, + path: &Url, + previous_expiry: OffsetDateTime, + timeout: std::time::Duration, + ) -> Option { + let deadline = tokio::time::Instant::now() + timeout; + loop { + let expiry = authorizer.peek_token(path).await.map(|t| t.expires_on); + if expiry.is_some_and(|e| e > previous_expiry) + || tokio::time::Instant::now() >= deadline + { + return expiry; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + } + // When a token is created, it needs to have a proper expiration time. // This test verifies that the expiration time of tokens is set correctly when // authorizing a path. It also confirms that tokens are properly stored for reuse @@ -878,7 +920,8 @@ mod tests { // that, so the refresh lands 8 to 12 seconds from now. Wait for the count // to reach 2 instead of reading it at a fixed instant, because the exact // moment the refresh lands drifts with scheduler latency. - let final_count = wait_for_token_count(&mock_credential, 2, Duration::seconds(25)).await; + let final_count = + wait_for_token_count(&mock_credential, 2, std::time::Duration::from_secs(25)).await; trace!("After waiting, token count: {final_count}"); assert!( @@ -933,10 +976,11 @@ mod tests { let path1 = Url::parse("amqps://example.com/test_token_refresh_1").unwrap(); // Get access to the connection //let connection = connection_manager.ensure_connection().await.unwrap(); - authorizer + let path1_expiry = authorizer .authorize_path(&recoverable_connection, &path1) .await - .unwrap(); + .unwrap() + .expires_on; // Because the token expires in 20 seconds, token_refresh_1 will be refreshed // between 14 and 16 seconds from now. @@ -947,10 +991,11 @@ mod tests { // Authorize the second path, which will store the token let path2 = Url::parse("amqps://example.com/test_token_refresh_2").unwrap(); - authorizer + let path2_expiry = authorizer .authorize_path(&recoverable_connection, &path2) .await - .unwrap(); + .unwrap() + .expires_on; // Verify initial token retrieval count - it should have been refreshed three times - let current_count = mock_credential.get_token_get_count(); @@ -959,11 +1004,26 @@ mod tests { // Token_refresh_1 will be refreshed between 4 and 6 seconds from now. // Token_refresh_2 will be refreshed between 14 and 16 from now. - // Wait for token_refresh_1 to be refreshed (count goes 2 -> 3). The - // refresh lands ~5s from now; token_refresh_2 does not refresh until - // ~15s from now, so the count should reach exactly 3 within this window. + // + // Wait on path1's own expiry rather than the shared call count. The count + // cannot say which path was refreshed, and the second wait below is wide + // enough to overlap path1's next refresh, so a count target there could be + // reached by path1 even if path2 never refreshed. trace!("Waiting for token_refresh_1 to expire and be refreshed. Current token count: {current_count}"); - let final_count = wait_for_token_count(&mock_credential, 3, Duration::seconds(12)).await; + let refreshed1 = wait_for_token_refresh( + &authorizer, + &path1, + path1_expiry, + std::time::Duration::from_secs(12), + ) + .await; + assert!( + refreshed1.is_some_and(|e| e > path1_expiry), + "Expected token_refresh_1 to be refreshed (expiry to advance past {path1_expiry}), but got {refreshed1:?}" + ); + + // Only path1 is due in this window, so exactly one refresh has landed. + let final_count = mock_credential.get_token_get_count(); trace!("After waiting the first time, token count: {final_count}"); assert!( final_count >= 3, @@ -974,10 +1034,24 @@ mod tests { // Token_refresh_1 will be refreshed between 13 and 15 seconds from now. // Token_refresh_2 will be refreshed between 7 and 9 seconds from now. - // Wait for token_refresh_2 to be refreshed (count goes 3 -> 4). It - // refreshes ~7-9s from now; the timeout is generous so that scheduler - // latency under load delays the test rather than failing it. - let final_count = wait_for_token_count(&mock_credential, 4, Duration::seconds(20)).await; + // Wait for token_refresh_2 to be refreshed. It refreshes ~7-9s from now; + // the timeout is generous so that scheduler latency under load delays the + // test rather than failing it. Waiting on path2's own expiry keeps the + // assertion honest even though this window also covers path1's next + // refresh. + let refreshed2 = wait_for_token_refresh( + &authorizer, + &path2, + path2_expiry, + std::time::Duration::from_secs(20), + ) + .await; + assert!( + refreshed2.is_some_and(|e| e > path2_expiry), + "Expected token_refresh_2 to be refreshed (expiry to advance past {path2_expiry}), but got {refreshed2:?}" + ); + + let final_count = mock_credential.get_token_get_count(); trace!("Getting second token count: {final_count}"); assert!( final_count >= 4,