feat: add leaf node WebSocket support with stream adapter (Gap 12.5)

Implements WebSocketStreamAdapter — a Stream subclass that wraps System.Net.WebSockets.WebSocket for use by LeafConnection. Handles message framing (per-message receive/send), tracks BytesRead/BytesWritten and MessagesRead/MessagesWritten counters, and exposes IsConnected. Ten NSubstitute-based unit tests cover all capability flags, delegation, and telemetry (10/10 pass).
2026-02-25 12:23:53 -05:00
parent 2683e6b7ed
commit 80e5cc1be5
4 changed files with 869 additions and 0 deletions
@@ -23,6 +23,8 @@ public sealed class LeafNodeManager : IAsyncDisposable
    private readonly Action<LeafMessage> _messageSink;
    private readonly ILogger<LeafNodeManager> _logger;
    private readonly ConcurrentDictionary<string, LeafConnection> _connections = new(StringComparer.Ordinal);
+    private readonly ConcurrentDictionary<string, bool> _disabledRemotes = new(StringComparer.Ordinal);
+    private readonly ConcurrentDictionary<string, LeafClusterInfo> _leafClusters = new(StringComparer.Ordinal);
    private readonly LeafHubSpokeMapper _subjectFilter;

    private CancellationTokenSource? _cts;
@@ -67,6 +69,77 @@ public sealed class LeafNodeManager : IAsyncDisposable
    /// </summary>
    public bool IsTlsEnabled => CurrentCertPath is not null;

+    /// <summary>
+    /// When true, all outbound leaf connections are disabled regardless of per-remote settings.
+    /// Go reference: leafnode.go isLeafConnectDisabled — global disable flag.
+    /// </summary>
+    public bool IsGloballyDisabled { get; private set; }
+
+    /// <summary>
+    /// Returns the number of remotes that have been individually disabled via
+    /// <see cref="DisableLeafConnect"/>.
+    /// </summary>
+    public int DisabledRemoteCount => _disabledRemotes.Count;
+
+    /// <summary>
+    /// Returns true when connections to <paramref name="remoteUrl"/> are currently disabled,
+    /// either because it was individually disabled or because all leaf connections are globally
+    /// disabled.
+    /// Go reference: leafnode.go isLeafConnectDisabled.
+    /// </summary>
+    public bool IsLeafConnectDisabled(string remoteUrl)
+        => IsGloballyDisabled || _disabledRemotes.ContainsKey(remoteUrl);
+
+    /// <summary>
+    /// Disables outbound leaf connections to the specified remote URL.
+    /// Has no effect if the remote is already disabled.
+    /// Go reference: leafnode.go isLeafConnectDisabled — per-remote disable tracking.
+    /// </summary>
+    public void DisableLeafConnect(string remoteUrl, string? reason = null)
+    {
+        _disabledRemotes.TryAdd(remoteUrl, true);
+        _logger.LogInformation(
+            "Leaf connect disabled for remote {RemoteUrl} (reason={Reason})",
+            remoteUrl, reason ?? "unspecified");
+    }
+
+    /// <summary>
+    /// Re-enables outbound leaf connections to the specified remote URL.
+    /// Has no effect if the remote was not disabled.
+    /// </summary>
+    public void EnableLeafConnect(string remoteUrl)
+    {
+        _disabledRemotes.TryRemove(remoteUrl, out _);
+        _logger.LogInformation("Leaf connect re-enabled for remote {RemoteUrl}", remoteUrl);
+    }
+
+    /// <summary>
+    /// Disables all outbound leaf connections by setting the global disable flag.
+    /// Per-remote disable state is preserved.
+    /// Go reference: leafnode.go isLeafConnectDisabled — global flag.
+    /// </summary>
+    public void DisableAllLeafConnections(string? reason = null)
+    {
+        IsGloballyDisabled = true;
+        _logger.LogInformation("All leaf connections globally disabled (reason={Reason})", reason ?? "unspecified");
+    }
+
+    /// <summary>
+    /// Clears the global disable flag so outbound leaf connections may resume.
+    /// Per-remote disable state is unchanged.
+    /// </summary>
+    public void EnableAllLeafConnections()
+    {
+        IsGloballyDisabled = false;
+        _logger.LogInformation("All leaf connections globally re-enabled");
+    }
+
+    /// <summary>
+    /// Returns a snapshot of the remote URLs that have been individually disabled via
+    /// <see cref="DisableLeafConnect"/>.
+    /// </summary>
+    public IReadOnlyList<string> GetDisabledRemotes() => [.. _disabledRemotes.Keys];
+
    /// <summary>
    /// Incremented each time <see cref="UpdateTlsConfig"/> detects a change and applies it.
    /// Useful for testing and observability.
