CrossRegionHedgingAvailabilityStrategy: Fixes ArgumentNullException race condition in hedging cancellation (#5613)

## CrossRegionHedgingAvailabilityStrategy: Fixes `ArgumentNullException`
race condition in hedging cancellation

### Bug

Multiple production customers reported unobserved
`ArgumentNullException: Value cannot be null. (Parameter 'request')`
crashes originating from
`CrossRegionHedgingAvailabilityStrategy.RequestSenderAndResultCheckAsync`.
The exception was surfaced as a `TaskScheduler_UnobservedTaskException`,
crashing the process. The affected code paths were:

- `ContainerCore.ReadItemAsync` → `ReadItemStreamAsync` →
`ProcessItemStreamAsync` → `RequestInvokerHandler.SendAsync` →
`CrossRegionHedgingAvailabilityStrategy.ExecuteAvailabilityStrategyAsync`

### Root Cause

A **race condition** caused by passing the wrong `CancellationToken` to
the sender delegate:

1. `ExecuteAvailabilityStrategyAsync` creates a
`hedgeRequestsCancellationTokenSource` (linked to the app-provided CT)
to coordinate hedge request lifecycle.
2. `CloneAndSendAsync` clones the request inside a `using` block and
calls `RequestSenderAndResultCheckAsync`.
3. **Bug:** `RequestSenderAndResultCheckAsync` called
`sender.Invoke(request, cancellationToken)` with the
**application-provided `CancellationToken`** — not
`hedgeRequestsCancellationTokenSource.Token`.
4. When hedge Region B returned a final result (e.g., 200 OK),
`hedgeRequestsCancellationTokenSource.Cancel()` was called.
5. **But the in-flight sender for Region A still held the app CT**
(e.g., `CancellationToken.None`), which was **never cancelled**.
6. The `CloneAndSendAsync` `using` block exited, **disposing the cloned
request**.
7. The Region A sender continued executing with a reference to the
now-disposed request → **`ArgumentNullException: Value cannot be null.
(Parameter 'request')`**.

A secondary issue: when the application CT was cancelled (e2e timeout),
the hedge timer (linked to app CT) would fire, and the old code would
blindly continue the loop attempting to clone and send new requests on a
cancelled path.

### Fix

Two changes in `CrossRegionHedgingAvailabilityStrategy.cs`:

**1. Pass `hedgeRequestsCancellationTokenSource.Token` to
`sender.Invoke()` instead of the app CT**

This ensures that when **any** hedge gets a final result and calls
`hedgeRequestsCancellationTokenSource.Cancel()`, **all** in-flight
senders immediately see their CT cancelled and stop before the cloned
request is disposed. The `CancellationTokenSource` and
`CancellationToken` parameters were also consolidated into a single
`hedgeRequestsCancellationTokenSource` parameter passed through
`CloneAndSendAsync` → `RequestSenderAndResultCheckAsync`.

**2. Add `do/while` loop to handle spurious timer completions on app CT
cancellation**

When the app CT is cancelled (e2e timeout), the hedge timer fires via
the linked CTS. The old code would `continue` the loop and try to clone
a new request. The `do/while` loop now detects
`applicationProvidedCancellationToken.IsCancellationRequested` and falls
through to consolidate existing request outcomes instead of spawning new
hedges.

### Tests Added (8 new unit tests)

| Test | Validates |
|---|---|
| `HedgeCancellationCancelsInFlightRequests_NoNullRef` | Slow primary
request's CT is cancelled when a hedge returns a final result — core
regression test |
| `SenderReceivesHedgeCancellationToken_NotAppToken` | Captures the
actual CT passed to each sender and asserts all are from the hedge CTS,
not the app CT |
| `AppCancellationDuringHedging_DoesNotSpawnNewHedgeRequests` | E2e
timeout (app CT cancelled) does not spawn new hedge requests — validates
the do/while loop fix |
| `MultiRegionHedging_RequestNotAccessedAfterDisposal` | Verifies the
cloned request is still accessible when cancellation fires — exact
scenario from the crash reports |
| `HedgeCancellation_StreamRequest_NoNullRef` | Tests the stream-based
code path (ReadItemStreamAsync) from the NullRef2/NullRef3 stack traces
|
| `PrimaryRequestFinalResult_NoAdditionalHedgesSent` | Fast primary
response skips hedging entirely |
| `AllHedgesTransientError_ReturnsLastResponse` | All regions return
transient errors — strategy returns last response without NullRef |
| `ConcurrentHedgingRequests_NoNullRef` | Stress test: 50 concurrent
hedging requests with random delays — no NullRef under concurrency |

### Type of change

- [x] Bug fix (non-breaking change which fixes an issue)

---------

Co-authored-by: Nalu Tripician <27316859+NaluTripician@users.noreply.github.com>

Author: FabianMeiswinkel

Microsoft.Azure.Cosmos/src/Routing/AvailabilityStrategy/CrossRegionHedgingAvailabilityStrategy.cs

@@ -125,23 +125,24 @@ internal bool ShouldHedge(RequestMessage request, CosmosClient client)
         /// <param name="sender"></param>
         /// <param name="client"></param>
         /// <param name="request"></param>
-        /// <param name="cancellationToken"></param>
+        /// <param name="applicationProvidedCancellationToken"></param>
         /// <returns>The response after executing cross region hedging</returns>
         internal override async Task<ResponseMessage> ExecuteAvailabilityStrategyAsync(
             Func<RequestMessage, CancellationToken, Task<ResponseMessage>> sender,
             CosmosClient client,
             RequestMessage request,
-            CancellationToken cancellationToken)
+            CancellationToken applicationProvidedCancellationToken)
         {
             if (!this.ShouldHedge(request, client)
                 || client.DocumentClient.GlobalEndpointManager.ReadEndpoints.Count == 1)
             {
-                return await sender(request, cancellationToken);
+                return await sender(request, applicationProvidedCancellationToken);
             }
 
             ITrace trace = request.Trace;
 
-            using (CancellationTokenSource cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken))
+            using (CancellationTokenSource hedgeRequestsCancellationTokenSource = 
+                CancellationTokenSource.CreateLinkedTokenSource(applicationProvidedCancellationToken))
             {
                 using (CloneableStream clonedBody = (CloneableStream)(request.Content == null
                     ? null
@@ -161,7 +162,7 @@ internal override async Task<ResponseMessage> ExecuteAvailabilityStrategyAsync(
                     {
                         TimeSpan awaitTime = requestNumber == 0 ? this.Threshold : this.ThresholdStep;
 
-                        using (CancellationTokenSource timerTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken))
+                        using (CancellationTokenSource timerTokenSource = CancellationTokenSource.CreateLinkedTokenSource(applicationProvidedCancellationToken))
                         {
                             CancellationToken timerToken = timerTokenSource.Token;
                             using (Task hedgeTimer = Task.Delay(awaitTime, timerToken))
@@ -173,32 +174,50 @@ internal override async Task<ResponseMessage> ExecuteAvailabilityStrategyAsync(
                                         hedgeRegions: hedgeRegions,
                                         requestNumber: requestNumber,
                                         trace: trace,
-                                        cancellationToken: cancellationToken,
-                                        cancellationTokenSource: cancellationTokenSource);
+                                        hedgeRequestsCancellationTokenSource: hedgeRequestsCancellationTokenSource);
 
                                 requestTasks.Add(requestTask);
                                 requestTasks.Add(hedgeTimer);
 
-                                Task completedTask = await Task.WhenAny(requestTasks);
-                                requestTasks.Remove(completedTask);
+                                Task completedTask;
+                                do
+                                {
+                                    completedTask = await Task.WhenAny(requestTasks);
+                                    requestTasks.Remove(completedTask);
+                                }
+                                while (
+                                    completedTask == hedgeTimer &&
+                                    // Ignore hedge timer signals if either the e2e timeout is hit 
+                                    // or the hedgeTimer task failed (or more commonly since this is a linked CTS was cancelled)
+                                    // in both of these cases we do not want to spawn new hedge requests
+                                    // but just consolidate the outcome of previous requests
+                                    (!completedTask.IsCompleted || applicationProvidedCancellationToken.IsCancellationRequested));
 
                                 if (completedTask == hedgeTimer)
                                 {
                                     continue;
                                 }
 
-                                timerTokenSource.Cancel();
                                 requestTasks.Remove(hedgeTimer);
+                                timerTokenSource.Cancel();
 
-                                if (completedTask.IsFaulted)
+                                if (completedTask.IsFaulted || completedTask.IsCanceled)
                                 {
-                                    AggregateException innerExceptions = completedTask.Exception.Flatten();
+                                    requestTasks.Remove(hedgeTimer);
+                                    timerTokenSource.Cancel();
+
+                                    if (applicationProvidedCancellationToken.IsCancellationRequested)
+                                    {
+                                        await (Task<HedgingResponse>)completedTask;
+                                    }
+
+                                    continue;
                                 }
 
                                 hedgeResponse = await (Task<HedgingResponse>)completedTask;
                                 if (hedgeResponse.IsNonTransient)
                                 {
-                                    cancellationTokenSource.Cancel();
+                                    hedgeRequestsCancellationTokenSource.Cancel();
 
                                     ((CosmosTraceDiagnostics)hedgeResponse.ResponseMessage.Diagnostics).Value.AddOrUpdateDatum(
                                         HedgeConfig,
@@ -227,12 +246,19 @@ internal override async Task<ResponseMessage> ExecuteAvailabilityStrategyAsync(
                         {
                             AggregateException innerExceptions = completedTask.Exception.Flatten();
                             lastException = innerExceptions.InnerExceptions.FirstOrDefault();
+                            continue;
+                        }
+
+                        if (completedTask.IsCanceled)
+                        {
+                            lastException = new OperationCanceledException();
+                            continue;
                         }
 
                         hedgeResponse = await (Task<HedgingResponse>)completedTask;
                         if (hedgeResponse.IsNonTransient || requestTasks.Count == 0)
                         {
-                            cancellationTokenSource.Cancel();
+                            hedgeRequestsCancellationTokenSource.Cancel();
                             ((CosmosTraceDiagnostics)hedgeResponse.ResponseMessage.Diagnostics).Value.AddOrUpdateDatum(
                                         HedgeConfig,
                                         this.HedgeConfigText);
@@ -251,7 +277,16 @@ internal override async Task<ResponseMessage> ExecuteAvailabilityStrategyAsync(
                         throw lastException;
                     }
 
-                    Debug.Assert(hedgeResponse != null);
+                    if (hedgeResponse == null)
+                    {
+                        if (applicationProvidedCancellationToken.IsCancellationRequested)
+                        {
+                            throw new CosmosOperationCanceledException(new OperationCanceledException(), trace);
+                        }
+
+                        throw new InvalidOperationException("Cross-region hedging completed without producing a response.");
+                    }
+
                     return hedgeResponse.ResponseMessage;
                 }
             }
@@ -264,8 +299,7 @@ private async Task<HedgingResponse> CloneAndSendAsync(
             IReadOnlyCollection<string> hedgeRegions,
             int requestNumber,
             ITrace trace,
-            CancellationToken cancellationToken,
-            CancellationTokenSource cancellationTokenSource)
+            CancellationTokenSource hedgeRequestsCancellationTokenSource)
         {
             RequestMessage clonedRequest;
 
@@ -287,8 +321,7 @@ private async Task<HedgingResponse> CloneAndSendAsync(
                     sender,
                     clonedRequest,
                     hedgeRegions.ElementAt(requestNumber),
-                    cancellationToken,
-                    cancellationTokenSource, 
+                    hedgeRequestsCancellationTokenSource, 
                     trace);
             }
         }
@@ -297,27 +330,30 @@ private async Task<HedgingResponse> RequestSenderAndResultCheckAsync(
             Func<RequestMessage, CancellationToken, Task<ResponseMessage>> sender,
             RequestMessage request,
             string targetRegionName,
-            CancellationToken cancellationToken,
-            CancellationTokenSource cancellationTokenSource,
+            CancellationTokenSource hedgeRequestsCancellationTokenSource,
             ITrace trace)
         {
             try
             {
-                ResponseMessage response = await sender.Invoke(request, cancellationToken);
+                ResponseMessage response = await sender.Invoke(request, hedgeRequestsCancellationTokenSource.Token);
                 if (IsFinalResult((int)response.StatusCode, (int)response.Headers.SubStatusCode))
                 {
-                    if (!cancellationToken.IsCancellationRequested)
+                    if (!hedgeRequestsCancellationTokenSource.IsCancellationRequested)
                     {
-                        cancellationTokenSource.Cancel();
+                        // App has not reached e2e timeout - we can cancel any still remaining
+                        // hedge requests since we have a final response now
+                        hedgeRequestsCancellationTokenSource.Cancel();
                     }
 
                     return new HedgingResponse(true, response, targetRegionName);
                 }
 
                 return new HedgingResponse(false, response, targetRegionName);
             }
-            catch (OperationCanceledException oce) when (cancellationTokenSource.IsCancellationRequested)
+            catch (OperationCanceledException oce) when (hedgeRequestsCancellationTokenSource.IsCancellationRequested)
             {
+                // hedgeRequestsCancellationTokenSource is a linked cancellation token source - so, would also signal
+                // cancellation on e2e timeout via app provided CT
                 throw new CosmosOperationCanceledException(oce, trace);
             }
             catch (Exception ex)