feat: complete final jetstream parity transport and runtime baselines

This commit is contained in:
Joseph Doherty
2026-02-23 11:04:43 -05:00
parent 53585012f3
commit 8bce096f55
61 changed files with 2655 additions and 129 deletions

View File

@@ -1,5 +1,6 @@
using System.Net.Sockets;
using System.Text;
using NATS.Server.Subscriptions;
namespace NATS.Server.Routes;
@@ -7,9 +8,14 @@ public sealed class RouteConnection(Socket socket) : IAsyncDisposable
{
private readonly Socket _socket = socket;
private readonly NetworkStream _stream = new(socket, ownsSocket: true);
private readonly SemaphoreSlim _writeGate = new(1, 1);
private readonly CancellationTokenSource _closedCts = new();
private Task? _frameLoopTask;
public string? RemoteServerId { get; private set; }
public string RemoteEndpoint => _socket.RemoteEndPoint?.ToString() ?? Guid.NewGuid().ToString("N");
public Func<RemoteSubscription, Task>? RemoteSubscriptionReceived { get; set; }
public Func<RouteMessage, Task>? RoutedMessageReceived { get; set; }
public async Task PerformOutboundHandshakeAsync(string serverId, CancellationToken ct)
{
@@ -25,27 +31,168 @@ public sealed class RouteConnection(Socket socket) : IAsyncDisposable
await WriteLineAsync($"ROUTE {serverId}", ct);
}
public void StartFrameLoop(CancellationToken ct)
{
if (_frameLoopTask != null)
return;
var linked = CancellationTokenSource.CreateLinkedTokenSource(ct, _closedCts.Token);
_frameLoopTask = Task.Run(() => ReadFramesAsync(linked.Token), linked.Token);
}
public async Task SendRsPlusAsync(string subject, string? queue, CancellationToken ct)
{
var frame = queue is { Length: > 0 }
? $"RS+ {subject} {queue}"
: $"RS+ {subject}";
await WriteLineAsync(frame, ct);
}
public async Task SendRsMinusAsync(string subject, string? queue, CancellationToken ct)
{
var frame = queue is { Length: > 0 }
? $"RS- {subject} {queue}"
: $"RS- {subject}";
await WriteLineAsync(frame, ct);
}
public async Task SendRmsgAsync(string subject, string? replyTo, ReadOnlyMemory<byte> payload, CancellationToken ct)
{
var replyToken = string.IsNullOrEmpty(replyTo) ? "-" : replyTo;
await _writeGate.WaitAsync(ct);
try
{
var control = Encoding.ASCII.GetBytes($"RMSG {subject} {replyToken} {payload.Length}\r\n");
await _stream.WriteAsync(control, ct);
if (!payload.IsEmpty)
await _stream.WriteAsync(payload, ct);
await _stream.WriteAsync("\r\n"u8.ToArray(), ct);
await _stream.FlushAsync(ct);
}
finally
{
_writeGate.Release();
}
}
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;
}
if (_frameLoopTask == null)
return;
using var linked = CancellationTokenSource.CreateLinkedTokenSource(ct, _closedCts.Token);
await _frameLoopTask.WaitAsync(linked.Token);
}
public async ValueTask DisposeAsync()
{
await _closedCts.CancelAsync();
if (_frameLoopTask != null)
await _frameLoopTask.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
_closedCts.Dispose();
_writeGate.Dispose();
await _stream.DisposeAsync();
}
private async Task ReadFramesAsync(CancellationToken ct)
{
while (!ct.IsCancellationRequested)
{
string line;
try
{
line = await ReadLineAsync(ct);
}
catch (OperationCanceledException)
{
break;
}
catch (IOException)
{
break;
}
if (line.StartsWith("RS+ ", StringComparison.Ordinal))
{
var parts = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
if (parts.Length >= 2 && RemoteSubscriptionReceived != null)
{
var queue = parts.Length >= 3 ? parts[2] : null;
await RemoteSubscriptionReceived(new RemoteSubscription(parts[1], queue, RemoteServerId ?? string.Empty));
}
continue;
}
if (line.StartsWith("RS- ", StringComparison.Ordinal))
{
var parts = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
if (parts.Length >= 2 && RemoteSubscriptionReceived != null)
{
var queue = parts.Length >= 3 ? parts[2] : null;
await RemoteSubscriptionReceived(RemoteSubscription.Removal(parts[1], queue, RemoteServerId ?? string.Empty));
}
continue;
}
if (!line.StartsWith("RMSG ", StringComparison.Ordinal))
continue;
var args = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
if (args.Length < 4)
continue;
var subject = args[1];
var reply = args[2] == "-" ? null : args[2];
if (!int.TryParse(args[3], out var size) || size < 0)
continue;
var payload = await ReadPayloadAsync(size, ct);
if (RoutedMessageReceived != null)
await RoutedMessageReceived(new RouteMessage(subject, reply, payload));
}
}
private async Task<ReadOnlyMemory<byte>> ReadPayloadAsync(int size, CancellationToken ct)
{
var payload = new byte[size];
var offset = 0;
while (offset < size)
{
var read = await _stream.ReadAsync(payload.AsMemory(offset, size - offset), ct);
if (read == 0)
throw new IOException("Route connection closed during payload read");
offset += read;
}
var trailer = new byte[2];
var trailerRead = 0;
while (trailerRead < 2)
{
var read = await _stream.ReadAsync(trailer.AsMemory(trailerRead, 2 - trailerRead), ct);
if (read == 0)
throw new IOException("Route connection closed during payload trailer read");
trailerRead += read;
}
if (trailer[0] != (byte)'\r' || trailer[1] != (byte)'\n')
throw new IOException("Invalid route payload trailer");
return payload;
}
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);
await _writeGate.WaitAsync(ct);
try
{
var bytes = Encoding.ASCII.GetBytes($"{line}\r\n");
await _stream.WriteAsync(bytes, ct);
await _stream.FlushAsync(ct);
}
finally
{
_writeGate.Release();
}
}
private async Task<string> ReadLineAsync(CancellationToken ct)
@@ -56,7 +203,7 @@ public sealed class RouteConnection(Socket socket) : IAsyncDisposable
{
var read = await _stream.ReadAsync(single, ct);
if (read == 0)
throw new IOException("Route connection closed during handshake");
throw new IOException("Route connection closed");
if (single[0] == (byte)'\n')
break;
@@ -79,3 +226,5 @@ public sealed class RouteConnection(Socket socket) : IAsyncDisposable
return id;
}
}
public sealed record RouteMessage(string Subject, string? ReplyTo, ReadOnlyMemory<byte> Payload);

