Merge pull request #8 from slang25/eager-load-services

Eagerly load services

Author: Blind-Striker

src/Aspire.Hosting.LocalStack/Container/LocalStackContainerOptions.cs

@@ -3,6 +3,7 @@
 
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.LocalStack.Container;
+using LocalStack.Client.Enums;
 
 namespace Aspire.Hosting.LocalStack;
 
@@ -17,7 +18,6 @@ public sealed class LocalStackContainerOptions
     /// <remarks>
     /// - <see cref="ContainerLifetime.Persistent"/>: Container survives application restarts (default for databases)
     /// - <see cref="ContainerLifetime.Session"/>: Container is cleaned up when application stops (recommended for LocalStack)
-    /// - <see cref="ContainerLifetime.Transient"/>: Container is recreated on each run
     /// </remarks>
     public ContainerLifetime Lifetime { get; set; } = ContainerLifetime.Persistent;
 
@@ -49,4 +49,7 @@ public sealed class LocalStackContainerOptions
     /// Gets or sets additional environment variables to pass to the LocalStack container.
     /// </summary>
     public IDictionary<string, string> AdditionalEnvironmentVariables { get; } = new Dictionary<string, string>(StringComparer.Ordinal);
+
+    /// <summary>A collection of services to eagerly start.</summary>
+    public IReadOnlyCollection<AwsService> EagerLoadedServices { get; set; } = [];
 }

src/Aspire.Hosting.LocalStack/Internal/LocalStackHealthCheck.cs

@@ -0,0 +1,69 @@
+using System.Net.Http.Json;
+using System.Text.Json.Nodes;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+
+namespace Aspire.Hosting.LocalStack.Internal;
+
+internal sealed class LocalStackHealthCheck(Uri uri, string[] services) : IHealthCheck, IDisposable
+{
+    private readonly HttpClient _client =
+        new(new SocketsHttpHandler { ActivityHeadersPropagator = null })
+        {
+            BaseAddress = uri, Timeout = TimeSpan.FromSeconds(1)
+        };
+
+    public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context,
+        CancellationToken cancellationToken = default)
+    {
+        try
+        {
+#pragma warning disable CA2234
+            using var response = await _client.GetAsync("_localstack/health", cancellationToken).ConfigureAwait(false);
+#pragma warning restore CA2234
+            if (response.IsSuccessStatusCode)
+            {
+                var responseJson =
+                    await response.Content.ReadFromJsonAsync<JsonNode>(cancellationToken: cancellationToken).ConfigureAwait(false);
+                var servicesNode = responseJson?["services"]?.AsObject();
+
+                if (servicesNode is null)
+                {
+                    return HealthCheckResult.Unhealthy(
+                        "LocalStack health response did not contain a 'services' object."
+                    );
+                }
+
+                var failingServices = services
+                    .Where(s =>
+                        !servicesNode.ContainsKey(s)
+                        || servicesNode[s]?.ToString() != "running"
+                    )
+                    .ToList();
+
+                if (failingServices.Count == 0)
+                {
+                    return HealthCheckResult.Healthy("LocalStack is healthy.");
+                }
+
+                var reason =
+                    $"The following required services are not running: {string.Join(", ", failingServices)}.";
+                return HealthCheckResult.Unhealthy(
+                    $"LocalStack is unhealthy. {reason}"
+                );
+            }
+
+            return HealthCheckResult.Unhealthy("LocalStack is unhealthy.");
+        }
+#pragma warning disable CA1031 // Do not catch general exception types
+        catch (Exception ex)
+#pragma warning restore CA1031 // Do not catch general exception types
+        {
+            return HealthCheckResult.Unhealthy("LocalStack is unhealthy.", ex);
+        }
+    }
+
+    public void Dispose()
+    {
+        _client.Dispose();
+    }
+}

src/Aspire.Hosting.LocalStack/LocalStackResourceBuilderExtensions.cs

