Alirexaa/qdrant waitfor (#6044)

* WaitFor for Qdrant * Address PR feedback * Address PR feedback * Don't use logger factory. --------- Co-authored-by: Alireza Baloochi <[email protected]>
dotnet · Oct 1, 2024 · f4e3bc3 · f4e3bc3
1 parent 865775b
commit f4e3bc3
Show file tree

Hide file tree

Showing 6 changed files with 144 additions and 2 deletions.
diff --git a/playground/Qdrant/Qdrant.AppHost/Program.cs b/playground/Qdrant/Qdrant.AppHost/Program.cs
@@ -8,6 +8,7 @@
 
 builder.AddProject<Projects.Qdrant_ApiService>("apiservice")
     .WithExternalHttpEndpoints()
-    .WithReference(qdrant);
+    .WithReference(qdrant)
+    .WaitFor(qdrant);
 
 builder.Build().Run();
diff --git a/src/Aspire.Hosting.Qdrant/Aspire.Hosting.Qdrant.csproj b/src/Aspire.Hosting.Qdrant/Aspire.Hosting.Qdrant.csproj
@@ -15,6 +15,7 @@
   <ItemGroup>
     <Compile Include="$(SharedDir)StringComparers.cs" Link="Utils\StringComparers.cs" />
     <Compile Include="$(SharedDir)VolumeNameGenerator.cs" Link="Utils\VolumeNameGenerator.cs" />
+    <Compile Include="$(ComponentsDir)Aspire.Qdrant.Client\QdrantHealthCheck.cs" Link="QdrantHealthCheck.cs"></Compile>
   </ItemGroup>
 
   <ItemGroup>

diff --git a/src/Aspire.Hosting.Qdrant/QdrantBuilderExtensions.cs b/src/Aspire.Hosting.Qdrant/QdrantBuilderExtensions.cs
@@ -1,9 +1,13 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Data.Common;
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Qdrant;
 using Aspire.Hosting.Utils;
+using Aspire.Qdrant.Client;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
 
 namespace Aspire.Hosting;
 
@@ -42,6 +46,25 @@ public static IResourceBuilder<QdrantServerResource> AddQdrant(this IDistributed
         var apiKeyParameter = apiKey?.Resource ??
             ParameterResourceBuilderExtensions.CreateDefaultPasswordParameter(builder, $"{name}-Key", special: false);
         var qdrant = new QdrantServerResource(name, apiKeyParameter);
+
+        HttpClient? httpClient = null;
+
+        builder.Eventing.Subscribe<ConnectionStringAvailableEvent>(qdrant, async (@event, ct) =>
+        {
+            var connectionString = await qdrant.HttpConnectionStringExpression.GetValueAsync(ct).ConfigureAwait(false)
+            ?? throw new DistributedApplicationException($"ConnectionStringAvailableEvent was published for the '{qdrant.Name}' resource but the connection string was null.");
+            httpClient = CreateQdrantHttpClient(connectionString);
+        });
+
+        var healthCheckKey = $"{name}_check";
+        builder.Services.AddHealthChecks()
+          .Add(new HealthCheckRegistration(
+              healthCheckKey,
+              sp => new QdrantHealthCheck(httpClient!),
+              failureStatus: default,
+              tags: default,
+              timeout: default));
+
         return builder.AddResource(qdrant)
             .WithImage(QdrantContainerImageTags.Image, QdrantContainerImageTags.Tag)
             .WithImageRegistry(QdrantContainerImageTags.Registry)
@@ -61,7 +84,8 @@ public static IResourceBuilder<QdrantServerResource> AddQdrant(this IDistributed
                 {
                     context.EnvironmentVariables[EnableStaticContentEnvVarName] = "0";
                 }
-            });
+            })
+            .WithHealthCheck(healthCheckKey);
     }
 
     /// <summary>
@@ -117,4 +141,45 @@ public static IResourceBuilder<TDestination> WithReference<TDestination>(this IR
 
         return builder;
     }
