diff --git a/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj b/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj index 0806b9a815e..a2aade08516 100644 --- a/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj +++ b/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj @@ -24,6 +24,7 @@ + diff --git a/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs b/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs index 782961c8b38..d754874a09d 100644 --- a/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs +++ b/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs @@ -14,6 +14,8 @@ using Azure.Provisioning.KeyVault; using Microsoft.Azure.Cosmos; using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Diagnostics.HealthChecks; +using Polly; namespace Aspire.Hosting; @@ -101,6 +103,8 @@ private static IResourceBuilder RunAsEmulator(this IResou cosmosClient = CreateCosmosClient(connectionString); }); + var creationState = HealthCheckResult.Unhealthy("Waiting for databases and containers to be created"); + builder.ApplicationBuilder.Eventing.Subscribe(builder.Resource, async (@event, ct) => { if (cosmosClient is null) @@ -108,28 +112,45 @@ private static IResourceBuilder RunAsEmulator(this IResou throw new InvalidOperationException("CosmosClient is not initialized."); } - await cosmosClient.ReadAccountAsync().WaitAsync(ct).ConfigureAwait(false); - - foreach (var database in builder.Resource.Databases) + try { - var db = (await cosmosClient.CreateDatabaseIfNotExistsAsync(database.DatabaseName, cancellationToken: ct).ConfigureAwait(false)).Database; + await cosmosClient.ReadAccountAsync().WaitAsync(ct).ConfigureAwait(false); - foreach (var container in database.Containers) + foreach (var database in builder.Resource.Databases) { - var containerProperties = container.ContainerProperties; + var db = (await cosmosClient.CreateDatabaseIfNotExistsAsync(database.DatabaseName, cancellationToken: ct).ConfigureAwait(false)).Database; + + foreach (var container in database.Containers) + { + var containerProperties = container.ContainerProperties ?? new ContainerProperties + { + Id = container.ContainerName, + PartitionKeyPaths = container.PartitionKeyPaths + }; - await db.CreateContainerIfNotExistsAsync(containerProperties, cancellationToken: ct).ConfigureAwait(false); + await db.CreateContainerIfNotExistsAsync(containerProperties, cancellationToken: ct).ConfigureAwait(false); + } } + creationState = HealthCheckResult.Healthy(); + } + catch (Exception ex) + { + creationState = HealthCheckResult.Degraded("Could not create databases and containers", ex); + throw; } }); var healthCheckKey = $"{builder.Resource.Name}_check"; - builder.ApplicationBuilder.Services.AddHealthChecks().AddAzureCosmosDB( - sp => cosmosClient ?? throw new InvalidOperationException("CosmosClient is not initialized."), - name: healthCheckKey - ); + var creationHealthCheckKey = $"{builder.Resource.Name}_databases"; + builder.ApplicationBuilder.Services.AddHealthChecks() + .AddAzureCosmosDB( + sp => cosmosClient ?? throw new InvalidOperationException("CosmosClient is not initialized."), + name: healthCheckKey + ) + .AddCheck(creationHealthCheckKey, () => creationState); - builder.WithHealthCheck(healthCheckKey); + builder.WithHealthCheck(healthCheckKey) + .WithHealthCheck(creationHealthCheckKey); if (configureContainer != null) { @@ -155,11 +176,37 @@ static CosmosClient CreateCosmosClient(string connectionString) { clientOptions.ConnectionMode = ConnectionMode.Gateway; clientOptions.LimitToEndpoint = true; + clientOptions.CustomHandlers.Add(new CosmosClientRetryHandler()); } return new CosmosClient(connectionString, clientOptions); } } + + } + class CosmosClientRetryHandler : RequestHandler + { + private static ResiliencePipeline Pipeline { get; } + = new ResiliencePipelineBuilder() + .AddRetry(new() + { + MaxRetryAttempts = 10, + Delay = TimeSpan.FromMilliseconds(500), + BackoffType = DelayBackoffType.Constant, + ShouldHandle = new PredicateBuilder() + .Handle() + .HandleResult(static result => !result.IsSuccessStatusCode), + }) + .Build(); + + public override async Task SendAsync( + RequestMessage request, + CancellationToken cancellationToken) + { + return await Pipeline + .ExecuteAsync(async ct => await base.SendAsync(request, ct).ConfigureAwait(false), cancellationToken) + .ConfigureAwait(false); + } } ///