diff --git a/src/Microsoft.Health.Fhir.CosmosDb.UnitTests/Features/Storage/Registry/CosmosDbSearchParameterStatusDataStoreTests.cs b/src/Microsoft.Health.Fhir.CosmosDb.UnitTests/Features/Storage/Registry/CosmosDbSearchParameterStatusDataStoreTests.cs new file mode 100644 index 0000000000..22db66f46d --- /dev/null +++ b/src/Microsoft.Health.Fhir.CosmosDb.UnitTests/Features/Storage/Registry/CosmosDbSearchParameterStatusDataStoreTests.cs @@ -0,0 +1,216 @@ +// ------------------------------------------------------------------------------------------------- +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License (MIT). See LICENSE in the repo root for license information. +// ------------------------------------------------------------------------------------------------- + +using System; +using System.Collections.Generic; +using System.Net; +using System.Text; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.Logging.Abstractions; +using Microsoft.Health.Core.Features.Context; +using Microsoft.Health.Extensions.DependencyInjection; +using Microsoft.Health.Fhir.Core.Features.Context; +using Microsoft.Health.Fhir.Core.Features.Search.Registry; +using Microsoft.Health.Fhir.CosmosDb.Core.Configs; +using Microsoft.Health.Fhir.CosmosDb.Features.Queries; +using Microsoft.Health.Fhir.CosmosDb.Features.Storage; +using Microsoft.Health.Fhir.CosmosDb.Features.Storage.Registry; +using Microsoft.Health.Fhir.Tests.Common; +using Microsoft.Health.Test.Utilities; +using NSubstitute; +using NSubstitute.ExceptionExtensions; +using Polly; +using Xunit; + +namespace Microsoft.Health.Fhir.CosmosDb.UnitTests.Features.Storage.Registry; + +[Trait(Traits.OwningTeam, OwningTeam.Fhir)] +[Trait(Traits.Category, Categories.Search)] +public class CosmosDbSearchParameterStatusDataStoreTests +{ + private readonly CosmosDbSearchParameterStatusDataStore _dataStore; + private readonly ICosmosQueryFactory _cosmosQueryFactory; + private readonly IScoped _containerScope; + private readonly CosmosDataStoreConfiguration _cosmosDataStoreConfiguration = new CosmosDataStoreConfiguration(); + private readonly IFhirRequestContext _fhirRequestContext; + private readonly RequestContextAccessor _requestContextAccessor; + private readonly RetryExceptionPolicyFactory _retryExceptionPolicyFactory; + + public CosmosDbSearchParameterStatusDataStoreTests() + { + _cosmosQueryFactory = Substitute.For(); + _containerScope = Substitute.For>(); + + _fhirRequestContext = Substitute.For(); + _requestContextAccessor = Substitute.For>(); + _requestContextAccessor.RequestContext.Returns(_fhirRequestContext); + _retryExceptionPolicyFactory = new RetryExceptionPolicyFactory(_cosmosDataStoreConfiguration, _requestContextAccessor, NullLogger.Instance); + + _dataStore = new CosmosDbSearchParameterStatusDataStore( + () => _containerScope, + new CosmosDataStoreConfiguration(), + _cosmosQueryFactory, + _retryExceptionPolicyFactory); + } + + [Fact] + public async Task GivenTransientError_WhenGettingSearchParameterStatusesAsBackgroundTask_ThenRetriesShouldBeAttempted() + { + // Arrange + _fhirRequestContext.IsBackgroundTask.Returns(true); + + var exception = new CosmosException( + message: "Service Unavailable", + statusCode: HttpStatusCode.ServiceUnavailable, + subStatusCode: 0, + activityId: Guid.NewGuid().ToString(), + requestCharge: 0); + + var mockQuery = Substitute.For>(); + int runs = 0; + + // Simulate failure on the first three attempts and success on fourth + mockQuery.ExecuteNextAsync(CancellationToken.None) + .ReturnsForAnyArgs(_ => + { + runs++; + if (runs < 4) + { + throw exception; + } + + // Return a mock FeedResponse on the fourth attempt + return Task.FromResult(CreateMockFeedResponse( + [ + new SearchParameterStatusWrapper { Uri = new Uri("http://example.com") }, + ])); + }); + + _cosmosQueryFactory.Create( + Arg.Any(), + Arg.Any()).Returns(mockQuery); + + // Act + await _dataStore.GetSearchParameterStatuses(CancellationToken.None); + + // Assert + await mockQuery.Received(4).ExecuteNextAsync(Arg.Any()); // 1 initial attempt + 3 retry + } + + [Fact] + public async Task GivenTransientError_WhenCheckingIfSearchParameterStatusUpdateIsRequiredAsBackgroundTask_ThenRetriesShouldBeAttempted() + { + // Arrange + _fhirRequestContext.IsBackgroundTask.Returns(true); + + var exception = new CosmosException( + message: "Service Unavailable", + statusCode: HttpStatusCode.ServiceUnavailable, + subStatusCode: 0, + activityId: Guid.NewGuid().ToString(), + requestCharge: 0); + + var mockQuery = Substitute.For>(); + int runs = 0; + + // Simulate failure on the first three attempts and success on the fourth + mockQuery.ExecuteNextAsync(Arg.Any()) + .ReturnsForAnyArgs(_ => + { + runs++; + if (runs < 4) + { + throw exception; + } + + // Return a mock FeedResponse on the fourth attempt + return Task.FromResult(CreateMockFeedResponse(new List + { + new() + { + Count = 5, + LastUpdated = DateTimeOffset.UtcNow.AddMinutes(1), + }, + })); + }); + + _cosmosQueryFactory.Create( + Arg.Any(), + Arg.Any()).Returns(mockQuery); + + var container = Substitute.For>(); + container.Value.Returns(Substitute.For()); + + // Act + var result = await _dataStore.CheckIfSearchParameterStatusUpdateRequiredAsync( + container, + currentCount: 4, + lastRefreshed: DateTimeOffset.UtcNow, + CancellationToken.None); + + // Assert + Assert.True(result); // Should return true due to mismatched count and timestamp + await mockQuery.Received(4).ExecuteNextAsync(Arg.Any()); // 1 initial attempt + 3 retries + } + + [Fact] + public async Task GivenBatchExecutionError_WhenUpsertingStatusesAsBackgroundTask_ThenRetriesShouldBeAttempted() + { + // Arrange + _fhirRequestContext.IsBackgroundTask.Returns(true); + + var exception = new CosmosException( + message: "Service Unavailable", + statusCode: HttpStatusCode.ServiceUnavailable, + subStatusCode: 0, + activityId: Guid.NewGuid().ToString(), + requestCharge: 0); + + var mockContainer = Substitute.For(); + var mockTransactionalBatch = Substitute.For(); + + int runs = 0; + + // Simulate failure on the first three attempts and success on the fourth + mockTransactionalBatch.ExecuteAsync(Arg.Any()) + .ReturnsForAnyArgs(_ => + { + runs++; + if (runs < 4) + { + throw exception; + } + + // Return a mock TransactionalBatchResponse on the fourth attempt + return Task.FromResult(Substitute.For()); + }); + + mockContainer.CreateTransactionalBatch(Arg.Any()) + .Returns(mockTransactionalBatch); + + _containerScope.Value.Returns(mockContainer); + + var statuses = new List + { + new() { Uri = new Uri("http://example.com") }, + }; + + // Act + await _dataStore.UpsertStatuses(statuses, CancellationToken.None); + + // Assert + await mockTransactionalBatch.Received(4).ExecuteAsync(Arg.Any()); // 3 failure + 1 success + } + + private static FeedResponse CreateMockFeedResponse(List items) + { + var feedResponse = Substitute.For>(); + feedResponse.Count.Returns(items.Count); + feedResponse.GetEnumerator().Returns(items.GetEnumerator()); + return feedResponse; + } +} diff --git a/src/Microsoft.Health.Fhir.CosmosDb/AssemblyInfo.cs b/src/Microsoft.Health.Fhir.CosmosDb/AssemblyInfo.cs index 84819ace40..78e2155040 100644 --- a/src/Microsoft.Health.Fhir.CosmosDb/AssemblyInfo.cs +++ b/src/Microsoft.Health.Fhir.CosmosDb/AssemblyInfo.cs @@ -11,5 +11,5 @@ [assembly: InternalsVisibleTo("Microsoft.Health.Fhir.R4.Tests.Integration")] [assembly: InternalsVisibleTo("Microsoft.Health.Fhir.R4B.Tests.Integration")] [assembly: InternalsVisibleTo("Microsoft.Health.Fhir.R5.Tests.Integration")] -[assembly: InternalsVisibleTo("DynamicProxyGenAssembly2, PublicKey=0024000004800000940000000602000000240000525341310004000001000100c547cac37abd99c8db225ef2f6c8a3602f3b3606cc9891605d02baa56104f4cfc0734aa39b93bf7852f7d9266654753cc297e7d2edfe0bac1cdcf9f717241550e0a7b191195b7667bb4f64bcb8e2121380fd1d9d46ad2d92d2d15605093924cceaf74c4861eff62abf69b9291ed0a340e113be11e6a7d3113e92484cf7045cc7")] +[assembly: InternalsVisibleTo("DynamicProxyGenAssembly2")] [assembly: NeutralResourcesLanguage("en-us")] diff --git a/src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/Registry/CosmosDbSearchParameterStatusDataStore.cs b/src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/Registry/CosmosDbSearchParameterStatusDataStore.cs index 302ccf07b3..5513e8fdf6 100644 --- a/src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/Registry/CosmosDbSearchParameterStatusDataStore.cs +++ b/src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/Registry/CosmosDbSearchParameterStatusDataStore.cs @@ -25,18 +25,22 @@ public sealed class CosmosDbSearchParameterStatusDataStore : ISearchParameterSta private DateTimeOffset? _lastRefreshed = null; private List _statusList = new(); private readonly SemaphoreSlim _statusListSemaphore = new(1, 1); + private readonly RetryExceptionPolicyFactory _retryExceptionPolicyFactory; public CosmosDbSearchParameterStatusDataStore( Func> containerScopeFactory, CosmosDataStoreConfiguration cosmosDataStoreConfiguration, - ICosmosQueryFactory queryFactory) + ICosmosQueryFactory queryFactory, + RetryExceptionPolicyFactory retryExceptionPolicyFactory) { EnsureArg.IsNotNull(containerScopeFactory, nameof(containerScopeFactory)); EnsureArg.IsNotNull(cosmosDataStoreConfiguration, nameof(cosmosDataStoreConfiguration)); EnsureArg.IsNotNull(queryFactory, nameof(queryFactory)); + EnsureArg.IsNotNull(retryExceptionPolicyFactory, nameof(retryExceptionPolicyFactory)); _containerScopeFactory = containerScopeFactory; _queryFactory = queryFactory; + _retryExceptionPolicyFactory = retryExceptionPolicyFactory; } public async Task> GetSearchParameterStatuses(CancellationToken cancellationToken) @@ -77,7 +81,10 @@ public async Task> GetSearchP do { - FeedResponse results = await query.ExecuteNextAsync(cancellationToken); + FeedResponse results = await _retryExceptionPolicyFactory.RetryPolicy.ExecuteAsync( + async ct => await query.ExecuteNextAsync( + ct), + cancellationToken); parameterStatus.AddRange(results.Select(x => x.ToSearchParameterStatus())); } @@ -101,7 +108,7 @@ public async Task> GetSearchP } } - private async Task CheckIfSearchParameterStatusUpdateRequiredAsync(IScoped container, int currentCount, DateTimeOffset lastRefreshed, CancellationToken cancellationToken) + internal async Task CheckIfSearchParameterStatusUpdateRequiredAsync(IScoped container, int currentCount, DateTimeOffset lastRefreshed, CancellationToken cancellationToken) { var lastUpdatedQuery = _queryFactory.Create( container.Value, @@ -113,7 +120,11 @@ private async Task CheckIfSearchParameterStatusUpdateRequiredAsync(IScoped MaxItemCount = 1, })); - FeedResponse lastUpdatedResponse = await lastUpdatedQuery.ExecuteNextAsync(cancellationToken); + FeedResponse lastUpdatedResponse = await _retryExceptionPolicyFactory.RetryPolicy.ExecuteAsync( + async ct => await lastUpdatedQuery.ExecuteNextAsync( + ct), + cancellationToken); + var result = lastUpdatedResponse?.FirstOrDefault(); if (result == null || result.Count != currentCount || result.LastUpdated > lastRefreshed) @@ -144,7 +155,9 @@ public async Task UpsertStatuses(IReadOnlyCollection await batch.ExecuteAsync(ct), + cancellationToken); } } @@ -158,7 +171,7 @@ public void Dispose() _statusListSemaphore?.Dispose(); } - private class CacheQueryResponse + internal class CacheQueryResponse { public int Count { get; set; } diff --git a/test/Microsoft.Health.Fhir.Shared.Tests.Integration/Persistence/CosmosDbFhirStorageTestsFixture.cs b/test/Microsoft.Health.Fhir.Shared.Tests.Integration/Persistence/CosmosDbFhirStorageTestsFixture.cs index aa7256ace8..696759508f 100644 --- a/test/Microsoft.Health.Fhir.Shared.Tests.Integration/Persistence/CosmosDbFhirStorageTestsFixture.cs +++ b/test/Microsoft.Health.Fhir.Shared.Tests.Integration/Persistence/CosmosDbFhirStorageTestsFixture.cs @@ -207,7 +207,8 @@ public virtual async Task InitializeAsync() _searchParameterStatusDataStore = new CosmosDbSearchParameterStatusDataStore( () => documentClient, _cosmosDataStoreConfiguration, - cosmosDocumentQueryFactory); + cosmosDocumentQueryFactory, + new RetryExceptionPolicyFactory(_cosmosDataStoreConfiguration, _fhirRequestContextAccessor, NullLogger.Instance)); var bundleConfiguration = new BundleConfiguration() { SupportsBundleOrchestrator = true }; var bundleOptions = Substitute.For>();