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);