Add an extra init wait for single node clusters

Author: brantburnett

src/Couchbase.Aspire.Hosting/Api/CouchbaseApi.cs

@@ -129,12 +129,27 @@ public async Task AddNodeAsync(CouchbaseServerResource server, string hostname,
         };
 
         var response = await SendRequestAsync(server.GetManagementEndpoint(),
-                HttpMethod.Post,
-                "/controller/addNode",
-                new FormUrlEncodedContent(parameters),
-                cancellationToken: cancellationToken).ConfigureAwait(false);
+            HttpMethod.Post,
+            "/controller/addNode",
+            new FormUrlEncodedContent(parameters),
+            cancellationToken: cancellationToken).ConfigureAwait(false);
+
+        await ThrowOnFailureAsync(response, cancellationToken).ConfigureAwait(false);
+    }
+
+    public async Task<NodeServices> GetNodeServicesAsync(CouchbaseServerResource server, CancellationToken cancellationToken = default)
+    {
+        ArgumentNullException.ThrowIfNull(server);
+
+        var response = await SendRequestAsync(server.GetManagementEndpoint(),
+            HttpMethod.Get,
+            "pools/default/nodeServices",
+            cancellationToken: cancellationToken).ConfigureAwait(false);
 
         await ThrowOnFailureAsync(response, cancellationToken).ConfigureAwait(false);
+
+        var nodeServicesResponse = await response.Content.ReadFromJsonAsync<NodeServicesResponse>(cancellationToken).ConfigureAwait(false);
+        return nodeServicesResponse!.NodesExt.First(p => p.ThisNode);
     }
 
     public async Task SetupAlternateAddressesAsync(CouchbaseServerResource server, string hostname, Dictionary<string, string> ports,

src/Couchbase.Aspire.Hosting/Api/ICouchbaseApi.cs

@@ -9,6 +9,7 @@ internal interface ICouchbaseApi
     Task<Bucket?> GetBucketAsync(CouchbaseServerResource server, string bucketName, CancellationToken cancellationToken);
     Task FlushBucketAsync(CouchbaseServerResource server, string bucketName, CancellationToken cancellationToken);
     Task<Pool> GetClusterNodesAsync(CouchbaseServerResource server, CancellationToken cancellationToken = default);
+    Task<NodeServices> GetNodeServicesAsync(CouchbaseServerResource server, CancellationToken cancellationToken = default);
     Task<bool> GetDefaultPoolAsync(CouchbaseServerResource server, bool preferInsecure = false, CancellationToken cancellationToken = default);
     Task<RebalanceStatus> GetRebalanceProgressAsync(CouchbaseServerResource server, CancellationToken cancellationToken = default);
     Task InitializeClusterAsync(CouchbaseServerResource server, CouchbaseClusterSettings settings, CancellationToken cancellationToken = default);

src/Couchbase.Aspire.Hosting/Api/Pool.cs

@@ -13,3 +13,21 @@ internal sealed class Node
     [JsonPropertyName("hostname")]
     public string Hostname { get; set; } = null!;
 }
+
+internal sealed class NodeServices
+{
+    [JsonPropertyName("hostname")]
+    public string Hostname { get; set; } = null!;
+
+    [JsonPropertyName("services")]
+    public Dictionary<string, int>? Services { get; set; }
+
+    [JsonPropertyName("thisNode")]
+    public bool ThisNode { get; set; }
+}
+
+internal sealed class NodeServicesResponse
+{
+    [JsonPropertyName("nodesExt")]
+    public List<NodeServices> NodesExt { get; set; } = null!;
+}

src/Couchbase.Aspire.Hosting/Orchestration/CouchbaseClusterOrchestrator.cs

@@ -382,6 +382,14 @@ await _resourceNotificationService.PublishUpdateAsync(cluster, s => s with
                         await RebalanceAsync(api, primaryServer, cancellationToken).ConfigureAwait(false);
                     }
                 }
+                else
+                {
+                    // A single node cluster may not have all services come online immediately, which can result in
+                    // health checks that never pass when the SDK bootstraps with a subset of services. Wait for
+                    // all services to appear in the ports list
+
+                    await WaitForServicesAsync(api, primaryServer, cancellationToken).ConfigureAwait(false);
+                }
 
                 await _orchestratorEvents.PublishAsync(new OnCouchbaseResourceStartedEvent(cluster), cancellationToken).ConfigureAwait(false);
 
@@ -453,6 +461,44 @@ public async Task InitializeClusterAsync(ICouchbaseApi api, CouchbaseServerResou
         await api.InitializeClusterAsync(primaryServer, settings, cancellationToken).ConfigureAwait(false);
     }
 
+    public async Task WaitForServicesAsync(ICouchbaseApi api, CouchbaseServerResource server, CancellationToken cancellationToken = default)
+    {
+        while (true)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            if (!server.TryGetEndpoints(out var endpoints))
+            {
+                throw new InvalidOperationException("Failed to get node endpoints.");
+            }
+
+            var node = await api.GetNodeServicesAsync(server, cancellationToken).ConfigureAwait(false);
+
+            if (node.Services is not null)
+            {
+                var allServicesFound = true;
+                foreach (var endpoint in endpoints)
+                {
+                    if (CouchbaseEndpointNames.EndpointNameServiceMappings.TryGetValue(endpoint.Name, out var serviceName))
+                    {
+                        if (!node.Services.ContainsKey(serviceName))
+                        {
+                            allServicesFound = false;
+                            break;
+                        }
+                    }
+                }
+
+                if (allServicesFound)
+                {
+                    return;
+                }
+            }
+
+            await Task.Delay(250, cancellationToken).ConfigureAwait(false);
+        }
+    }
+
     /// <returns><c>true</c> if the node was added, <c>false</c> if it is already part of the cluster.</returns>
     private async Task<bool> AddNodeAsync(ICouchbaseApi api, CouchbaseServerResource primaryServer, CouchbaseServerResource addServer,
         List<string> existingNodes, CancellationToken cancellationToken)