@@ -13,8 +13,11 @@
 using Aspire.Hosting.LocalStack.Container;
 using Aspire.Hosting.LocalStack.Internal;
 using LocalStack.Client.Contracts;
+using LocalStack.Client.Enums;
 using LocalStack.Client.Options;
 using Microsoft.Extensions.Configuration;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
 
 namespace Aspire.Hosting;
 
@@ -168,7 +171,6 @@ public static IDistributedApplicationBuilder UseLocalStack(this IDistributedAppl
             .WithImageRegistry(LocalStackContainerImageTags.Registry)
             .WithImageTag(LocalStackContainerImageTags.Tag)
             .WithHttpEndpoint(targetPort: Constants.DefaultContainerPort, name: LocalStackResource.PrimaryEndpointName)
-            .WithHttpHealthCheck("/_localstack/health", 200, LocalStackResource.PrimaryEndpointName)
             .WithLifetime(containerOptions.Lifetime)
             .WithEnvironment("DEBUG", containerOptions.DebugLevel.ToString(CultureInfo.InvariantCulture))
             .WithEnvironment("LS_LOG", containerOptions.LogLevel.ToEnvironmentValue())
@@ -182,6 +184,31 @@ public static IDistributedApplicationBuilder UseLocalStack(this IDistributedAppl
             resourceBuilder = resourceBuilder.WithEnvironment(key, value);
         }
 
+        if (containerOptions.EagerLoadedServices.Count == 0)
+        {
+            resourceBuilder = resourceBuilder.WithHttpHealthCheck("/_localstack/health", 200, LocalStackResource.PrimaryEndpointName);
+        }
+        else
+        {
+            resourceBuilder = resourceBuilder.WithEnvironment("EAGER_SERVICE_LOADING", "1");
+
+            List<string> serviceNames = [];
+            foreach (var awsService in containerOptions.EagerLoadedServices)
+            {
+                var serviceName = AwsServiceEndpointMetadata.ByEnum(awsService)!.CliName;
+                if (serviceName is null)
+                {
+                    throw new InvalidOperationException($"Eager loaded service '{awsService}' is not supported by LocalStack.");
+                }
+                serviceNames.Add(serviceName);
+            }
+
+            var servicesValue = string.Join(',', serviceNames);
+            resourceBuilder = resourceBuilder
+                .WithEnvironment("SERVICES", servicesValue)
+                .WithLocalStackHealthCheck(serviceNames.ToArray());
+        }
+
         // Configure callback for dynamic resource configuration
         var callback = LocalStackConnectionStringAvailableCallback.CreateCallback(builder);
         resourceBuilder.OnConnectionStringAvailable(callback);
@@ -263,4 +290,54 @@ public static ILocalStackOptions AddLocalStackOptions(this IDistributedApplicati
 
         return options;
     }
+
+    private static IResourceBuilder<T> WithLocalStackHealthCheck<T>(this IResourceBuilder<T> builder, string[] services) where T : IResourceWithEndpoints
+    {
+        ArgumentNullException.ThrowIfNull(builder);
+
+        var endpoint = builder.Resource.GetEndpoint(LocalStackResource.PrimaryEndpointName);
+        if (endpoint.Scheme != "http")
+        {
+            throw new DistributedApplicationException($"Could not create HTTP health check for resource '{builder.Resource.Name}' as the endpoint with name '{endpoint.EndpointName}' and scheme '{endpoint.Scheme}' is not an HTTP endpoint.");
+        }
+
+        builder.EnsureEndpointIsAllocated(endpoint);
+
+        Uri? baseUri = null;
+        builder.ApplicationBuilder.Eventing.Subscribe<BeforeResourceStartedEvent>(builder.Resource, (@event, ct) =>
+        {
+            baseUri = new Uri(endpoint.Url, UriKind.Absolute);
+            return Task.CompletedTask;
+        });
+
+        var healthCheckKey = $"{builder.Resource.Name}_localstack_check";
+
+        builder.ApplicationBuilder.Services.AddHealthChecks().Add(new HealthCheckRegistration(healthCheckKey,
+            _ =>
+            {
+                return baseUri switch
+                {
+                    null => throw new DistributedApplicationException(
+                        "The URI for the health check is not set. Ensure that the resource has been allocated before the health check is executed."),
+                    _ => new LocalStackHealthCheck(baseUri!, services)
+                };
+            }, failureStatus: null, tags: null));
+
+        builder.WithHealthCheck(healthCheckKey);
+
+        return builder;
+    }
+
+    private static void EnsureEndpointIsAllocated<T>(this IResourceBuilder<T> builder, EndpointReference endpoint)  where T : IResourceWithEndpoints
+    {
+        var endpointName = endpoint.EndpointName;
+
+        builder.OnResourceEndpointsAllocated((_, _, _) =>
+            endpoint.Exists switch
+            {
+                true => Task.CompletedTask,
+                false => throw new DistributedApplicationException(
+                    $"The endpoint '{endpointName}' does not exist on the resource '{builder.Resource.Name}'.")
+            });
+    }
 }

