Update health check to ensure blob containers created at right time (#9159)

* Update health check to ensure blob containers created at right time

Resolves #9139
Resolves #9145

* fixup! Update health check to ensure blob containers created at right time

* fixup! Update health check to ensure blob containers created at right time

* fixup! Update health check to ensure blob containers created at right time

* Prevent multiple container checks

* Use better variable name

* Move container creation to RunAsEmulator

* Register single hc for blobs

* Reuse blobserviceclient

* Fix registrations

* Remove custom heathcheck

* Remove custom healthcheck

* Remove specific health checks

* Add blob and container health checks

* Fix test

* Feedback

* Improve comment

---------

Co-authored-by: Sebastien Ros <sebastienros@gmail.com>

Author: RussKie

src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs

@@ -121,7 +121,6 @@ private static IResourceBuilder<AzureCosmosDBResource> RunAsEmulator(this IResou
             }
         });
 
-        // Use custom health check that also seeds the databases and containers
         var healthCheckKey = $"{builder.Resource.Name}_check";
         builder.ApplicationBuilder.Services.AddHealthChecks().AddAzureCosmosDB(
             sp => cosmosClient ?? throw new InvalidOperationException("CosmosClient is not initialized."),

src/Aspire.Hosting.Azure.Storage/AzureStorageExtensions.cs

@@ -95,6 +95,7 @@ public static IResourceBuilder<AzureStorageResource> AddAzureStorage(this IDistr
         };
 
         var resource = new AzureStorageResource(name, configureInfrastructure);
+
         return builder.AddResource(resource)
             .WithDefaultRoleAssignments(StorageBuiltInRole.GetBuiltInRoleName,
                 StorageBuiltInRole.StorageBlobDataContributor,
@@ -131,34 +132,30 @@ public static IResourceBuilder<AzureStorageResource> RunAsEmulator(this IResourc
                });
 
         BlobServiceClient? blobServiceClient = null;
-
         builder.ApplicationBuilder.Eventing.Subscribe<BeforeResourceStartedEvent>(builder.Resource, async (@event, ct) =>
         {
-            var connectionString = await builder.Resource.GetBlobConnectionString().GetValueAsync(ct).ConfigureAwait(false);
-            if (connectionString is null)
-            {
-                throw new DistributedApplicationException($"BeforeResourceStartedEvent was published for the '{builder.Resource.Name}' resource but the connection string was null.");
-            }
+            // The BlobServiceClient is created before the health check is run.
+            // We can't use ConnectionStringAvailableEvent here because the resource doesn't have a connection string, so
+            // we use BeforeResourceStartedEvent
 
+            var connectionString = await builder.Resource.GetBlobConnectionString().GetValueAsync(ct).ConfigureAwait(false) ?? throw new DistributedApplicationException($"{nameof(ConnectionStringAvailableEvent)} was published for the '{builder.Resource.Name}' resource but the connection string was null.");
             blobServiceClient = CreateBlobServiceClient(connectionString);
         });
 
         builder.ApplicationBuilder.Eventing.Subscribe<ResourceReadyEvent>(builder.Resource, async (@event, ct) =>
         {
-            if (blobServiceClient is null)
-            {
-                throw new DistributedApplicationException($"BlobServiceClient was not created for the '{builder.Resource.Name}' resource.");
-            }
+            // The ResourceReadyEvent of a resource is triggered after its health check is healthy.
+            // This means we can safely use this event to create the blob containers.
 
-            var connectionString = await builder.Resource.GetBlobConnectionString().GetValueAsync(ct).ConfigureAwait(false);
-            if (connectionString is null)
+            if (blobServiceClient is null)
             {
-                throw new DistributedApplicationException($"ResourceReadyEvent was published for the '{builder.Resource.Name}' resource but the connection string was null.");
+                throw new InvalidOperationException("BlobServiceClient is not initialized.");
             }
 
-            foreach (var blobContainer in builder.Resource.BlobContainers)
+            foreach (var container in builder.Resource.BlobContainers)
             {
-                await blobServiceClient.GetBlobContainerClient(blobContainer.BlobContainerName).CreateIfNotExistsAsync(cancellationToken: ct).ConfigureAwait(false);
+                var blobContainerClient = blobServiceClient.GetBlobContainerClient(container.BlobContainerName);
+                await blobContainerClient.CreateIfNotExistsAsync(cancellationToken: ct).ConfigureAwait(false);
             }
         });
 
@@ -182,18 +179,6 @@ public static IResourceBuilder<AzureStorageResource> RunAsEmulator(this IResourc
         configureContainer?.Invoke(surrogateBuilder);
 
         return builder;
-
-        static BlobServiceClient CreateBlobServiceClient(string connectionString)
-        {
-            if (Uri.TryCreate(connectionString, UriKind.Absolute, out var uri))
-            {
-                return new BlobServiceClient(uri, new DefaultAzureCredential());
-            }
-            else
-            {
-                return new BlobServiceClient(connectionString);
-            }
-        }
     }
 
     /// <summary>
@@ -308,7 +293,22 @@ public static IResourceBuilder<AzureBlobStorageResource> AddBlobs(this IResource
         ArgumentException.ThrowIfNullOrEmpty(name);
 
         var resource = new AzureBlobStorageResource(name, builder.Resource);
-        return builder.ApplicationBuilder.AddResource(resource);
+
+        string? connectionString = null;
+        builder.ApplicationBuilder.Eventing.Subscribe<ConnectionStringAvailableEvent>(resource, async (@event, ct) =>
+        {
+            connectionString = await resource.ConnectionStringExpression.GetValueAsync(ct).ConfigureAwait(false);
+        });
+
+        var healthCheckKey = $"{resource.Name}_check";
+
+        BlobServiceClient? blobServiceClient = null;
+        builder.ApplicationBuilder.Services.AddHealthChecks().AddAzureBlobStorage(sp =>
+        {
+            return blobServiceClient ??= CreateBlobServiceClient(connectionString ?? throw new InvalidOperationException("Connection string is not initialized."));
+        }, name: healthCheckKey);
+
+        return builder.ApplicationBuilder.AddResource(resource).WithHealthCheck(healthCheckKey);
     }
 
     /// <summary>
@@ -326,10 +326,24 @@ public static IResourceBuilder<AzureBlobStorageContainerResource> AddBlobContain
         blobContainerName ??= name;
 
         AzureBlobStorageContainerResource resource = new(name, blobContainerName, builder.Resource);
-
         builder.Resource.Parent.BlobContainers.Add(resource);
 
-        return builder.ApplicationBuilder.AddResource(resource);
+        string? connectionString = null;
+        builder.ApplicationBuilder.Eventing.Subscribe<ConnectionStringAvailableEvent>(resource, async (@event, ct) =>
+        {
+            connectionString = await resource.Parent.ConnectionStringExpression.GetValueAsync(ct).ConfigureAwait(false);
+        });
+
+        var healthCheckKey = $"{resource.Name}_check";
+
+        BlobServiceClient? blobServiceClient = null;
+        builder.ApplicationBuilder.Services.AddHealthChecks().AddAzureBlobStorage(
+            sp => blobServiceClient ??= CreateBlobServiceClient(connectionString ?? throw new InvalidOperationException("Connection string is not initialized.")),
+            optionsFactory: sp => new HealthChecks.Azure.Storage.Blobs.AzureBlobStorageHealthCheckOptions { ContainerName = blobContainerName },
+            name: healthCheckKey);
+
+        return builder.ApplicationBuilder
+            .AddResource(resource).WithHealthCheck(healthCheckKey);
     }
 
     /// <summary>
