feat: complete final jetstream parity transport and runtime baselines
This commit is contained in:
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user