Merge branch 'codex/parser-span-retention'

This commit is contained in:
Joseph Doherty
2026-03-13 10:11:42 -04:00
12 changed files with 594 additions and 81 deletions

View File

@@ -10,6 +10,7 @@
<InternalsVisibleTo Include="NATS.Server.Monitoring.Tests" />
<InternalsVisibleTo Include="NATS.Server.Auth.Tests" />
<InternalsVisibleTo Include="NATS.Server.JetStream.Tests" />
<InternalsVisibleTo Include="NATS.Server.Benchmark.Tests" />
</ItemGroup>
<ItemGroup>
<FrameworkReference Include="Microsoft.AspNetCore.App" />

View File

@@ -407,23 +407,23 @@ public sealed class NatsClient : INatsClient, IDisposable
long localInMsgs = 0;
long localInBytes = 0;
while (_parser.TryParse(ref buffer, out var cmd))
while (_parser.TryParseView(ref buffer, out var cmdView))
{
Interlocked.Exchange(ref _lastIn, Environment.TickCount64);
// Handle Pub/HPub inline to allow ref parameter passing for stat batching.
// DispatchCommandAsync is async and cannot accept ref parameters.
if (cmd.Type is CommandType.Pub or CommandType.HPub
if (cmdView.Type is CommandType.Pub or CommandType.HPub
&& (!_authService.IsAuthRequired || ConnectReceived))
{
Interlocked.Exchange(ref _lastActivityTicks, DateTime.UtcNow.Ticks);
ProcessPub(cmd, ref localInMsgs, ref localInBytes);
ProcessPub(cmdView, ref localInMsgs, ref localInBytes);
if (ClientOpts?.Verbose == true)
WriteProtocol(NatsProtocol.OkBytes);
}
else
{
await DispatchCommandAsync(cmd, ct);
await DispatchCommandAsync(cmdView.Materialize(), ct);
}
}
@@ -702,46 +702,50 @@ public sealed class NatsClient : INatsClient, IDisposable
server.OnLocalUnsubscription(Account?.Name ?? Account.GlobalAccountName, sub.Subject, sub.Queue);
}
private void ProcessPub(ParsedCommand cmd, ref long localInMsgs, ref long localInBytes)
private void ProcessPub(ParsedCommandView cmd, ref long localInMsgs, ref long localInBytes)
{
var payloadMemory = cmd.GetPayloadMemory();
localInMsgs++;
localInBytes += cmd.Payload.Length;
localInBytes += payloadMemory.Length;
// Max payload validation (always, hard close)
if (cmd.Payload.Length > _options.MaxPayload)
if (payloadMemory.Length > _options.MaxPayload)
{
_logger.LogWarning("Client {ClientId} exceeded max payload: {Size} > {MaxPayload}",
Id, cmd.Payload.Length, _options.MaxPayload);
Id, payloadMemory.Length, _options.MaxPayload);
_ = SendErrAndCloseAsync(NatsProtocol.ErrMaxPayloadViolation, ClientClosedReason.MaxPayloadExceeded);
return;
}
var subject = Encoding.ASCII.GetString(cmd.Subject.Span);
// Pedantic mode: validate publish subject
if (ClientOpts?.Pedantic == true && !SubjectMatch.IsValidPublishSubject(cmd.Subject!))
if (ClientOpts?.Pedantic == true && !SubjectMatch.IsValidPublishSubject(subject))
{
_logger.LogDebug("Client {ClientId} invalid publish subject: {Subject}", Id, cmd.Subject);
_logger.LogDebug("Client {ClientId} invalid publish subject: {Subject}", Id, subject);
SendErr(NatsProtocol.ErrInvalidPublishSubject);
return;
}
// Permission check for publish
if (_permissions != null && !_permissions.IsPublishAllowed(cmd.Subject!))
if (_permissions != null && !_permissions.IsPublishAllowed(subject))
{
_logger.LogDebug("Client {ClientId} publish permission denied for {Subject}", Id, cmd.Subject);
_logger.LogDebug("Client {ClientId} publish permission denied for {Subject}", Id, subject);
SendErr(NatsProtocol.ErrPermissionsPublish);
return;
}
ReadOnlyMemory<byte> headers = default;
ReadOnlyMemory<byte> payload = cmd.Payload;
ReadOnlyMemory<byte> payload = payloadMemory;
if (cmd.Type == CommandType.HPub && cmd.HeaderSize > 0)
{
headers = cmd.Payload[..cmd.HeaderSize];
payload = cmd.Payload[cmd.HeaderSize..];
headers = payloadMemory[..cmd.HeaderSize];
payload = payloadMemory[cmd.HeaderSize..];
}
Router?.ProcessMessage(cmd.Subject!, cmd.ReplyTo, headers, payload, this);
var replyTo = cmd.ReplyTo.IsEmpty ? null : Encoding.ASCII.GetString(cmd.ReplyTo.Span);
Router?.ProcessMessage(subject, replyTo, headers, payload, this);
}
public void RecordJetStreamPubAck(PubAck ack)

