CrossRegionHedging: Fixes NullReference Bug

Azure · Nov 5, 2024 · f5dc6c0 · f5dc6c0
1 parent 408ee12
commit f5dc6c0
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 5 deletions.
diff --git a/...t.Azure.Cosmos/src/Routing/AvailabilityStrategy/CrossRegionHedgingAvailabilityStrategy.cs b/...t.Azure.Cosmos/src/Routing/AvailabilityStrategy/CrossRegionHedgingAvailabilityStrategy.cs
@@ -143,7 +143,7 @@ internal override async Task<ResponseMessage> ExecuteAvailabilityStrategyAsync(
                                         request,
                                         hedgeRegions.ElementAt(requestNumber),
                                         cancellationToken,
-                                        cancellationTokenSource);
+                                        cancellationTokenSource, trace);
 
                                     requestTasks.Add(primaryRequest);
                                 }
@@ -262,7 +262,7 @@ private async Task<HedgingResponse> CloneAndSendAsync(
                     clonedRequest,
                     region,
                     cancellationToken,
-                    cancellationTokenSource);
+                    cancellationTokenSource, trace);
             }
         }
 
@@ -271,7 +271,7 @@ private async Task<HedgingResponse> RequestSenderAndResultCheckAsync(
             RequestMessage request,
             string hedgedRegion,
             CancellationToken cancellationToken,
-            CancellationTokenSource cancellationTokenSource)
+            CancellationTokenSource cancellationTokenSource, ITrace trace)
         {
             try
             {
@@ -288,9 +288,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)
             {

diff --git a/...zure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosAvailabilityStrategyTests.cs b/...zure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosAvailabilityStrategyTests.cs
@@ -840,6 +840,60 @@ public async Task AvailabilityStrategyStepTests(string operation, string condito
             }
         }
 
+        [TestMethod]
+        [TestCategory("MultiRegion")]
+        [ExpectedException(typeof(CosmosOperationCanceledException))]
+        public async Task AvailabilityStrategyWithCancellationTokenThrowsExceptionTest()
+        {
+            FaultInjectionRule responseDelay = new FaultInjectionRuleBuilder(
+                id: "responseDely",
+                condition:
+                    new FaultInjectionConditionBuilder()
+                        .WithRegion("Central US")
+                        .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>() { "Central US", "North Central US" },
+                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);
+
+                ItemResponse<AvailabilityStrategyTestObject> ir = await container.ReadItemAsync<AvailabilityStrategyTestObject>(
+                "testId",
+                    new PartitionKey("pk"), cancellationToken: cts.Token
+                    );
+
+            }
+
+        }
+
         private static async Task HandleChangesAsync(
             ChangeFeedProcessorContext context,
             IReadOnlyCollection<AvailabilityStrategyTestObject> changes,