Fixing a missing retry in CosmosFhirDataStore. (#4831)

microsoft · Feb 20, 2025 · ceba9ce · ceba9ce
1 parent 718597c
commit ceba9ce
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 1 deletion.
diff --git a/src/Microsoft.Health.Fhir.Core.UnitTests/Features/Operations/Export/ExportJobTaskTests.cs b/src/Microsoft.Health.Fhir.Core.UnitTests/Features/Operations/Export/ExportJobTaskTests.cs
@@ -16,6 +16,7 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using Microsoft.Extensions.Options;
+using Microsoft.Health.Abstractions.Exceptions;
 using Microsoft.Health.Core.Features.Context;
 using Microsoft.Health.Extensions.DependencyInjection;
 using Microsoft.Health.Fhir.Core.Configs;
@@ -2086,6 +2087,30 @@ public async Task GivenAnExportJobWithInvalidStorageAccount_WhenExecuted_ThenAnE
             Assert.Equal(errorMessage, _exportJobRecord.FailureDetails.FailureReason);
         }
 
+        [Fact]
+        public async Task GivenAnExportJob_WhenSearchFailedWithRequestRateExceeded_ThenJobStatusShouldBeUpdatedToFailed()
+        {
+            _searchService.SearchAsync(
+                Arg.Any<string>(),
+                Arg.Any<IReadOnlyList<Tuple<string, string>>>(),
+                _cancellationToken,
+                true)
+                .Throws(new RequestRateExceededException(null));
+
+            SetupExportJobRecordAndOperationDataStore(_exportJobRecord, CancellationToken.None);
+
+            DateTimeOffset endTimestamp = DateTimeOffset.UtcNow;
+
+            using (Mock.Property(() => ClockResolver.TimeProvider, new Microsoft.Extensions.Time.Testing.FakeTimeProvider(endTimestamp)))
+            {
+                await _exportJobTask.ExecuteAsync(_exportJobRecord, _weakETag, _cancellationToken);
+            }
+
+            Assert.NotNull(_lastExportJobOutcome);
+            Assert.Equal(OperationStatus.Failed, _lastExportJobOutcome.JobRecord.Status);
+            Assert.Equal(endTimestamp, _lastExportJobOutcome.JobRecord.EndTime);
+        }
+
         private async Task RunTypeFilterTest(IList<ExportJobFilter> filters, string resourceTypes)
         {
             var exportJobRecordWithFormat = CreateExportJobRecord(

diff --git a/src/Microsoft.Health.Fhir.Core/Features/Operations/Export/ExportJobTask.cs b/src/Microsoft.Health.Fhir.Core/Features/Operations/Export/ExportJobTask.cs
@@ -231,6 +231,9 @@ public async Task ExecuteAsync(ExportJobRecord exportJobRecord, WeakETag weakETa
             catch (RequestRateExceededException rree)
             {
                 _logger.LogWarning(rree, "[JobId:{JobId}] Job failed due to RequestRateExceeded.", _exportJobRecord.Id);
+
+                _exportJobRecord.FailureDetails = new JobFailureDetails(rree.Message, HttpStatusCode.TooManyRequests);
+                await CompleteJobAsync(OperationStatus.Failed, cancellationToken);
             }
             catch (DestinationConnectionException dce)
             {

diff --git a/src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/CosmosFhirDataStore.cs b/src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/CosmosFhirDataStore.cs
@@ -679,7 +679,7 @@ public async Task<ResourceWrapper> UpdateSearchParameterIndicesAsync(ResourceWra
                 try
                 {
                     var prevPage = page;
-                    page = await cosmosQuery.ExecuteNextAsync(linkedTokenSource.Token);
+                    page = await _retryExceptionPolicyFactory.RetryPolicy.ExecuteAsync(() => cosmosQuery.ExecuteNextAsync(linkedTokenSource.Token));
 
                     if (mustNotExceedMaxItemCount && (page.Count + results.Count > totalDesiredCount))
                     {