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:
@@ -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);
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user