View File

@@ -15,6 +15,7 @@ public sealed class RouteManager : IAsyncDisposable
private readonly string _serverId;
private readonly ILogger<RouteManager> _logger;
private readonly Action<RemoteSubscription> _remoteSubSink;
private readonly Action<RouteMessage> _routedMessageSink;
private readonly ConcurrentDictionary<string, RouteConnection> _routes = new(StringComparer.Ordinal);
private readonly ConcurrentDictionary<string, byte> _connectedServerIds = new(StringComparer.Ordinal);
@@ -29,12 +30,14 @@ public sealed class RouteManager : IAsyncDisposable
ServerStats stats,
string serverId,
Action<RemoteSubscription> remoteSubSink,
Action<RouteMessage> routedMessageSink,
ILogger<RouteManager> logger)
{
_options = options;
_stats = stats;
_serverId = serverId;
_remoteSubSink = remoteSubSink;
_routedMessageSink = routedMessageSink;
_logger = logger;
}
@@ -51,8 +54,12 @@ public sealed class RouteManager : IAsyncDisposable
_options.Port = ((IPEndPoint)_listener.LocalEndPoint!).Port;
_acceptLoopTask = Task.Run(() => AcceptLoopAsync(_cts.Token));
var poolSize = Math.Max(_options.PoolSize, 1);
foreach (var route in _options.Routes.Distinct(StringComparer.OrdinalIgnoreCase))
_ = Task.Run(() => ConnectToRouteWithRetryAsync(route, _cts.Token));
{
for (var i = 0; i < poolSize; i++)
_ = Task.Run(() => ConnectToRouteWithRetryAsync(route, _cts.Token));
}
return Task.CompletedTask;
}
@@ -81,17 +88,33 @@ public sealed class RouteManager : IAsyncDisposable
public void PropagateLocalSubscription(string subject, string? queue)
{
if (_connectedServerIds.IsEmpty)
if (_routes.IsEmpty)
return;
var remoteSub = new RemoteSubscription(subject, queue, _serverId);
foreach (var peerId in _connectedServerIds.Keys)
foreach (var route in _routes.Values)
{
if (Managers.TryGetValue(peerId, out var peer))
peer.ReceiveRemoteSubscription(remoteSub);
_ = route.SendRsPlusAsync(subject, queue, _cts?.Token ?? CancellationToken.None);
}
}
public void PropagateLocalUnsubscription(string subject, string? queue)
{
if (_routes.IsEmpty)
return;
foreach (var route in _routes.Values)
_ = route.SendRsMinusAsync(subject, queue, _cts?.Token ?? CancellationToken.None);
}
public async Task ForwardRoutedMessageAsync(string subject, string? replyTo, ReadOnlyMemory<byte> payload, CancellationToken ct)
{
if (_routes.IsEmpty)
return;
foreach (var route in _routes.Values)
await route.SendRmsgAsync(subject, replyTo, payload, ct);
}
private async Task AcceptLoopAsync(CancellationToken ct)
{
while (!ct.IsCancellationRequested)
@@ -170,7 +193,7 @@ public sealed class RouteManager : IAsyncDisposable
private void Register(RouteConnection route)
{
var key = $"{route.RemoteServerId}:{route.RemoteEndpoint}";
var key = $"{route.RemoteServerId}:{route.RemoteEndpoint}:{Guid.NewGuid():N}";
if (!_routes.TryAdd(key, route))
{
_ = route.DisposeAsync();
@@ -180,6 +203,18 @@ public sealed class RouteManager : IAsyncDisposable
if (route.RemoteServerId is { Length: > 0 } remoteServerId)
_connectedServerIds[remoteServerId] = 0;
route.RemoteSubscriptionReceived = sub =>
{
_remoteSubSink(sub);
return Task.CompletedTask;
};
route.RoutedMessageReceived = msg =>
{
_routedMessageSink(msg);
return Task.CompletedTask;
};
route.StartFrameLoop(_cts!.Token);
Interlocked.Increment(ref _stats.Routes);
_ = Task.Run(() => WatchRouteAsync(key, route, _cts!.Token));
}
@@ -217,8 +252,5 @@ public sealed class RouteManager : IAsyncDisposable
return new IPEndPoint(IPAddress.Parse(parts[0]), int.Parse(parts[1]));
}
private void ReceiveRemoteSubscription(RemoteSubscription sub)
{
_remoteSubSink(sub);
}
public int RouteCount => _routes.Count;
}