Add retry mechanism to CosmosDB Database creation

src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj

@@ -24,6 +24,7 @@
     <PackageReference Include="Newtonsoft.Json" /> <!-- Required by Microsoft.Azure.Cosmos -->
     <PackageReference Include="Azure.Provisioning" />
     <PackageReference Include="Azure.Provisioning.CosmosDB" />
+    <PackageReference Include="Polly.Core" />
   </ItemGroup>
 
   <ItemGroup>

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,35 +103,54 @@ private static IResourceBuilder<AzureCosmosDBResource> RunAsEmulator(this IResou
             cosmosClient = CreateCosmosClient(connectionString);
         });
 
+        var creationState = HealthCheckResult.Unhealthy("Waiting for databases and containers to be created");
+
         builder.ApplicationBuilder.Eventing.Subscribe<ResourceReadyEvent>(builder.Resource, async (@event, ct) =>
         {
             if (cosmosClient is null)
             {
                 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<ResponseMessage> Pipeline { get; }
+            = new ResiliencePipelineBuilder<ResponseMessage>()
+                .AddRetry(new()
+                {
+                    MaxRetryAttempts = 10,
+                    Delay = TimeSpan.FromMilliseconds(500),
+                    BackoffType = DelayBackoffType.Constant,
+                    ShouldHandle = new PredicateBuilder<ResponseMessage>()
+                        .Handle<CosmosException>()
+                        .HandleResult(static result => !result.IsSuccessStatusCode),
+                })
+                .Build();
+
+        public override async Task<ResponseMessage> SendAsync(
+            RequestMessage request,
+            CancellationToken cancellationToken)
+        {
+            return await Pipeline
+                .ExecuteAsync(async ct => await base.SendAsync(request, ct).ConfigureAwait(false), cancellationToken)
+                .ConfigureAwait(false);
+        }
     }
 
     /// <summary>