feat: add route handshake lifecycle
This commit is contained in:
Joseph Doherty
81
src/NATS.Server/Routes/RouteConnection.cs
Normal file
81
src/NATS.Server/Routes/RouteConnection.cs
Normal file
@@ -0,0 +1,81 @@
|
||||
using System.Net.Sockets;
|
||||
using System.Text;
|
||||
|
||||
namespace NATS.Server.Routes;
|
||||
|
||||
public sealed class RouteConnection(Socket socket) : IAsyncDisposable
|
||||
{
|
||||
private readonly Socket _socket = socket;
|
||||
private readonly NetworkStream _stream = new(socket, ownsSocket: true);
|
||||
|
||||
public string? RemoteServerId { get; private set; }
|
||||
public string RemoteEndpoint => _socket.RemoteEndPoint?.ToString() ?? Guid.NewGuid().ToString("N");
|
||||
|
||||
public async Task PerformOutboundHandshakeAsync(string serverId, CancellationToken ct)
|
||||
{
|
||||
await WriteLineAsync($"ROUTE {serverId}", ct);
|
||||
var line = await ReadLineAsync(ct);
|
||||
RemoteServerId = ParseHandshake(line);
|
||||
}
|
||||
|
||||
public async Task PerformInboundHandshakeAsync(string serverId, CancellationToken ct)
|
||||
{
|
||||
var line = await ReadLineAsync(ct);
|
||||
RemoteServerId = ParseHandshake(line);
|
||||
await WriteLineAsync($"ROUTE {serverId}", ct);
|
||||
}
|
||||
|
||||
public async Task WaitUntilClosedAsync(CancellationToken ct)
|
||||
{
|
||||
var buffer = new byte[1024];
|
||||
while (!ct.IsCancellationRequested)
|
||||
{
|
||||
var bytesRead = await _stream.ReadAsync(buffer, ct);
|
||||
if (bytesRead == 0)
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
public async ValueTask DisposeAsync()
|
||||
{
|
||||
await _stream.DisposeAsync();
|
||||
}
|
||||
|
||||
private async Task WriteLineAsync(string line, CancellationToken ct)
|
||||
{
|
||||
var bytes = Encoding.ASCII.GetBytes($"{line}\r\n");
|
||||
await _stream.WriteAsync(bytes, ct);
|
||||
await _stream.FlushAsync(ct);
|
||||
}
|
||||
|
||||
private async Task<string> ReadLineAsync(CancellationToken ct)
|
||||
{
|
||||
var bytes = new List<byte>(64);
|
||||
var single = new byte[1];
|
||||
while (true)
|
||||
{
|
||||
var read = await _stream.ReadAsync(single, ct);
|
||||
if (read == 0)
|
||||
throw new IOException("Route connection closed during handshake");
|
||||
|
||||
if (single[0] == (byte)'\n')
|
||||
break;
|
||||
if (single[0] != (byte)'\r')
|
||||
bytes.Add(single[0]);
|
||||
}
|
||||
|
||||
return Encoding.ASCII.GetString([.. bytes]);
|
||||
}
|
||||
|
||||
private static string ParseHandshake(string line)
|
||||
{
|
||||
if (!line.StartsWith("ROUTE ", StringComparison.OrdinalIgnoreCase))
|
||||
throw new InvalidOperationException("Invalid route handshake");
|
||||
|
||||
var id = line[6..].Trim();
|
||||
if (id.Length == 0)
|
||||
throw new InvalidOperationException("Route handshake missing server id");
|
||||
|
||||
return id;
|
||||
}
|
||||
}
|
||||
190
src/NATS.Server/Routes/RouteManager.cs
Normal file
190
src/NATS.Server/Routes/RouteManager.cs
Normal file
@@ -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]));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user