Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add retry to Cosmos search param operations #4790

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add retry to Cosmos search param operations #4790

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
216 changes: 216 additions & 0 deletions ...smosDb.UnitTests/Features/Storage/Registry/CosmosDbSearchParameterStatusDataStoreTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<Container> _containerScope;
private readonly CosmosDataStoreConfiguration _cosmosDataStoreConfiguration = new CosmosDataStoreConfiguration();
private readonly IFhirRequestContext _fhirRequestContext;
private readonly RequestContextAccessor<IFhirRequestContext> _requestContextAccessor;
private readonly RetryExceptionPolicyFactory _retryExceptionPolicyFactory;

public CosmosDbSearchParameterStatusDataStoreTests()
{
_cosmosQueryFactory = Substitute.For<ICosmosQueryFactory>();
_containerScope = Substitute.For<IScoped<Container>>();

_fhirRequestContext = Substitute.For<IFhirRequestContext>();
_requestContextAccessor = Substitute.For<RequestContextAccessor<IFhirRequestContext>>();
_requestContextAccessor.RequestContext.Returns(_fhirRequestContext);
_retryExceptionPolicyFactory = new RetryExceptionPolicyFactory(_cosmosDataStoreConfiguration, _requestContextAccessor, NullLogger<RetryExceptionPolicyFactory>.Instance);

_dataStore = new CosmosDbSearchParameterStatusDataStore(
() => _containerScope,
new CosmosDataStoreConfiguration(),
_cosmosQueryFactory,
_retryExceptionPolicyFactory);
}

[Fact]
public async Task GivenTransientError_WhenGettingSearchParameterStatusesAsBackgroundTask_ThenRetriesShouldBeAttempted()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Your subject under test is a TransientError?

{
// 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<ICosmosQuery<SearchParameterStatusWrapper>>();
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<SearchParameterStatusWrapper>(
Arg.Any<Container>(),
Arg.Any<CosmosQueryContext>()).Returns(mockQuery);

// Act
await _dataStore.GetSearchParameterStatuses(CancellationToken.None);

// Assert
await mockQuery.Received(4).ExecuteNextAsync(Arg.Any<CancellationToken>()); // 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<ICosmosQuery<CosmosDbSearchParameterStatusDataStore.CacheQueryResponse>>();
int runs = 0;

// Simulate failure on the first three attempts and success on the fourth
mockQuery.ExecuteNextAsync(Arg.Any<CancellationToken>())
.ReturnsForAnyArgs(_ =>
{
runs++;
if (runs < 4)
{
throw exception;
}

// Return a mock FeedResponse on the fourth attempt
return Task.FromResult(CreateMockFeedResponse(new List<CosmosDbSearchParameterStatusDataStore.CacheQueryResponse>
{
new()
{
Count = 5,
LastUpdated = DateTimeOffset.UtcNow.AddMinutes(1),
},
}));
});

_cosmosQueryFactory.Create<CosmosDbSearchParameterStatusDataStore.CacheQueryResponse>(
Arg.Any<Container>(),
Arg.Any<CosmosQueryContext>()).Returns(mockQuery);

var container = Substitute.For<IScoped<Container>>();
container.Value.Returns(Substitute.For<Container>());

// 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<CancellationToken>()); // 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<Container>();
var mockTransactionalBatch = Substitute.For<TransactionalBatch>();

int runs = 0;

// Simulate failure on the first three attempts and success on the fourth
mockTransactionalBatch.ExecuteAsync(Arg.Any<CancellationToken>())
.ReturnsForAnyArgs(_ =>
{
runs++;
if (runs < 4)
{
throw exception;
}

// Return a mock TransactionalBatchResponse on the fourth attempt
return Task.FromResult(Substitute.For<TransactionalBatchResponse>());
});

mockContainer.CreateTransactionalBatch(Arg.Any<PartitionKey>())
.Returns(mockTransactionalBatch);

_containerScope.Value.Returns(mockContainer);

var statuses = new List<ResourceSearchParameterStatus>
{
new() { Uri = new Uri("http://example.com") },
};

// Act
await _dataStore.UpsertStatuses(statuses, CancellationToken.None);

// Assert
await mockTransactionalBatch.Received(4).ExecuteAsync(Arg.Any<CancellationToken>()); // 3 failure + 1 success
}