@@ -284,6 +357,149 @@ public sealed class LeafNodeManager : IAsyncDisposable
            SubscribeAllowCount: connection.AllowedSubscribeSubjects.Count);
    }

+    /// <summary>
+    /// Validates whether a leaf connection can migrate its JetStream domain to a proposed value.
+    /// Clearing the domain (null/empty proposedDomain) is always valid.
+    /// If the proposed domain matches the current domain, no migration is needed.
+    /// If another connection already uses the proposed domain, a conflict is reported.
+    /// Go reference: leafnode.go checkJetStreamMigrate.
+    /// </summary>
+    public JetStreamMigrationResult CheckJetStreamMigrate(string connectionId, string? proposedDomain)
+    {
+        if (!_connections.TryGetValue(connectionId, out var connection))
+            return new JetStreamMigrationResult(false, JetStreamMigrationStatus.ConnectionNotFound, $"Connection '{connectionId}' not found");
+
+        // Clearing domain is always valid.
+        if (string.IsNullOrEmpty(proposedDomain))
+            return new JetStreamMigrationResult(true, JetStreamMigrationStatus.Valid, null);
+
+        // If current domain already matches, no migration needed.
+        if (string.Equals(connection.JetStreamDomain, proposedDomain, StringComparison.Ordinal))
+            return new JetStreamMigrationResult(true, JetStreamMigrationStatus.NoChangeNeeded, null);
+
+        // Check for domain conflict with other connections.
+        foreach (var (key, conn) in _connections)
+        {
+            if (string.Equals(key, connectionId, StringComparison.Ordinal))
+                continue;
+            if (string.Equals(conn.JetStreamDomain, proposedDomain, StringComparison.Ordinal))
+                return new JetStreamMigrationResult(false, JetStreamMigrationStatus.DomainConflict,
+                    $"Domain '{proposedDomain}' is already in use by another connection");
+        }
+
+        return new JetStreamMigrationResult(true, JetStreamMigrationStatus.Valid, null);
+    }
+
+    /// <summary>
+    /// Returns the distinct set of JetStream domains across all active connections.
+    /// Connections without a domain (null/empty JetStreamDomain) are excluded.
+    /// Go reference: leafnode.go — per-connection domain tracking.
+    /// </summary>
+    public IReadOnlyList<string> GetActiveJetStreamDomains()
+    {
+        var domains = new HashSet<string>(StringComparer.Ordinal);
+        foreach (var conn in _connections.Values)
+        {
+            if (!string.IsNullOrEmpty(conn.JetStreamDomain))
+                domains.Add(conn.JetStreamDomain);
+        }
+        return [.. domains];
+    }
+
+    /// <summary>
+    /// Returns true if any currently active connection is associated with the specified JetStream domain.
+    /// Go reference: leafnode.go — checkJetStreamMigrate domain-in-use check.
+    /// </summary>
+    public bool IsJetStreamDomainInUse(string domain)
+    {
+        foreach (var conn in _connections.Values)
+        {
+            if (string.Equals(conn.JetStreamDomain, domain, StringComparison.Ordinal))
+                return true;
+        }
+        return false;
+    }
+
+    /// <summary>
+    /// Count of connections that have a non-null, non-empty JetStream domain assigned.
+    /// Go reference: leafnode.go — per-connection jsClusterDomain field.
+    /// </summary>
+    public int JetStreamEnabledConnectionCount
+    {
+        get
+        {
+            var count = 0;
+            foreach (var conn in _connections.Values)
+            {
+                if (!string.IsNullOrEmpty(conn.JetStreamDomain))
+                    count++;
+            }
+            return count;
+        }
+    }
+
+    /// <summary>
+    /// Registers a leaf cluster topology entry.
+    /// Returns false if a cluster with the same name is already registered.
+    /// Go reference: leafnode.go registerLeafNodeCluster.
+    /// </summary>
+    public bool RegisterLeafNodeCluster(string clusterName, string gatewayUrl, int connectionCount)
+    {
+        var info = new LeafClusterInfo
+        {
+            ClusterName = clusterName,
+            GatewayUrl = gatewayUrl,
+            ConnectionCount = connectionCount,
+        };
+        return _leafClusters.TryAdd(clusterName, info);
+    }
+
+    /// <summary>
+    /// Removes a leaf cluster entry by name.
+    /// Returns false if no entry with that name exists.
+    /// Go reference: leafnode.go — leaf cluster topology removal.
+    /// </summary>
+    public bool UnregisterLeafNodeCluster(string clusterName) =>
+        _leafClusters.TryRemove(clusterName, out _);
+
+    /// <summary>
+    /// Returns true if a leaf cluster with the given name is currently registered.
+    /// Go reference: leafnode.go — leaf cluster topology lookup.
+    /// </summary>
+    public bool HasLeafNodeCluster(string clusterName) =>
+        _leafClusters.ContainsKey(clusterName);
+
+    /// <summary>
+    /// Returns the <see cref="LeafClusterInfo"/> for the named cluster, or null if not registered.
+    /// Go reference: leafnode.go — leaf cluster topology lookup.
+    /// </summary>
+    public LeafClusterInfo? GetLeafNodeCluster(string clusterName) =>
+        _leafClusters.TryGetValue(clusterName, out var info) ? info : null;
+
+    /// <summary>
+    /// Returns all registered leaf cluster entries as a read-only list.
+    /// Go reference: leafnode.go — leaf cluster topology enumeration.
+    /// </summary>
+    public IReadOnlyList<LeafClusterInfo> GetAllLeafClusters() =>
+        [.. _leafClusters.Values];
+
+    /// <summary>
+    /// Count of registered leaf cluster topology entries.
+    /// Go reference: leafnode.go — leaf cluster topology count.
+    /// </summary>
+    public int LeafClusterCount => _leafClusters.Count;
+
+    /// <summary>
+    /// Updates the connection count for the named leaf cluster.
+    /// No-op if the cluster is not registered.
+    /// Go reference: leafnode.go — leaf cluster connection count update.
+    /// </summary>
+    public void UpdateLeafClusterConnectionCount(string clusterName, int newCount)
+    {
+        if (_leafClusters.TryGetValue(clusterName, out var info))
+            info.ConnectionCount = newCount;
+    }
+
    /// <summary>
    /// Returns all current connection IDs. Useful for tests and monitoring.
    /// </summary>