@@ -362,6 +376,18 @@ public static IResourceBuilder<AzureQueueStorageResource> AddQueues(this IResour
         return builder.ApplicationBuilder.AddResource(resource);
     }
 
+    private static BlobServiceClient CreateBlobServiceClient(string connectionString)
+    {
+        if (Uri.TryCreate(connectionString, UriKind.Absolute, out var uri))
+        {
+            return new BlobServiceClient(uri, new DefaultAzureCredential());
+        }
+        else
+        {
+            return new BlobServiceClient(connectionString);
+        }
+    }
+
     /// <summary>
     /// Assigns the specified roles to the given resource, granting it the necessary permissions
     /// on the target Azure Storage account. This replaces the default role assignments for the resource.

src/Components/Aspire.Azure.Storage.Blobs/AspireBlobStorageExtensions.BlobStorageContainerComponent.AzureBlobStorageContainerHealthCheck.cs

@@ -6,6 +6,8 @@
 using Azure.Core;
 using Azure.Core.Extensions;
 using Azure.Storage.Blobs;
+using Azure.Storage.Blobs.Specialized;
+using HealthChecks.Azure.Storage.Blobs;
 using Microsoft.Extensions.Azure;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
@@ -55,7 +57,7 @@ protected override void BindSettingsToConfiguration(AzureBlobStorageContainerSet
         }
 
         protected override IHealthCheck CreateHealthCheck(BlobContainerClient client, AzureBlobStorageContainerSettings settings)
