WaitFor integration with Health Check Service

playground/TestShop/TestShop.AppHost/Program.cs

@@ -1,9 +1,10 @@
 var builder = DistributedApplication.CreateBuilder(args);
 
-var catalogDb = builder.AddPostgres("postgres")
-                       .WithDataVolume()
-                       .WithPgAdmin()
-                       .AddDatabase("catalogdb");
+var postgres = builder.AddPostgres("postgres")
+                      .WithDataVolume()
+                      .WithPgAdmin();
+
+var catalogDb = postgres.AddDatabase("catalogdb");
 
 var basketCache = builder.AddRedis("basketcache")
                          .WithDataVolume();
@@ -16,7 +17,8 @@
 #endif
 
 var catalogDbApp = builder.AddProject<Projects.CatalogDb>("catalogdbapp")
-                          .WithReference(catalogDb);
+                          .WithReference(catalogDb)
+                          .WaitFor(postgres);
 
 var catalogService = builder.AddProject<Projects.CatalogService>("catalogservice")
                             .WithReference(catalogDb)

src/Aspire.Hosting.PostgreSQL/Aspire.Hosting.PostgreSQL.csproj

@@ -17,6 +17,10 @@
     <Compile Include="$(SharedDir)VolumeNameGenerator.cs" Link="Utils\VolumeNameGenerator.cs" />
   </ItemGroup>
 
+  <ItemGroup>
+    <PackageReference Include="AspNetCore.HealthChecks.NpgSql" />
+  </ItemGroup>
+
   <ItemGroup>
     <ProjectReference Include="..\Aspire.Hosting\Aspire.Hosting.csproj" />
   </ItemGroup>

src/Aspire.Hosting.PostgreSQL/PostgresBuilderExtensions.cs

@@ -6,6 +6,7 @@
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Postgres;
 using Aspire.Hosting.Utils;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace Aspire.Hosting;
 
@@ -38,6 +39,17 @@ public static IResourceBuilder<PostgresServerResource> AddPostgres(this IDistrib
         var passwordParameter = password?.Resource ?? ParameterResourceBuilderExtensions.CreateDefaultPasswordParameter(builder, $"{name}-password");
 
         var postgresServer = new PostgresServerResource(name, userName?.Resource, passwordParameter);
+
+        string? connectionString = null;
+
+        builder.Eventing.Subscribe<AfterEndpointsAllocatedEvent>(async (@event, ct) =>
+        {
+            connectionString = await postgresServer.GetConnectionStringAsync(ct).ConfigureAwait(false);
+        });
+
+        builder.Services.AddHealthChecks()
+                        .AddNpgSql(sp => connectionString!, name: $"{name}_check");
+
         return builder.AddResource(postgresServer)
                       .WithEndpoint(port: port, targetPort: 5432, name: PostgresServerResource.PrimaryEndpointName) // Internal port is always 5432.
                       .WithImage(PostgresContainerImageTags.Image, PostgresContainerImageTags.Tag)
@@ -48,7 +60,8 @@ public static IResourceBuilder<PostgresServerResource> AddPostgres(this IDistrib
                       {
                           context.EnvironmentVariables[UserEnvVarName] = postgresServer.UserNameReference;
                           context.EnvironmentVariables[PasswordEnvVarName] = postgresServer.PasswordParameter;
-                      });
+                      })
+                      .WithAnnotation(new HealthCheckAnnotation($"{postgresServer.Name}_check"));
     }
 
     /// <summary>

src/Aspire.Hosting.Redis/Aspire.Hosting.Redis.csproj

@@ -16,6 +16,10 @@
     <Compile Include="$(SharedDir)VolumeNameGenerator.cs" Link="Utils\VolumeNameGenerator.cs" />
   </ItemGroup>
 
+  <ItemGroup>
+    <PackageReference Include="AspNetCore.HealthChecks.Redis" />
+  </ItemGroup>
+
   <ItemGroup>
     <ProjectReference Include="..\Aspire.Hosting\Aspire.Hosting.csproj" />
   </ItemGroup>

