From 77ec78938f138c3e779aa7dbe4aba4e77971e20c Mon Sep 17 00:00:00 2001 From: Alireza Baloochi Date: Tue, 17 Sep 2024 21:03:11 +0330 Subject: [PATCH 1/2] WaitFor for Nats --- .../MilvusBuilderExtensions.cs | 9 ++-- .../Aspire.Hosting.Nats.csproj | 5 +++ .../NatsBuilderExtensions.cs | 35 +++++++++++++++- .../NatsFunctionalTests.cs | 42 +++++++++++++++++++ 4 files changed, 84 insertions(+), 7 deletions(-) diff --git a/src/Aspire.Hosting.Milvus/MilvusBuilderExtensions.cs b/src/Aspire.Hosting.Milvus/MilvusBuilderExtensions.cs index 098bc619692..9eb1260c90a 100644 --- a/src/Aspire.Hosting.Milvus/MilvusBuilderExtensions.cs +++ b/src/Aspire.Hosting.Milvus/MilvusBuilderExtensions.cs @@ -57,16 +57,13 @@ public static IResourceBuilder AddMilvus(this IDistributed var milvus = new MilvusServerResource(name, apiKeyParameter); - string? connectionString = null; MilvusClient? milvusClient = null; builder.Eventing.Subscribe(milvus, async (@event, ct) => { - connectionString = await milvus.ConnectionStringExpression.GetValueAsync(ct).ConfigureAwait(false); - if (connectionString is null) - { - throw new DistributedApplicationException($"ConnectionStringAvailableEvent was published for the '{milvus.Name}' resource but the connection string was null."); - } + var connectionString = await milvus.ConnectionStringExpression.GetValueAsync(ct).ConfigureAwait(false) + ?? throw new DistributedApplicationException($"ConnectionStringAvailableEvent was published for the '{milvus.Name}' resource but the connection string was null."); + milvusClient = CreateMilvusClient(@event.Services, connectionString); var lookup = builder.Resources.OfType().ToDictionary(d => d.Name); foreach (var databaseName in milvus.Databases) diff --git a/src/Aspire.Hosting.Nats/Aspire.Hosting.Nats.csproj b/src/Aspire.Hosting.Nats/Aspire.Hosting.Nats.csproj index 19928f9302b..a2c0ab96400 100644 --- a/src/Aspire.Hosting.Nats/Aspire.Hosting.Nats.csproj +++ b/src/Aspire.Hosting.Nats/Aspire.Hosting.Nats.csproj @@ -14,10 +14,15 @@ + + + + + diff --git a/src/Aspire.Hosting.Nats/NatsBuilderExtensions.cs b/src/Aspire.Hosting.Nats/NatsBuilderExtensions.cs index 1929ef06ddd..b61df531f5e 100644 --- a/src/Aspire.Hosting.Nats/NatsBuilderExtensions.cs +++ b/src/Aspire.Hosting.Nats/NatsBuilderExtensions.cs @@ -4,6 +4,11 @@ using Aspire.Hosting.ApplicationModel; using Aspire.Hosting.Nats; using Aspire.Hosting.Utils; +using Aspire.NATS.Net; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Diagnostics.HealthChecks; +using Microsoft.Extensions.Logging; +using NATS.Client.Core; namespace Aspire.Hosting; @@ -25,10 +30,38 @@ public static IResourceBuilder AddNats(this IDistributedAppl ArgumentNullException.ThrowIfNull(name); var nats = new NatsServerResource(name); + + NatsConnection? natsConnection = null; + + builder.Eventing.Subscribe(nats, async (@event, ct) => + { + var connectionString = await nats.ConnectionStringExpression.GetValueAsync(ct).ConfigureAwait(false) + ?? throw new DistributedApplicationException($"ConnectionStringAvailableEvent was published for the '{nats.Name}' resource but the connection string was null."); + + var options = NatsOpts.Default with + { + LoggerFactory = @event.Services.GetRequiredService(), + }; + + options = options with { Url = connectionString }; + + natsConnection = new NatsConnection(options); + }); + + var healthCheckKey = $"{name}_check"; + builder.Services.AddHealthChecks() + .Add(new HealthCheckRegistration( + healthCheckKey, + sp => new NatsHealthCheck(natsConnection!), + failureStatus: default, + tags: default, + timeout: default)); + return builder.AddResource(nats) .WithEndpoint(targetPort: 4222, port: port, name: NatsServerResource.PrimaryEndpointName) .WithImage(NatsContainerImageTags.Image, NatsContainerImageTags.Tag) - .WithImageRegistry(NatsContainerImageTags.Registry); + .WithImageRegistry(NatsContainerImageTags.Registry) + .WithHealthCheck(healthCheckKey); } /// diff --git a/tests/Aspire.Hosting.Nats.Tests/NatsFunctionalTests.cs b/tests/Aspire.Hosting.Nats.Tests/NatsFunctionalTests.cs index e107f538e52..7f0cacace22 100644 --- a/tests/Aspire.Hosting.Nats.Tests/NatsFunctionalTests.cs +++ b/tests/Aspire.Hosting.Nats.Tests/NatsFunctionalTests.cs @@ -11,6 +11,8 @@ using NATS.Client.Core; using NATS.Client.JetStream; using NATS.Client.JetStream.Models; +using Aspire.Hosting.ApplicationModel; +using Microsoft.Extensions.Diagnostics.HealthChecks; namespace Aspire.Hosting.Nats.Tests; public class NatsFunctionalTests(ITestOutputHelper testOutputHelper) @@ -236,4 +238,44 @@ private static async Task CreateTestData(INatsJSContext jetStream, CancellationT ack.EnsureSuccess(); } } + + [Fact] + [RequiresDocker] + public async Task VerifyWaitForOnNatsBlocksDependentResources() + { + var cts = new CancellationTokenSource(TimeSpan.FromMinutes(3)); + using var builder = TestDistributedApplicationBuilder.CreateWithTestContainerRegistry(testOutputHelper); + + var healthCheckTcs = new TaskCompletionSource(); + builder.Services.AddHealthChecks().AddAsyncCheck("blocking_check", () => + { + return healthCheckTcs.Task; + }); + + var resource = builder.AddNats("resource") + .WithHealthCheck("blocking_check"); + + var dependentResource = builder.AddNats("dependentresource") + .WaitFor(resource); + + using var app = builder.Build(); + + var pendingStart = app.StartAsync(cts.Token); + + var rns = app.Services.GetRequiredService(); + + await rns.WaitForResourceAsync(resource.Resource.Name, KnownResourceStates.Running, cts.Token); + + await rns.WaitForResourceAsync(dependentResource.Resource.Name, KnownResourceStates.Waiting, cts.Token); + + healthCheckTcs.SetResult(HealthCheckResult.Healthy()); + + await rns.WaitForResourceAsync(resource.Resource.Name, (re => re.Snapshot.HealthStatus == HealthStatus.Healthy), cts.Token); + + await rns.WaitForResourceAsync(dependentResource.Resource.Name, KnownResourceStates.Running, cts.Token); + + await pendingStart; + + await app.StopAsync(); + } } From 0f36635e32702aa4db004dd1ce3b923f26e26ce7 Mon Sep 17 00:00:00 2001 From: Alireza Baloochi Date: Tue, 17 Sep 2024 21:11:32 +0330 Subject: [PATCH 2/2] Update plaground app --- playground/nats/Nats.AppHost/Program.cs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/playground/nats/Nats.AppHost/Program.cs b/playground/nats/Nats.AppHost/Program.cs index ecd1da87d21..fbe6d421472 100644 --- a/playground/nats/Nats.AppHost/Program.cs +++ b/playground/nats/Nats.AppHost/Program.cs @@ -5,9 +5,11 @@ builder.AddProject("api") .WithExternalHttpEndpoints() - .WithReference(nats); + .WithReference(nats) + .WaitFor(nats); builder.AddProject("backend") - .WithReference(nats); + .WithReference(nats) + .WaitFor(nats); builder.Build().Run();