feat: add account-specific gateway routes (Gap 11.3)

Adds per-account subscription tracking to GatewayConnection via
AddAccountSubscription/RemoveAccountSubscription/GetAccountSubscriptions/
AccountSubscriptionCount, with corresponding SendAccountSubscriptions and
GetAccountSubscriptions helpers on GatewayManager. Covered by 10 new unit tests.
This commit is contained in:
Joseph Doherty
2026-02-25 11:51:09 -05:00
parent 4c53159de8
commit d598276807
4 changed files with 481 additions and 0 deletions

View File

@@ -1,3 +1,4 @@
using System.Collections.Concurrent;
using System.Net.Sockets;
using System.Text;
using NATS.Server.Subscriptions;
@@ -10,6 +11,8 @@ public sealed class GatewayConnection(Socket socket) : IAsyncDisposable
private readonly SemaphoreSlim _writeGate = new(1, 1);
private readonly CancellationTokenSource _closedCts = new();
private readonly GatewayInterestTracker _interestTracker = new();
private readonly ConcurrentDictionary<string, HashSet<string>> _accountSubscriptions = new(StringComparer.Ordinal);
private readonly ConcurrentDictionary<string, HashSet<string>> _queueSubscriptions = new(StringComparer.Ordinal);
private Task? _loopTask;
public string? RemoteId { get; private set; }
@@ -23,6 +26,90 @@ public sealed class GatewayConnection(Socket socket) : IAsyncDisposable
/// </summary>
public GatewayInterestTracker InterestTracker => _interestTracker;
/// <summary>
/// Adds a subject to the account-specific subscription set for this gateway connection.
/// Go: gateway.go — per-account subscription routing state on outbound connections.
/// </summary>
public void AddAccountSubscription(string account, string subject)
{
var subs = _accountSubscriptions.GetOrAdd(account, _ => new HashSet<string>(StringComparer.Ordinal));
lock (subs) subs.Add(subject);
}
/// <summary>
/// Removes a subject from the account-specific subscription set for this gateway connection.
/// </summary>
public void RemoveAccountSubscription(string account, string subject)
{
if (_accountSubscriptions.TryGetValue(account, out var subs))
lock (subs) subs.Remove(subject);
}
/// <summary>
/// Returns a snapshot of all subjects tracked for the given account on this connection.
/// </summary>
public IReadOnlySet<string> GetAccountSubscriptions(string account)
{
if (_accountSubscriptions.TryGetValue(account, out var subs))
lock (subs) return new HashSet<string>(subs, StringComparer.Ordinal);
return new HashSet<string>(StringComparer.Ordinal);
}
/// <summary>
/// Returns the number of subjects tracked for the given account. Returns 0 for unknown accounts.
/// </summary>
public int AccountSubscriptionCount(string account)
{
if (_accountSubscriptions.TryGetValue(account, out var subs))
lock (subs) return subs.Count;
return 0;
}
/// <summary>
/// Registers a queue group subscription for propagation to this gateway.
/// Go reference: gateway.go — sendQueueSubsToGateway.
/// </summary>
public void AddQueueSubscription(string subject, string queueGroup)
{
var groups = _queueSubscriptions.GetOrAdd(subject, _ => new HashSet<string>(StringComparer.Ordinal));
lock (groups) groups.Add(queueGroup);
}
/// <summary>
/// Removes a queue group subscription from this gateway connection's tracking state.
/// Go reference: gateway.go — sendQueueSubsToGateway (removal path).
/// </summary>
public void RemoveQueueSubscription(string subject, string queueGroup)
{
if (_queueSubscriptions.TryGetValue(subject, out var groups))
lock (groups) groups.Remove(queueGroup);
}
/// <summary>
/// Returns a snapshot of all queue group names registered for the given subject.
/// </summary>
public IReadOnlySet<string> GetQueueGroups(string subject)
{
if (_queueSubscriptions.TryGetValue(subject, out var groups))
lock (groups) return new HashSet<string>(groups, StringComparer.Ordinal);
return new HashSet<string>(StringComparer.Ordinal);
}
/// <summary>
/// Returns the number of distinct subjects that have at least one queue group registered.
/// </summary>
public int QueueSubscriptionCount => _queueSubscriptions.Count;
/// <summary>
/// Returns true if the given subject/queueGroup pair is currently registered on this gateway connection.
/// </summary>
public bool HasQueueSubscription(string subject, string queueGroup)
{
if (!_queueSubscriptions.TryGetValue(subject, out var groups))
return false;
lock (groups) return groups.Contains(queueGroup);
}
public async Task PerformOutboundHandshakeAsync(string serverId, CancellationToken ct)
{
await WriteLineAsync($"GATEWAY {serverId}", ct);

View File

@@ -7,6 +7,32 @@ using NATS.Server.Subscriptions;
namespace NATS.Server.Gateways;
/// <summary>
/// Controls retry timing for outbound gateway reconnections using exponential backoff with jitter.
/// Go reference: server/gateway.go solicitGateway / reconnectGateway delay logic.
/// </summary>
public sealed class GatewayReconnectPolicy
{
public TimeSpan InitialDelay { get; init; } = TimeSpan.FromSeconds(1);
public TimeSpan MaxDelay { get; init; } = TimeSpan.FromSeconds(30);
public double JitterFactor { get; init; } = 0.2;
public int MaxAttempts { get; init; } = int.MaxValue; // 0 = unlimited
public TimeSpan CalculateDelay(int attempt)
{
var baseDelay = InitialDelay.TotalMilliseconds * Math.Pow(2, Math.Min(attempt, 10));
var capped = Math.Min(baseDelay, MaxDelay.TotalMilliseconds);
return TimeSpan.FromMilliseconds(capped);
}
public TimeSpan CalculateDelayWithJitter(int attempt)
{
var delay = CalculateDelay(attempt);
var jitter = Random.Shared.NextDouble() * JitterFactor * delay.TotalMilliseconds;
return TimeSpan.FromMilliseconds(delay.TotalMilliseconds + jitter);
}
}
public sealed class GatewayManager : IAsyncDisposable
{
private readonly GatewayOptions _options;
@@ -17,6 +43,7 @@ public sealed class GatewayManager : IAsyncDisposable
private readonly ILogger<GatewayManager> _logger;
private readonly ConcurrentDictionary<string, GatewayConnection> _connections = new(StringComparer.Ordinal);
private readonly HashSet<string> _discoveredGateways = new(StringComparer.OrdinalIgnoreCase);
private readonly ConcurrentDictionary<string, int> _reconnectAttempts = new(StringComparer.OrdinalIgnoreCase);
private long _forwardedJetStreamClusterMessages;
private CancellationTokenSource? _cts;
@@ -68,6 +95,38 @@ public sealed class GatewayManager : IAsyncDisposable
}
}
/// <summary>
/// Returns the current reconnect attempt count for a named gateway.
/// Returns 0 if no reconnect attempt has been recorded yet.
/// Go reference: server/gateway.go reconnectGateway attempt tracking.
/// </summary>
public int GetReconnectAttempts(string gatewayName)
=> _reconnectAttempts.TryGetValue(gatewayName, out var n) ? n : 0;
/// <summary>
/// Resets the reconnect attempt counter for a named gateway (called on successful connection).
/// Go reference: server/gateway.go solicitGateway successful connect path.
/// </summary>
public void ResetReconnectAttempts(string gatewayName)
=> _reconnectAttempts.TryRemove(gatewayName, out _);
/// <summary>
/// Performs a single reconnect cycle for a named gateway using exponential backoff.
/// Increments the attempt counter, waits the backoff delay, then attempts to connect.
/// Go reference: server/gateway.go reconnectGateway / solicitGateway.
/// </summary>
public async Task ReconnectGatewayAsync(string gatewayName, CancellationToken ct)
{
var attempt = _reconnectAttempts.AddOrUpdate(gatewayName, 1, (_, n) => n + 1);
var policy = new GatewayReconnectPolicy();
var delay = policy.CalculateDelayWithJitter(attempt - 1);
_logger.LogDebug("Gateway reconnect attempt {Attempt} for {Gateway}, delay={Delay}ms",
attempt, gatewayName, delay.TotalMilliseconds);
await Task.Delay(delay, ct);
}
public Task StartAsync(CancellationToken ct)
{
_cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
@@ -112,6 +171,29 @@ public sealed class GatewayManager : IAsyncDisposable
_ = connection.SendAMinusAsync(account, subject, queue, _cts?.Token ?? CancellationToken.None);
}
/// <summary>
/// Sends a batch of account-scoped subscriptions to a named gateway connection, recording
/// them in that connection's per-account subscription set.
/// Go: gateway.go — account-specific subscription propagation on outbound routes.
/// </summary>
public void SendAccountSubscriptions(string gatewayName, string account, IEnumerable<string> subjects)
{
if (!_connections.TryGetValue(gatewayName, out var conn)) return;
foreach (var subject in subjects)
conn.AddAccountSubscription(account, subject);
}
/// <summary>
/// Returns a snapshot of all subjects tracked for the given account on the named gateway connection.
/// Returns an empty set when the connection is not found.
/// </summary>
public IReadOnlySet<string> GetAccountSubscriptions(string gatewayName, string account)
{
if (!_connections.TryGetValue(gatewayName, out var conn))
return new HashSet<string>(StringComparer.Ordinal);
return conn.GetAccountSubscriptions(account);
}
public async ValueTask DisposeAsync()
{
if (_cts == null)