src/Aspire.Hosting.Redis/RedisBuilderExtensions.cs

@@ -6,6 +6,7 @@
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Redis;
 using Aspire.Hosting.Utils;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace Aspire.Hosting;
 
@@ -29,10 +30,25 @@ public static IResourceBuilder<RedisResource> AddRedis(this IDistributedApplicat
         ArgumentNullException.ThrowIfNull(builder);
 
         var redis = new RedisResource(name);
+
+        builder.Eventing.Subscribe<AfterEndpointsAllocatedEvent>(async (@event, ct) =>
+        {
+            var connectionString = await redis.GetConnectionStringAsync(ct).ConfigureAwait(false);
+            builder.Configuration["hackedinconnectionstring"] = connectionString;
+        });
+
+        var hcb = builder.Services.AddHealthChecks().AddRedis(sp =>
+            {
+                return builder.Configuration["hackedinconnectionstring"]!;
+            },
+            name: $"{redis.Name}_check"
+        );
+
         return builder.AddResource(redis)
                       .WithEndpoint(port: port, targetPort: 6379, name: RedisResource.PrimaryEndpointName)
                       .WithImage(RedisContainerImageTags.Image, RedisContainerImageTags.Tag)
-                      .WithImageRegistry(RedisContainerImageTags.Registry);
+                      .WithImageRegistry(RedisContainerImageTags.Registry)
+                      .WithAnnotation(new HealthCheckAnnotation($"{redis.Name}_check"));
     }
 
     /// <summary>

src/Aspire.Hosting/ApplicationModel/CustomResourceSnapshot.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Immutable;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
 
 namespace Aspire.Hosting.ApplicationModel;
 
@@ -35,6 +36,11 @@ public sealed record CustomResourceSnapshot
     /// </summary>
     public int? ExitCode { get; init; }
 
+    /// <summary>
+    /// The health status of the resource.
+    /// </summary>
+    public HealthStatus HealthStatus { get; init; } = HealthStatus.Unhealthy;
+
     /// <summary>
     /// The environment variables that should show up in the dashboard for this resource.
     /// </summary>

src/Aspire.Hosting/ApplicationModel/HealthCheckAnnotation.cs

@@ -0,0 +1,16 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Aspire.Hosting.ApplicationModel;
+
+/// <summary>
+/// An annotation which tracks the name of the health check used to detect to health of a resource.
+/// </summary>
+/// <param name="key">TODO     </param>
+public class HealthCheckAnnotation(string key) : IResourceAnnotation
+{
+    /// <summary>
+    /// TODO
+    /// </summary>
+    public string Key => key;
+}

src/Aspire.Hosting/DistributedApplicationBuilder.cs

@@ -7,11 +7,13 @@
 using Aspire.Hosting.Dashboard;
 using Aspire.Hosting.Dcp;
 using Aspire.Hosting.Eventing;
+using Aspire.Hosting.Health;
 using Aspire.Hosting.Lifecycle;
 using Aspire.Hosting.Publishing;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.DependencyInjection.Extensions;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
@@ -159,8 +161,10 @@ public DistributedApplicationBuilder(DistributedApplicationOptions options)
         _innerBuilder.Services.AddHostedService<DistributedApplicationRunner>();
         _innerBuilder.Services.AddSingleton(options);
         _innerBuilder.Services.AddSingleton<ResourceNotificationService>();
+        _innerBuilder.Services.AddSingleton<IHealthCheckPublisher, ResourceNotificationHealthCheckPublisher>();
         _innerBuilder.Services.AddSingleton<ResourceLoggerService>();
         _innerBuilder.Services.AddSingleton<IDistributedApplicationEventing>(Eventing);
