Improved Qdrant server startup & client initialization (#770)

app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.cs

@@ -0,0 +1,180 @@
+using AIStudio.Tools.Databases.Qdrant;
+using AIStudio.Tools.Rust;
+using AIStudio.Tools.Services;
+
+namespace AIStudio.Tools.Databases;
+
+public sealed class DatabaseClientProvider(RustService rustService, ILoggerFactory loggerFactory) : IDisposable
+{
+    private readonly Dictionary<DatabaseRole, DatabaseClient> clients = new();
+    private readonly Dictionary<DatabaseRole, SemaphoreSlim> locks = new();
+    private readonly Lock locksLock = new();
+    private readonly ILogger<DatabaseClientProvider> logger = loggerFactory.CreateLogger<DatabaseClientProvider>();
+    private readonly ILogger<DatabaseClient> databaseClientLogger = loggerFactory.CreateLogger<DatabaseClient>();
+
+    public async Task<DatabaseClient> GetClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken = default)
+    {
+        var databaseLock = this.GetLock(databaseRole);
+        await databaseLock.WaitAsync(cancellationToken);
+        try
+        {
+            if (this.clients.TryGetValue(databaseRole, out var cachedClient) && cachedClient.IsAvailable)
+                return cachedClient;
+
+            var client = await this.CreateClientAsync(databaseRole, cancellationToken);
+            return this.CacheIfAvailable(databaseRole, client);
+        }
+        finally
+        {
+            databaseLock.Release();
+        }
+    }
+
+    public async Task<DatabaseClient> RefreshClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken = default)
+    {
+        var databaseLock = this.GetLock(databaseRole);
+        await databaseLock.WaitAsync(cancellationToken);
+        try
+        {
+            var client = await this.CreateClientAsync(databaseRole, cancellationToken);
+            return this.CacheIfAvailable(databaseRole, client);
+        }
+        finally
+        {
+            databaseLock.Release();
+        }
+    }
+
+    private DatabaseClient CacheIfAvailable(DatabaseRole databaseRole, DatabaseClient client)
+    {
+        if (!client.IsAvailable)
+            return client;
+
+        if (this.clients.TryGetValue(databaseRole, out var cachedClient))
+        {
+            if (IsSameClient(cachedClient, client))
+            {
+                client.Dispose();
+                return cachedClient;
+            }
+
+            cachedClient.Dispose();
+        }
+
+        this.clients[databaseRole] = client;
+        return client;
+    }
+
+    private SemaphoreSlim GetLock(DatabaseRole databaseRole)
+    {
+        lock (this.locksLock)
+        {
+            if (this.locks.TryGetValue(databaseRole, out var databaseLock))
+                return databaseLock;
+
+            databaseLock = new SemaphoreSlim(1, 1);
+            this.locks[databaseRole] = databaseLock;
+            return databaseLock;
+        }
+    }
+
+    private async Task<DatabaseClient> CreateClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken) => databaseRole switch
+    {
+        DatabaseRole.VECTOR_STORE => await this.CreateQdrantClientAsync(cancellationToken),
+        _ => new NoDatabaseClient(databaseRole.ToString(), "The requested database role is not supported.")
+    };
+
+    private async Task<DatabaseClient> CreateQdrantClientAsync(CancellationToken cancellationToken)
+    {
+        var qdrantInfo = await rustService.GetQdrantInfo(cancellationToken);
+        if (qdrantInfo.Status is QdrantStatus.STARTING)
+        {
+            return this.CreateNoDatabaseClient(
+                "Qdrant",
+                "Qdrant is starting. Details will appear shortly.",
+                DatabaseClientStatus.STARTING);
+        }
+
+        if (!qdrantInfo.IsAvailable || qdrantInfo.Status is QdrantStatus.UNAVAILABLE)
+        {
+            var reason = qdrantInfo.UnavailableReason ?? "unknown";
+            this.logger.LogWarning("Qdrant is not available. Starting without vector database. Reason: '{Reason}'.", reason);
+            return this.CreateNoDatabaseClient("Qdrant", qdrantInfo.UnavailableReason, DatabaseClientStatus.UNAVAILABLE);
+        }
+
+        if (!HasValidQdrantConnectionInfo(qdrantInfo, out var invalidReason))
+            return this.CreateNoDatabaseClient("Qdrant", invalidReason, DatabaseClientStatus.UNAVAILABLE);
+
+        var client = new QdrantClientImplementation("Qdrant", qdrantInfo.Path, qdrantInfo.PortHttp, qdrantInfo.PortGrpc, qdrantInfo.Fingerprint, qdrantInfo.ApiToken);
+        client.SetLogger(this.databaseClientLogger);
+
+        try
+        {
+            await client.CheckAvailabilityAsync();
+            return client;
+        }
+        catch (Exception e)
+        {
+            client.Dispose();
+            this.logger.LogWarning(e, "Qdrant reported as available by Rust, but the health check failed.");
+            return this.CreateNoDatabaseClient("Qdrant", e.Message, DatabaseClientStatus.STARTING);
+        }
+    }
+
+    private static bool HasValidQdrantConnectionInfo(QdrantInfo qdrantInfo, out string invalidReason)
+    {
+        if (qdrantInfo.Path == string.Empty)
+        {
+            invalidReason = "Failed to get the Qdrant path from Rust.";
+            return false;
+        }
+
+        if (qdrantInfo.PortHttp == 0)
+        {
+            invalidReason = "Failed to get the Qdrant HTTP port from Rust.";
+            return false;
+        }
+
+        if (qdrantInfo.PortGrpc == 0)
+        {
+            invalidReason = "Failed to get the Qdrant gRPC port from Rust.";
+            return false;
+        }
+
+        if (qdrantInfo.Fingerprint == string.Empty)
+        {
+            invalidReason = "Failed to get the Qdrant fingerprint from Rust.";
+            return false;
+        }
+
+        if (qdrantInfo.ApiToken == string.Empty)
+        {
+            invalidReason = "Failed to get the Qdrant API token from Rust.";
+            return false;
+        }
+
+        invalidReason = string.Empty;
+        return true;
+    }
+
+    private NoDatabaseClient CreateNoDatabaseClient(string name, string? unavailableReason, DatabaseClientStatus status)
+    {
+        var client = new NoDatabaseClient(name, unavailableReason, status);
+        client.SetLogger(this.databaseClientLogger);
+        return client;
+    }
+
+    private static bool IsSameClient(DatabaseClient left, DatabaseClient right) =>
+        left.IsAvailable
+        && right.IsAvailable
+        && left.CacheKey == right.CacheKey;
+
+    public void Dispose()
+    {
+        foreach (var client in this.clients.Values)
+            client.Dispose();
+
+        foreach (var databaseLock in this.locks.Values)
+            databaseLock.Dispose();
+    }
+}