+
+    private static HttpClient CreateQdrantHttpClient(string? connectionString)
+    {
+        if (connectionString is null)
+        {
+            throw new InvalidOperationException("Connection string is unavailable");
+        }
+
+        Uri? endpoint = null;
+        string? key = null;
+
+        if (Uri.TryCreate(connectionString, UriKind.Absolute, out var uri))
+        {
+            endpoint = uri;
+        }
+        else
+        {
+            var connectionBuilder = new DbConnectionStringBuilder
+            {
+                ConnectionString = connectionString
+            };
+
+            if (connectionBuilder.TryGetValue("Endpoint", out var endpointValue) && Uri.TryCreate(endpointValue.ToString(), UriKind.Absolute, out var serviceUri))
+            {
+                endpoint = serviceUri;
+            }
+
+            if (connectionBuilder.TryGetValue("Key", out var keyValue))
+            {
+                key = keyValue.ToString();
+            }
+        }
+
+        var client = new HttpClient();
+        client.BaseAddress = endpoint;
+        if (key is not null)
+        {
+            client.DefaultRequestHeaders.Add("Api-Key", key);
+        }
+        return client;
+    }
 }
diff --git a/src/Components/Aspire.Qdrant.Client/Aspire.Qdrant.Client.csproj b/src/Components/Aspire.Qdrant.Client/Aspire.Qdrant.Client.csproj
@@ -20,6 +20,7 @@
     <PackageReference Include="Qdrant.Client" />
     <PackageReference Include="Microsoft.Extensions.Configuration.Binder" />
     <PackageReference Include="Microsoft.Extensions.Hosting.Abstractions" />
+    <PackageReference Include="Microsoft.Extensions.Diagnostics.HealthChecks" />
   </ItemGroup>
 
 </Project>
diff --git a/src/Components/Aspire.Qdrant.Client/QdrantHealthCheck.cs b/src/Components/Aspire.Qdrant.Client/QdrantHealthCheck.cs
@@ -0,0 +1,32 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+
+namespace Aspire.Qdrant.Client;
+internal sealed class QdrantHealthCheck : IHealthCheck
+{
+    private readonly HttpClient _client;
+
+    public QdrantHealthCheck(HttpClient client)
+    {
+        ArgumentNullException.ThrowIfNull(client, nameof(client));
+        _client = client;
+    }
+
+    public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
+    {
+        try
+        {
+            var response = await _client.GetAsync("/readyz", cancellationToken).ConfigureAwait(false);
+
+            return response.IsSuccessStatusCode
+                ? HealthCheckResult.Healthy()
+                : new HealthCheckResult(HealthStatus.Unhealthy);
+        }
+        catch (Exception ex)
+        {
+            return new HealthCheckResult(context.Registration.FailureStatus, exception: ex);
+        }
+    }
+}
diff --git a/tests/Aspire.Hosting.Qdrant.Tests/QdrantFunctionalTests.cs b/tests/Aspire.Hosting.Qdrant.Tests/QdrantFunctionalTests.cs
@@ -2,10 +2,12 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using Aspire.Components.Common.Tests;
+using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Utils;
 using Grpc.Core;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
 using Microsoft.Extensions.Hosting;
 using Polly;
 using Qdrant.Client;
@@ -216,4 +218,44 @@ await pipeline.ExecuteAsync(async token =>
             }
         }
     }
+
+    [Fact]
+    [RequiresDocker]
+    public async Task VerifyWaitForOnQdrantBlocksDependentResources()
+    {
+        var cts = new CancellationTokenSource(TimeSpan.FromMinutes(3));
+        using var builder = TestDistributedApplicationBuilder.CreateWithTestContainerRegistry(testOutputHelper);
+
+        var healthCheckTcs = new TaskCompletionSource<HealthCheckResult>();
+        builder.Services.AddHealthChecks().AddAsyncCheck("blocking_check", () =>
+        {
+            return healthCheckTcs.Task;
+        });
+
+        var resource = builder.AddQdrant("resource")
+                              .WithHealthCheck("blocking_check");
+
+        var dependentResource = builder.AddQdrant("dependentresource")
+                                       .WaitFor(resource);
+
+        using var app = builder.Build();
+
+        var pendingStart = app.StartAsync(cts.Token);
+
+        var rns = app.Services.GetRequiredService<ResourceNotificationService>();
+
+        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();
+    }
 }