+        _innerBuilder.Services.AddHealthChecks();
 
         if (ExecutionContext.IsRunMode)
         {

src/Aspire.Hosting/Health/ResourceNotificationHealthCheckPublisher.cs

@@ -0,0 +1,37 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Aspire.Hosting.ApplicationModel;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+
+namespace Aspire.Hosting.Health;
+
+internal class ResourceNotificationHealthCheckPublisher(DistributedApplicationModel model, ResourceNotificationService resourceNotificationService) : IHealthCheckPublisher
+{
+    public async Task PublishAsync(HealthReport report, CancellationToken cancellationToken)
+    {
+        _ = model;
+        _ = resourceNotificationService;
+
+        foreach (var resource in model.Resources)
+        {
+            if (resource.TryGetAnnotationsOfType<HealthCheckAnnotation>(out var annotations))
+            {
+                var resourceEntries = report.Entries.Where(e => annotations.Any(a => a.Key == e.Key));
+                var status = resourceEntries.All(e => e.Value.Status == HealthStatus.Healthy) ? HealthStatus.Healthy : HealthStatus.Unhealthy;
+
+                await resourceNotificationService.PublishUpdateAsync(resource, s => s with
+                {
+                    HealthStatus = status
+                }).ConfigureAwait(false);
+            }
+            else
+            {
+                await resourceNotificationService.PublishUpdateAsync(resource, s => s with
+                {
+                    HealthStatus = HealthStatus.Healthy
+                }).ConfigureAwait(false);
+            }
+        }
+    }
+}

src/Aspire.Hosting/PublicAPI.Unshipped.txt

@@ -22,6 +22,11 @@ Aspire.Hosting.ApplicationModel.ContainerLifetimeAnnotation.LifetimeType.set ->
 Aspire.Hosting.ApplicationModel.ContainerLifetimeType
 Aspire.Hosting.ApplicationModel.ContainerLifetimeType.Default = 0 -> Aspire.Hosting.ApplicationModel.ContainerLifetimeType
 Aspire.Hosting.ApplicationModel.ContainerLifetimeType.Persistent = 1 -> Aspire.Hosting.ApplicationModel.ContainerLifetimeType
+Aspire.Hosting.ApplicationModel.CustomResourceSnapshot.HealthStatus.get -> Microsoft.Extensions.Diagnostics.HealthChecks.HealthStatus
+Aspire.Hosting.ApplicationModel.CustomResourceSnapshot.HealthStatus.init -> void
+Aspire.Hosting.ApplicationModel.HealthCheckAnnotation
+Aspire.Hosting.ApplicationModel.HealthCheckAnnotation.HealthCheckAnnotation(string! key) -> void
+Aspire.Hosting.ApplicationModel.HealthCheckAnnotation.Key.get -> string!
 Aspire.Hosting.ApplicationModel.ResourceNotificationService.ResourceNotificationService(Microsoft.Extensions.Logging.ILogger<Aspire.Hosting.ApplicationModel.ResourceNotificationService!>! logger, Microsoft.Extensions.Hosting.IHostApplicationLifetime! hostApplicationLifetime) -> void
 Aspire.Hosting.ApplicationModel.ResourceNotificationService.WaitForResourceAsync(string! resourceName, System.Func<Aspire.Hosting.ApplicationModel.ResourceEvent!, bool>! predicate, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task<Aspire.Hosting.ApplicationModel.ResourceEvent!>!
 Aspire.Hosting.DistributedApplicationBuilder.Eventing.get -> Aspire.Hosting.Eventing.IDistributedApplicationEventing!

src/Aspire.Hosting/ResourceBuilderExtensions.cs

@@ -5,6 +5,7 @@
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Publishing;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
 using Microsoft.Extensions.Logging;
 
 namespace Aspire.Hosting;
@@ -607,6 +608,9 @@ public static IResourceBuilder<T> WaitFor<T>(this IResourceBuilder<T> builder, I
                     $"Resource '{dependency.Resource.Name}' has entered the '{snapshot.State.Text}' state prematurely."
                     );
             }
+
+            // If we get to here the resource is running so we just need to check on whether it is healthy or not.
+            await rns.WaitForResourceAsync(dependency.Resource.Name, re => re.Snapshot.HealthStatus == HealthStatus.Healthy, cancellationToken: ct).ConfigureAwait(false);
         });
 
         return builder;