tests/Aspire.Hosting.LocalStack.Integration.Tests/Aspire.Hosting.LocalStack.Integration.Tests.csproj

@@ -1,15 +1,19 @@
 <Project Sdk="Microsoft.NET.Sdk">
 
+  <Sdk Name="Aspire.AppHost.Sdk" Version="9.4.0" />
+
   <PropertyGroup>
     <TargetFramework>$(DefaultTargetFramework)</TargetFramework>
     <OutputType>Exe</OutputType>
     <ContinuousIntegrationBuild>false</ContinuousIntegrationBuild>
+    <IsAspireHost>false</IsAspireHost>
     <IsTestProject>true</IsTestProject>
     <NoWarn>$(NoWarn);CA2007;CA1515;CA5399</NoWarn>
     <NoError>$(NoWarn);CA2007;CA1515;CA5399</NoError>
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="Aspire.Hosting.AppHost" />
     <PackageReference Include="AWSSDK.Core"/>
     <PackageReference Include="AWSSDK.DynamoDBv2"/>
     <PackageReference Include="AWSSDK.SQS"/>

tests/Aspire.Hosting.LocalStack.Integration.Tests/EagerLoadedServices/EagerLoadedServicesTests.cs

@@ -0,0 +1,106 @@
+using System.Net.Http.Json;
+using System.Text.Json.Nodes;
+using Amazon;
+using Amazon.SQS.Model;
+using Aspire.Hosting.LocalStack.Container;
+using LocalStack.Client.Enums;
+
+namespace Aspire.Hosting.LocalStack.Integration.Tests.EagerLoadedServices;
+
+public class EagerLoadedServicesTests
+{
+    [Fact]
+    public async Task LocalStack_Should_Lazy_Load_Services_By_Default_Async()
+    {
+        using var parentCts = new CancellationTokenSource(TimeSpan.FromMinutes(5));
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(parentCts.Token, TestContext.Current.CancellationToken);
+
+#pragma warning disable CA1849
+        await using var builder = DistributedApplicationTestingBuilder.Create("LocalStack:UseLocalStack=true");
+#pragma warning restore CA1849
+
+        var awsConfig = builder.AddAWSSDKConfig().WithRegion(RegionEndpoint.EUCentral1);
+        builder.AddLocalStack(awsConfig: awsConfig, configureContainer: container =>
+        {
+            container.Lifetime = ContainerLifetime.Session;
+            container.DebugLevel = 1;
+            container.LogLevel = LocalStackLogLevel.Debug;
+        });
+
+        await using var app = await builder.BuildAsync(cts.Token);
+        await app.StartAsync(cts.Token);
+
+        var resourceNotificationService = app.Services.GetRequiredService<ResourceNotificationService>();
+
+        await resourceNotificationService.WaitForResourceHealthyAsync("localstack", cts.Token);
+
+        using var httpClient = app.CreateHttpClient("localstack", "http");
+        var healthResponse = await httpClient.GetAsync(new Uri("/_localstack/health", UriKind.Relative), cts.Token);
+        var healthContent = await healthResponse.Content.ReadFromJsonAsync<JsonNode>(cts.Token);
+        Assert.Equal(HttpStatusCode.OK, healthResponse.StatusCode);
+
+        var servicesNode = healthContent?["services"]?.AsObject();
+        Assert.NotNull(servicesNode);
+        Assert.True(servicesNode.ContainsKey("sqs"));
+        Assert.NotEqual("running", servicesNode["sqs"]?.ToString());
+
+        var connectionString = await app.GetConnectionStringAsync("localstack", cancellationToken: cts.Token);
+        Assert.NotNull(connectionString);
+        Assert.NotEmpty(connectionString);
+
+        var connectionStringUri = new Uri(connectionString);
+
+        var configOptions = new ConfigOptions(connectionStringUri.Host, edgePort: connectionStringUri.Port);
+        var sessionOptions = new SessionOptions(regionName: awsConfig.Region!.SystemName);
+        var session = SessionStandalone.Init().WithSessionOptions(sessionOptions).WithConfigurationOptions(configOptions).Create();
+
+        var sqsClient = session.CreateClientByImplementation<AmazonSQSClient>();
+        await sqsClient.ListQueuesAsync(new ListQueuesRequest(), cts.Token);
+
+        var laterHealthResponse = await httpClient.GetAsync(new Uri("/_localstack/health", UriKind.Relative), cts.Token);
+        var laterHealthContent = await laterHealthResponse.Content.ReadFromJsonAsync<JsonNode>(cts.Token);
+        Assert.Equal(HttpStatusCode.OK, laterHealthResponse.StatusCode);
+
+        var sqsServicesNode = laterHealthContent?["services"]?.AsObject();
+        Assert.NotNull(sqsServicesNode);
+        Assert.True(sqsServicesNode.ContainsKey("sqs"));
+        Assert.Equal("running", sqsServicesNode["sqs"]?.ToString());
+    }
+
+    [Fact]
+    public async Task LocalStack_Should_Eagerly_Load_Services_When_Configured_Async()
+    {
+        using var parentCts = new CancellationTokenSource(TimeSpan.FromMinutes(5));
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(parentCts.Token, TestContext.Current.CancellationToken);
+
+#pragma warning disable CA1849
+        await using var builder = DistributedApplicationTestingBuilder.Create("LocalStack:UseLocalStack=true");
+#pragma warning restore CA1849
+
+        var awsConfig = builder.AddAWSSDKConfig().WithRegion(RegionEndpoint.EUCentral1);
+        builder.AddLocalStack(awsConfig: awsConfig, configureContainer: container =>
+        {
+            container.Lifetime = ContainerLifetime.Session;
+            container.DebugLevel = 1;
+            container.LogLevel = LocalStackLogLevel.Debug;
+            container.EagerLoadedServices = [AwsService.Sqs];
+        });
+
+        await using var app = await builder.BuildAsync(cts.Token);
+        await app.StartAsync(cts.Token);
+
+        var resourceNotificationService = app.Services.GetRequiredService<ResourceNotificationService>();
+
+        await resourceNotificationService.WaitForResourceHealthyAsync("localstack", cts.Token);
+
+        using var httpClient = app.CreateHttpClient("localstack", "http");
+        var healthResponse = await httpClient.GetAsync(new Uri("/_localstack/health", UriKind.Relative), cts.Token);
+        var healthContent = await healthResponse.Content.ReadFromJsonAsync<JsonNode>(cts.Token);
+        Assert.Equal(HttpStatusCode.OK, healthResponse.StatusCode);
+
+        var servicesNode = healthContent?["services"]?.AsObject();
+        Assert.NotNull(servicesNode);
+        Assert.True(servicesNode.ContainsKey("sqs"));
+        Assert.Equal("running", servicesNode["sqs"]?.ToString());
+    }
+}