feat: add route handshake lifecycle

2026-02-23 05:46:59 -05:00
parent 44d426a7c5
commit 5f98e53d62
6 changed files with 355 additions and 0 deletions
--- a/src/NATS.Server/Routes/RouteManager.cs
+++ b/src/NATS.Server/Routes/RouteManager.cs
@@ -0,0 +1,190 @@
+using System.Collections.Concurrent;
+using System.Net;
+using System.Net.Sockets;
+using Microsoft.Extensions.Logging;
+using NATS.Server.Configuration;
+
+namespace NATS.Server.Routes;
+
+public sealed class RouteManager : IAsyncDisposable
+{
+    private readonly ClusterOptions _options;
+    private readonly ServerStats _stats;
+    private readonly string _serverId;
+    private readonly ILogger<RouteManager> _logger;
+    private readonly ConcurrentDictionary<string, RouteConnection> _routes = new(StringComparer.Ordinal);
+
+    private CancellationTokenSource? _cts;
+    private Socket? _listener;
+    private Task? _acceptLoopTask;
+
+    public string ListenEndpoint => $"{_options.Host}:{_options.Port}";
+
+    public RouteManager(ClusterOptions options, ServerStats stats, string serverId, ILogger<RouteManager> logger)
+    {
+        _options = options;
+        _stats = stats;
+        _serverId = serverId;
+        _logger = logger;
+    }
+
+    public Task StartAsync(CancellationToken ct)
+    {
+        _cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
+        _listener = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
+        _listener.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, true);
+        _listener.Bind(new IPEndPoint(IPAddress.Parse(_options.Host), _options.Port));
+        _listener.Listen(128);
+
+        if (_options.Port == 0)
+            _options.Port = ((IPEndPoint)_listener.LocalEndPoint!).Port;
+
+        _acceptLoopTask = Task.Run(() => AcceptLoopAsync(_cts.Token));
+        foreach (var route in _options.Routes.Distinct(StringComparer.OrdinalIgnoreCase))
+            _ = Task.Run(() => ConnectToRouteWithRetryAsync(route, _cts.Token));
+
+        return Task.CompletedTask;
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        if (_cts == null)
+            return;
+
+        await _cts.CancelAsync();
+        _listener?.Dispose();
+
+        if (_acceptLoopTask != null)
+            await _acceptLoopTask.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
+
+        foreach (var route in _routes.Values)
+            await route.DisposeAsync();
+
+        _routes.Clear();
+        Interlocked.Exchange(ref _stats.Routes, 0);
+        _cts.Dispose();
+        _cts = null;
+    }
+
+    private async Task AcceptLoopAsync(CancellationToken ct)
+    {
+        while (!ct.IsCancellationRequested)
+        {
+            Socket socket;
+            try
+            {
+                socket = await _listener!.AcceptAsync(ct);
+            }
+            catch (OperationCanceledException)
+            {
+                break;
+            }
+            catch (ObjectDisposedException)
+            {
+                break;
+            }
+            catch (Exception ex)
+            {
+                _logger.LogDebug(ex, "Route accept loop error");
+                break;
+            }
+
+            _ = Task.Run(() => HandleInboundRouteAsync(socket, ct), ct);
+        }
+    }
+
+    private async Task HandleInboundRouteAsync(Socket socket, CancellationToken ct)
+    {
+        var route = new RouteConnection(socket);
+        try
+        {
+            await route.PerformInboundHandshakeAsync(_serverId, ct);
+            Register(route);
+        }
+        catch (Exception ex)
+        {
+            _logger.LogDebug(ex, "Inbound route handshake failed");
+            await route.DisposeAsync();
+        }
+    }
+
+    private async Task ConnectToRouteWithRetryAsync(string route, CancellationToken ct)
+    {
+        while (!ct.IsCancellationRequested)
+        {
+            try
+            {
+                var endPoint = ParseRouteEndpoint(route);
+                var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
+                await socket.ConnectAsync(endPoint.Address, endPoint.Port, ct);
+                var connection = new RouteConnection(socket);
+                await connection.PerformOutboundHandshakeAsync(_serverId, ct);
+                Register(connection);
+                return;
+            }
+            catch (OperationCanceledException)
+            {
+                return;
+            }
+            catch (Exception ex)
+            {
+                _logger.LogDebug(ex, "Failed to connect route seed {Route}", route);
+            }
+
+            try
+            {
+                await Task.Delay(250, ct);
+            }
+            catch (OperationCanceledException)
+            {
+                return;
+            }
+        }
+    }
+
+    private void Register(RouteConnection route)
+    {
+        var key = $"{route.RemoteServerId}:{route.RemoteEndpoint}";
+        if (!_routes.TryAdd(key, route))
+        {
+            _ = route.DisposeAsync();
+            return;
+        }
+
+        Interlocked.Increment(ref _stats.Routes);
+        _ = Task.Run(() => WatchRouteAsync(key, route, _cts!.Token));
+    }
+
+    private async Task WatchRouteAsync(string key, RouteConnection route, CancellationToken ct)
+    {
+        try
+        {
+            await route.WaitUntilClosedAsync(ct);
+        }
+        catch (OperationCanceledException)
+        {
+            // Shutdown path.
+        }
+        catch (Exception ex)
+        {
+            _logger.LogDebug(ex, "Route {RouteKey} closed with error", key);
+        }
+        finally
+        {
+            if (_routes.TryRemove(key, out _))
+                Interlocked.Decrement(ref _stats.Routes);
+
+            await route.DisposeAsync();
+        }
+    }
+
+    private static IPEndPoint ParseRouteEndpoint(string route)
+    {
+        var trimmed = route.Trim();
+        var parts = trimmed.Split(':', 2, StringSplitOptions.TrimEntries | StringSplitOptions.RemoveEmptyEntries);
+        if (parts.Length != 2)
+            throw new FormatException($"Invalid route endpoint: '{route}'");
+
+        return new IPEndPoint(IPAddress.Parse(parts[0]), int.Parse(parts[1]));
+    }
+}