private static FeedResponse<T> CreateMockFeedResponse<T>(List<T> items)
{
var feedResponse = Substitute.For<FeedResponse<T>>();
feedResponse.Count.Returns(items.Count);
feedResponse.GetEnumerator().Returns(items.GetEnumerator());
return feedResponse;
}
}
2 changes: 1 addition & 1 deletion src/Microsoft.Health.Fhir.CosmosDb/AssemblyInfo.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
25 changes: 19 additions & 6 deletions ....Health.Fhir.CosmosDb/Features/Storage/Registry/CosmosDbSearchParameterStatusDataStore.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,18 +25,22 @@ public sealed class CosmosDbSearchParameterStatusDataStore : ISearchParameterSta
private DateTimeOffset? _lastRefreshed = null;
private List<ResourceSearchParameterStatus> _statusList = new();
private readonly SemaphoreSlim _statusListSemaphore = new(1, 1);
private readonly RetryExceptionPolicyFactory _retryExceptionPolicyFactory;

public CosmosDbSearchParameterStatusDataStore(
Func<IScoped<Container>> 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<IReadOnlyCollection<ResourceSearchParameterStatus>> GetSearchParameterStatuses(CancellationToken cancellationToken)
Expand Down Expand Up @@ -77,7 +81,10 @@ public async Task<IReadOnlyCollection<ResourceSearchParameterStatus>> GetSearchP

do
{
FeedResponse<SearchParameterStatusWrapper> results = await query.ExecuteNextAsync(cancellationToken);
FeedResponse<SearchParameterStatusWrapper> results = await _retryExceptionPolicyFactory.RetryPolicy.ExecuteAsync(
async ct => await query.ExecuteNextAsync(
ct),
cancellationToken);

parameterStatus.AddRange(results.Select(x => x.ToSearchParameterStatus()));
}
Expand All @@ -101,7 +108,7 @@ public async Task<IReadOnlyCollection<ResourceSearchParameterStatus>> GetSearchP
}
}

private async Task<bool> CheckIfSearchParameterStatusUpdateRequiredAsync(IScoped<Container> container, int currentCount, DateTimeOffset lastRefreshed, CancellationToken cancellationToken)
internal async Task<bool> CheckIfSearchParameterStatusUpdateRequiredAsync(IScoped<Container> container, int currentCount, DateTimeOffset lastRefreshed, CancellationToken cancellationToken)
{
var lastUpdatedQuery = _queryFactory.Create<CacheQueryResponse>(
container.Value,
Expand All @@ -113,7 +120,11 @@ private async Task<bool> CheckIfSearchParameterStatusUpdateRequiredAsync(IScoped
MaxItemCount = 1,
}));

FeedResponse<CacheQueryResponse> lastUpdatedResponse = await lastUpdatedQuery.ExecuteNextAsync(cancellationToken);
FeedResponse<CacheQueryResponse> 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)
Expand Down Expand Up @@ -144,7 +155,9 @@ public async Task UpsertStatuses(IReadOnlyCollection<ResourceSearchParameterStat
batch.UpsertItem(status);
}

await batch.ExecuteAsync(cancellationToken);
await _retryExceptionPolicyFactory.RetryPolicy.ExecuteAsync(
async ct => await batch.ExecuteAsync(ct),
cancellationToken);
}
}

Expand All @@ -158,7 +171,7 @@ public void Dispose()
_statusListSemaphore?.Dispose();
}

private class CacheQueryResponse
internal class CacheQueryResponse
{
public int Count { get; set; }

Expand Down
3 changes: 2 additions & 1 deletion ...osoft.Health.Fhir.Shared.Tests.Integration/Persistence/CosmosDbFhirStorageTestsFixture.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,8 @@ public virtual async Task InitializeAsync()
_searchParameterStatusDataStore = new CosmosDbSearchParameterStatusDataStore(
() => documentClient,
_cosmosDataStoreConfiguration,
cosmosDocumentQueryFactory);
cosmosDocumentQueryFactory,
new RetryExceptionPolicyFactory(_cosmosDataStoreConfiguration, _fhirRequestContextAccessor, NullLogger<RetryExceptionPolicyFactory>.Instance));

var bundleConfiguration = new BundleConfiguration() { SupportsBundleOrchestrator = true };
var bundleOptions = Substitute.For<IOptions<BundleConfiguration>>();
Expand Down
Loading