-            => new AzureBlobStorageContainerHealthCheck(client);
+            => new AzureBlobStorageHealthCheck(client.GetParentBlobServiceClient(), new AzureBlobStorageHealthCheckOptions { ContainerName = client.Name });
 
         protected override bool GetHealthCheckEnabled(AzureBlobStorageContainerSettings settings)
             => !settings.DisableHealthChecks;

src/Components/Aspire.Azure.Storage.Blobs/AspireBlobStorageExtensions.BlobStorageContainerComponent.cs

@@ -1503,6 +1503,10 @@ public async Task AddAzureStorageViaPublishMode()
                 };
             });
 
+        var blob = storage.AddBlobs("blob");
+        var queue = storage.AddQueues("queue");
+        var table = storage.AddTables("table");
+
         storage.Resource.Outputs["blobEndpoint"] = "https://myblob";
         storage.Resource.Outputs["queueEndpoint"] = "https://myqueue";
         storage.Resource.Outputs["tableEndpoint"] = "https://mytable";
@@ -1578,7 +1582,6 @@ param principalId string
         Assert.Equal(expectedBicep, storageRolesManifest.BicepText);
 
         // Check blob resource.
-        var blob = storage.AddBlobs("blob");
 
         var connectionStringBlobResource = (IResourceWithConnectionString)blob.Resource;
 
@@ -1593,8 +1596,6 @@ param principalId string
         Assert.Equal(expectedBlobManifest, blobManifest.ToString());
 
         // Check queue resource.
-        var queue = storage.AddQueues("queue");
-
         var connectionStringQueueResource = (IResourceWithConnectionString)queue.Resource;
 
         Assert.Equal("https://myqueue", await connectionStringQueueResource.GetConnectionStringAsync());
@@ -1608,8 +1609,6 @@ param principalId string
         Assert.Equal(expectedQueueManifest, queueManifest.ToString());
 
         // Check table resource.
-        var table = storage.AddTables("table");
-
         var connectionStringTableResource = (IResourceWithConnectionString)table.Resource;
 
         Assert.Equal("https://mytable", await connectionStringTableResource.GetConnectionStringAsync());

tests/Aspire.Hosting.Azure.Tests/AzureBicepResourceTests.cs

@@ -135,9 +135,10 @@ public async Task VerifyAzureStorageEmulatorResource()
 
     [Fact]
     [RequiresDocker]
-    [QuarantinedTest("https://github.com/dotnet/aspire/issues/9139")]
     public async Task VerifyAzureStorageEmulator_blobcontainer_auto_created()
     {
+        var cts = new CancellationTokenSource(TimeSpan.FromMinutes(3));
+
         using var builder = TestDistributedApplicationBuilder.Create().WithTestAndResourceLogging(testOutputHelper);
         var storage = builder.AddAzureStorage("storage").RunAsEmulator();
         var blobs = storage.AddBlobs("BlobConnection");
@@ -146,16 +147,16 @@ public async Task VerifyAzureStorageEmulator_blobcontainer_auto_created()
         using var app = builder.Build();
         await app.StartAsync();
 
+        var rns = app.Services.GetRequiredService<ResourceNotificationService>();
+        await rns.WaitForResourceHealthyAsync(blobContainer.Resource.Name, cancellationToken: cts.Token);
+
         var hb = Host.CreateApplicationBuilder();
         hb.Configuration["ConnectionStrings:BlobConnection"] = await blobs.Resource.ConnectionStringExpression.GetValueAsync(CancellationToken.None);
         hb.AddAzureBlobClient("BlobConnection");
 
         using var host = hb.Build();
         await host.StartAsync();
 
-        var rns = app.Services.GetRequiredService<ResourceNotificationService>();
-        await rns.WaitForResourceHealthyAsync(blobContainer.Resource.Name, CancellationToken.None);
-
         var serviceClient = host.Services.GetRequiredService<BlobServiceClient>();
         var blobContainerClient = serviceClient.GetBlobContainerClient("testblobcontainer");