CrossRegionHedging: Fixes NullReference Exception Bug

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

@@ -160,7 +160,8 @@ internal override async Task<ResponseMessage> ExecuteAvailabilityStrategyAsync(
                                         request,
                                         hedgeRegions.ElementAt(requestNumber),
                                         cancellationToken,
-                                        cancellationTokenSource);
+                                        cancellationTokenSource, 
+                                        trace);
 
                                     requestTasks.Add(primaryRequest);
                                 }
@@ -279,7 +280,8 @@ private async Task<HedgingResponse> CloneAndSendAsync(
                     clonedRequest,
                     region,
                     cancellationToken,
-                    cancellationTokenSource);
+                    cancellationTokenSource, 
+                    trace);
             }
         }
 
@@ -288,7 +290,8 @@ private async Task<HedgingResponse> RequestSenderAndResultCheckAsync(
             RequestMessage request,
             string hedgedRegion,
             CancellationToken cancellationToken,
-            CancellationTokenSource cancellationTokenSource)
+            CancellationTokenSource cancellationTokenSource,
+            ITrace trace)
         {
             try
             {
@@ -305,9 +308,9 @@ private async Task<HedgingResponse> RequestSenderAndResultCheckAsync(
 
                 return new HedgingResponse(false, response, hedgedRegion);
             }
-            catch (OperationCanceledException) when (cancellationTokenSource.IsCancellationRequested)
+            catch (OperationCanceledException oce ) when (cancellationTokenSource.IsCancellationRequested)
             {
-                return new HedgingResponse(false, null, hedgedRegion);
+                throw new CosmosOperationCanceledException(oce, trace);
             }
             catch (Exception ex)
             {

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosAvailabilityStrategyTests.cs

@@ -1219,6 +1219,60 @@ await this.container.DeleteItemAsync<AvailabilityStrategyTestObject>(
             }
         }
 
+        [TestMethod]
+        [TestCategory("MultiRegion")]
+        public async Task AvailabilityStrategyWithCancellationTokenThrowsExceptionTest()
+        {
+            FaultInjectionRule responseDelay = new FaultInjectionRuleBuilder(
+                id: "responseDely",
+                condition:
+                    new FaultInjectionConditionBuilder()
+                        .WithRegion(region1)
+                        .WithOperationType(FaultInjectionOperationType.ReadItem)
+                        .Build(),
+                result:
+                    FaultInjectionResultBuilder.GetResultBuilder(FaultInjectionServerErrorType.ResponseDelay)
+                        .WithDelay(TimeSpan.FromMilliseconds(6000))
+                        .Build())
+                .WithDuration(TimeSpan.FromMinutes(90))
+                .WithHitLimit(2)
+                .Build();
+
+            List<FaultInjectionRule> rules = new List<FaultInjectionRule>() { responseDelay };
+            FaultInjector faultInjector = new FaultInjector(rules);
+
+            responseDelay.Disable();
+
+            CosmosClientOptions clientOptions = new CosmosClientOptions()
+            {
+                ConnectionMode = ConnectionMode.Direct,
+                ApplicationPreferredRegions = new List<string>() { region1, region2 },
+                AvailabilityStrategy = AvailabilityStrategy.CrossRegionHedgingStrategy(
+                        threshold: TimeSpan.FromMilliseconds(300),
+                        thresholdStep: null),
+                Serializer = this.cosmosSystemTextJsonSerializer
+            };
+
+            using (CosmosClient faultInjectionClient = new CosmosClient(
+                connectionString: this.connectionString,
+                clientOptions: faultInjector.GetFaultInjectionClientOptions(clientOptions)))
+            {
+                CancellationTokenSource cts = new CancellationTokenSource();
+                cts.Cancel();
+
+                Database database = faultInjectionClient.GetDatabase(CosmosAvailabilityStrategyTests.dbName);
+                Container container = database.GetContainer(CosmosAvailabilityStrategyTests.containerName);
+
+                CosmosOperationCanceledException cancelledException = await Assert.ThrowsExceptionAsync<CosmosOperationCanceledException>(() =>
+                        container.ReadItemAsync<AvailabilityStrategyTestObject>(
+                            "testId",
+                            new PartitionKey("pk"), cancellationToken: cts.Token
+                    ));
+
+            }
+
+        }
+
         private static async Task HandleChangesAsync(
             ChangeFeedProcessorContext context,
             IReadOnlyCollection<AvailabilityStrategyTestObject> changes,

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/AvailabilityStrategyUnitTests.cs

@@ -1,10 +1,11 @@
 namespace Microsoft.Azure.Cosmos.Tests
 {
     using System;
-    using System.Collections.Concurrent;
     using System.Collections.Generic;
+    using System.Collections.ObjectModel;
     using System.IO;
     using System.Net.Http;
+    using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos;
     using Microsoft.Azure.Documents;
@@ -58,5 +59,42 @@ public async Task RequestMessageCloneTests()
                 Assert.AreEqual(httpRequest.DatabaseId, clone.DatabaseId);
             }
         }
+
+        [TestMethod]
+        public async Task CancellationTokenThrowsExceptionTest()
+        {
+            //Arrange
+            CrossRegionHedgingAvailabilityStrategy availabilityStrategy = new CrossRegionHedgingAvailabilityStrategy(
+                                                                   threshold: TimeSpan.FromMilliseconds(100),
+                                                                   thresholdStep: TimeSpan.FromMilliseconds(50));
+
+            RequestMessage request = new RequestMessage
+            {
+                ResourceType = ResourceType.Document,
+                OperationType = OperationType.Read
+            };
+
+            CancellationTokenSource cts = new CancellationTokenSource();
+            cts.Cancel();
+
+            AccountProperties databaseAccount = new AccountProperties()
+            {
+                ReadLocationsInternal = new Collection<AccountRegion>()
+                {
+                    { new AccountRegion() { Name = "US East", Endpoint = new Uri("https://location1.documents.azure.com").ToString() } },
+                    { new AccountRegion() { Name = "US West", Endpoint = new Uri("https://location2.documents.azure.com").ToString() } },
+
+                }
+            };
+            using CosmosClient mockCosmosClient = MockCosmosUtil.CreateMockCosmosClient();
+            mockCosmosClient.DocumentClient.GlobalEndpointManager.InitializeAccountPropertiesAndStartBackgroundRefresh(databaseAccount);
+
+            Func<RequestMessage, CancellationToken, Task<ResponseMessage>> sender = (request, token) => throw new OperationCanceledException("operation cancellation requested");
+
+            CosmosOperationCanceledException cancelledException = await Assert.ThrowsExceptionAsync<CosmosOperationCanceledException>(() =>
+                       availabilityStrategy.ExecuteAvailabilityStrategyAsync(sender, mockCosmosClient, request, cts.Token));
+        }
+
+
     }
 }