View File

@@ -37,23 +37,21 @@ public readonly struct ParsedCommand
public sealed class NatsParser
{
private static readonly byte[] CrLfBytes = "\r\n"u8.ToArray();
private readonly int _maxPayload;
private static ReadOnlySpan<byte> CrLfBytes => "\r\n"u8;
private ILogger? _logger;
public ILogger? Logger { set => _logger = value; }
// State for split-packet payload reading
private bool _awaitingPayload;
private int _expectedPayloadSize;
private string? _pendingSubject;
private string? _pendingReplyTo;
private byte[]? _pendingSubject;
private byte[]? _pendingReplyTo;
private int _pendingHeaderSize;
private CommandType _pendingType;
private string _pendingOperation = string.Empty;
public NatsParser(int maxPayload = NatsProtocol.MaxPayloadSize, ILogger? logger = null)
{
_maxPayload = maxPayload;
_logger = logger;
}
@@ -71,12 +69,23 @@ public sealed class NatsParser
{
command = default;
if (!TryParseView(ref buffer, out var view))
return false;
command = view.Materialize();
return true;
}
internal bool TryParseView(ref ReadOnlySequence<byte> buffer, out ParsedCommandView command)
{
command = default;
if (_awaitingPayload)
return TryReadPayload(ref buffer, out command);
// Look for \r\n to find control line
var reader = new SequenceReader<byte>(buffer);
if (!reader.TryReadTo(out ReadOnlySequence<byte> line, CrLfBytes.AsSpan()))
if (!reader.TryReadTo(out ReadOnlySequence<byte> line, CrLfBytes))
return false;
// Control line size check
@@ -114,7 +123,7 @@ public sealed class NatsParser
case (byte)'p':
if (b1 == (byte)'i') // PING
{
command = ParsedCommand.Simple(CommandType.Ping, "PING");
command = ParsedCommandView.Simple(CommandType.Ping, "PING");
buffer = buffer.Slice(reader.Position);
TraceInOp("PING");
return true;
@@ -122,7 +131,7 @@ public sealed class NatsParser
if (b1 == (byte)'o') // PONG
{
command = ParsedCommand.Simple(CommandType.Pong, "PONG");
command = ParsedCommandView.Simple(CommandType.Pong, "PONG");
buffer = buffer.Slice(reader.Position);
TraceInOp("PONG");
return true;
@@ -168,7 +177,7 @@ public sealed class NatsParser
case (byte)'c':
if (b1 == (byte)'o') // CONNECT
{
command = ParseConnect(lineSpan);
command = ParseConnect(line);
buffer = buffer.Slice(reader.Position);
TraceInOp("CONNECT");
return true;
@@ -179,7 +188,7 @@ public sealed class NatsParser
case (byte)'i':
if (b1 == (byte)'n') // INFO
{
command = ParseInfo(lineSpan);
command = ParseInfo(line);
buffer = buffer.Slice(reader.Position);
TraceInOp("INFO");
return true;
@@ -188,13 +197,13 @@ public sealed class NatsParser
break;
case (byte)'+': // +OK
command = ParsedCommand.Simple(CommandType.Ok, "+OK");
command = ParsedCommandView.Simple(CommandType.Ok, "+OK");
buffer = buffer.Slice(reader.Position);
TraceInOp("+OK");
return true;
case (byte)'-': // -ERR
command = ParsedCommand.Simple(CommandType.Err, "-ERR");
command = ParsedCommandView.Simple(CommandType.Err, "-ERR");
buffer = buffer.Slice(reader.Position);
TraceInOp("-ERR");
return true;
@@ -227,7 +236,7 @@ public sealed class NatsParser
Span<byte> line,
ref ReadOnlySequence<byte> buffer,
SequencePosition afterLine,
out ParsedCommand command)
out ParsedCommandView command)
{
command = default;
@@ -236,19 +245,19 @@ public sealed class NatsParser
var argsSpan = line[4..];
int argCount = SplitArgs(argsSpan, ranges);
string subject;
string? reply = null;
byte[] subject;
byte[]? reply = null;
int size;
if (argCount == 2)
{
subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]);
subject = argsSpan[ranges[0]].ToArray();
size = ParseSize(argsSpan[ranges[1]]);
}
else if (argCount == 3)
{
subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]);
reply = Encoding.ASCII.GetString(argsSpan[ranges[1]]);
subject = argsSpan[ranges[0]].ToArray();
reply = argsSpan[ranges[1]].ToArray();
size = ParseSize(argsSpan[ranges[2]]);
}
else
@@ -277,7 +286,7 @@ public sealed class NatsParser
Span<byte> line,
ref ReadOnlySequence<byte> buffer,
SequencePosition afterLine,
out ParsedCommand command)
out ParsedCommandView command)
{
command = default;
@@ -286,20 +295,20 @@ public sealed class NatsParser
var argsSpan = line[5..];
int argCount = SplitArgs(argsSpan, ranges);
string subject;
string? reply = null;
byte[] subject;
byte[]? reply = null;
int hdrSize, totalSize;
if (argCount == 3)
{
subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]);
subject = argsSpan[ranges[0]].ToArray();
hdrSize = ParseSize(argsSpan[ranges[1]]);
totalSize = ParseSize(argsSpan[ranges[2]]);
}
else if (argCount == 4)
{
subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]);
reply = Encoding.ASCII.GetString(argsSpan[ranges[1]]);
subject = argsSpan[ranges[0]].ToArray();
reply = argsSpan[ranges[1]].ToArray();
hdrSize = ParseSize(argsSpan[ranges[2]]);
totalSize = ParseSize(argsSpan[ranges[3]]);
}
@@ -324,7 +333,7 @@ public sealed class NatsParser
return TryReadPayload(ref buffer, out command);
}
private bool TryReadPayload(ref ReadOnlySequence<byte> buffer, out ParsedCommand command)
private bool TryReadPayload(ref ReadOnlySequence<byte> buffer, out ParsedCommandView command)
{
command = default;
@@ -333,10 +342,7 @@ public sealed class NatsParser
if (buffer.Length < needed)
return false;
// Extract payload
var payloadSlice = buffer.Slice(0, _expectedPayloadSize);
var payload = new byte[_expectedPayloadSize];
payloadSlice.CopyTo(payload);
// Verify \r\n after payload
var trailer = buffer.Slice(_expectedPayloadSize, 2);
@@ -345,23 +351,25 @@ public sealed class NatsParser
if (trailerBytes[0] != (byte)'\r' || trailerBytes[1] != (byte)'\n')
throw new ProtocolViolationException("Expected \\r\\n after payload");
command = new ParsedCommand
command = new ParsedCommandView
{
Type = _pendingType,
Operation = _pendingOperation,
Subject = _pendingSubject,
ReplyTo = _pendingReplyTo,
Payload = payload,
Payload = payloadSlice,
HeaderSize = _pendingHeaderSize,
MaxMessages = -1,
};
buffer = buffer.Slice(buffer.GetPosition(needed));
_awaitingPayload = false;
_pendingSubject = null;
_pendingReplyTo = null;
return true;
}
private static ParsedCommand ParseSub(Span<byte> line)
private static ParsedCommandView ParseSub(Span<byte> line)
{
// SUB subject [queue] sid -- skip "SUB "
if (line.Length < 5)
@@ -372,28 +380,28 @@ public sealed class NatsParser
return argCount switch
{
2 => new ParsedCommand
2 => new ParsedCommandView
{
Type = CommandType.Sub,
Operation = "SUB",
Subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]),
Sid = Encoding.ASCII.GetString(argsSpan[ranges[1]]),
Subject = CopyBytes(argsSpan[ranges[0]]),
Sid = CopyBytes(argsSpan[ranges[1]]),
MaxMessages = -1,
},
3 => new ParsedCommand
3 => new ParsedCommandView
{
Type = CommandType.Sub,
Operation = "SUB",
Subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]),
Queue = Encoding.ASCII.GetString(argsSpan[ranges[1]]),
Sid = Encoding.ASCII.GetString(argsSpan[ranges[2]]),
Subject = CopyBytes(argsSpan[ranges[0]]),
Queue = CopyBytes(argsSpan[ranges[1]]),
Sid = CopyBytes(argsSpan[ranges[2]]),
MaxMessages = -1,
},
_ => throw new ProtocolViolationException("Invalid SUB arguments"),
};
}
private static ParsedCommand ParseUnsub(Span<byte> line)
private static ParsedCommandView ParseUnsub(Span<byte> line)
{
// UNSUB sid [max_msgs] -- skip "UNSUB "
if (line.Length < 7)
@@ -404,58 +412,59 @@ public sealed class NatsParser
return argCount switch
{
1 => new ParsedCommand
1 => new ParsedCommandView
{
Type = CommandType.Unsub,
Operation = "UNSUB",
Sid = Encoding.ASCII.GetString(argsSpan[ranges[0]]),
Sid = CopyBytes(argsSpan[ranges[0]]),
MaxMessages = -1,
},
2 => new ParsedCommand
2 => new ParsedCommandView
{
Type = CommandType.Unsub,
Operation = "UNSUB",
Sid = Encoding.ASCII.GetString(argsSpan[ranges[0]]),
Sid = CopyBytes(argsSpan[ranges[0]]),
MaxMessages = ParseSize(argsSpan[ranges[1]]),
},
_ => throw new ProtocolViolationException("Invalid UNSUB arguments"),
};
}
private static ParsedCommand ParseConnect(Span<byte> line)
private static ParsedCommandView ParseConnect(ReadOnlySequence<byte> line)
{
// CONNECT {json} -- find first space after command
int spaceIdx = line.IndexOf((byte)' ');
if (spaceIdx < 0)
var reader = new SequenceReader<byte>(line);
if (!reader.TryAdvanceTo((byte)' ', advancePastDelimiter: true))
throw new ProtocolViolationException("Invalid CONNECT");
var json = line[(spaceIdx + 1)..];
return new ParsedCommand
var json = line.Slice(reader.Position);
return new ParsedCommandView
{
Type = CommandType.Connect,
Operation = "CONNECT",
Payload = json.ToArray(),
Payload = json,
MaxMessages = -1,
};
}
private static ParsedCommand ParseInfo(Span<byte> line)
private static ParsedCommandView ParseInfo(ReadOnlySequence<byte> line)
{
// INFO {json} -- find first space after command
int spaceIdx = line.IndexOf((byte)' ');
if (spaceIdx < 0)
var reader = new SequenceReader<byte>(line);
if (!reader.TryAdvanceTo((byte)' ', advancePastDelimiter: true))
throw new ProtocolViolationException("Invalid INFO");
var json = line[(spaceIdx + 1)..];
return new ParsedCommand
var json = line.Slice(reader.Position);
return new ParsedCommandView
{
Type = CommandType.Info,
Operation = "INFO",
Payload = json.ToArray(),
Payload = json,
MaxMessages = -1,
};
}
private static ReadOnlyMemory<byte> CopyBytes(ReadOnlySpan<byte> value) =>
value.IsEmpty ? ReadOnlyMemory<byte>.Empty : value.ToArray();
/// <summary>
/// Parse a decimal integer from ASCII bytes. Returns -1 on error.
/// </summary>