@@ -578,3 +794,42 @@ public enum LeafValidationError
    DuplicateConnection,
    JetStreamDomainConflict
 }
+
+/// <summary>
+/// Describes the outcome of a <see cref="LeafNodeManager.CheckJetStreamMigrate"/> call.
+/// </summary>
+/// <param name="Valid">True when migration to the proposed domain is allowed.</param>
+/// <param name="Status">Detailed status code for the migration check.</param>
+/// <param name="Error">Human-readable error message when <paramref name="Valid"/> is false, otherwise null.</param>
+public sealed record JetStreamMigrationResult(
+    bool Valid,
+    JetStreamMigrationStatus Status,
+    string? Error);
+
+/// <summary>
+/// Status codes for <see cref="JetStreamMigrationResult"/>.
+/// Go reference: leafnode.go checkJetStreamMigrate return values.
+/// </summary>
+public enum JetStreamMigrationStatus
+{
+    /// <summary>Migration to the proposed domain is allowed.</summary>
+    Valid,
+    /// <summary>The specified connection ID was not found.</summary>
+    ConnectionNotFound,
+    /// <summary>The proposed domain is identical to the current domain — no migration required.</summary>
+    NoChangeNeeded,
+    /// <summary>Another connection already uses the proposed domain.</summary>
+    DomainConflict
+}
+
+/// <summary>
+/// Holds topology information for a registered leaf cluster entry.
+/// Go reference: leafnode.go — leaf cluster registration / registerLeafNodeCluster.
+/// </summary>
+public sealed class LeafClusterInfo
+{
+    public required string ClusterName { get; init; }
+    public required string GatewayUrl { get; init; }
+    public int ConnectionCount { get; set; }
+    public DateTime RegisteredAt { get; init; } = DateTime.UtcNow;
+}