View File

@@ -0,0 +1,42 @@
using System.Buffers;
using System.Text;
namespace NATS.Server.Protocol;
public readonly struct ParsedCommandView
{
public CommandType Type { get; init; }
public string? Operation { get; init; }
public ReadOnlyMemory<byte> Subject { get; init; }
public ReadOnlyMemory<byte> ReplyTo { get; init; }
public ReadOnlyMemory<byte> Queue { get; init; }
public ReadOnlyMemory<byte> Sid { get; init; }
public int MaxMessages { get; init; }
public int HeaderSize { get; init; }
public ReadOnlySequence<byte> Payload { get; init; }
public static ParsedCommandView Simple(CommandType type, string operation) =>
new() { Type = type, Operation = operation, MaxMessages = -1 };
public ReadOnlyMemory<byte> GetPayloadMemory() =>
Payload.IsEmpty ? ReadOnlyMemory<byte>.Empty
: Payload.IsSingleSegment ? Payload.First
: Payload.ToArray();
public ParsedCommand Materialize() =>
new()
{
Type = Type,
Operation = Operation,
Subject = DecodeAsciiOrNull(Subject),
ReplyTo = DecodeAsciiOrNull(ReplyTo),
Queue = DecodeAsciiOrNull(Queue),
Sid = DecodeAsciiOrNull(Sid),
MaxMessages = MaxMessages,
HeaderSize = HeaderSize,
Payload = GetPayloadMemory(),
};
private static string? DecodeAsciiOrNull(ReadOnlyMemory<byte> value) =>
value.IsEmpty ? null : Encoding.ASCII.GetString(value.Span);
}