AyCode.Core/AyCode.Services/SignalRs/AcBinaryHubProtocol.cs

633 lines
20 KiB
C#

using System.Buffers;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.CompilerServices;
using System.Text;
using AyCode.Core.Serializers.Binaries;
using Microsoft.AspNetCore.Connections;
using Microsoft.AspNetCore.SignalR;
using Microsoft.AspNetCore.SignalR.Protocol;
namespace AyCode.Services.SignalRs;
/// <summary>
/// Custom SignalR hub protocol using AcBinarySerializer for wire format.
/// Eliminates JSON+Base64 overhead by serializing all HubMessages directly to binary.
///
/// Wire format per message:
/// [4 bytes: payload length (little-endian)] [payload bytes]
///
/// Payload structure:
/// [1 byte: message type] [message-specific fields serialized via AcBinary]
///
/// Message types map 1:1 to SignalR HubMessageType values.
/// Arguments are serialized individually with an INT32 length prefix each,
/// enabling deferred deserialization via IHubProtocol's binder pattern.
///
/// All writes go through BufferWriterBinaryOutput for zero virtual dispatch
/// on the hot path. Argument payloads are serialized directly to the pipe
/// via AcBinarySerializer (zero-copy). Length prefixes are patched in-place.
/// </summary>
public sealed class AcBinaryHubProtocol : IHubProtocol
{
private const int LengthPrefixSize = 4;
// Message type markers (matching HubMessageType enum values)
private const byte MsgInvocation = 1;
private const byte MsgStreamItem = 2;
private const byte MsgCompletion = 3;
private const byte MsgStreamInvocation = 4;
private const byte MsgCancelInvocation = 5;
private const byte MsgPing = 6;
private const byte MsgClose = 7;
private const byte MsgAck = 8;
private const byte MsgSequence = 9;
private volatile AcBinarySerializerOptions _options;
public AcBinaryHubProtocol() : this(AcBinarySerializerOptions.Default) { }
public AcBinaryHubProtocol(AcBinarySerializerOptions options)
{
_options = options;
}
/// <summary>
/// Runtime-replaceable serializer options.
/// Thread-safe: uses volatile field, callers see the new options on next message.
/// </summary>
public AcBinarySerializerOptions Options
{
get => _options;
set => _options = value;
}
public string Name => "acbinary";
public int Version => 1;
public TransferFormat TransferFormat => TransferFormat.Binary;
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public bool IsVersionSupported(int version) => version <= Version;
#region WriteMessage
public ReadOnlyMemory<byte> GetMessageBytes(HubMessage message)
{
// +LengthPrefixSize: prevents ArrayBufferWriter resize on first GetMemory,
// which would invalidate the length prefix span obtained before Advance.
var writer = new ArrayBufferWriter<byte>(_options.BufferWriterChunkSize + LengthPrefixSize);
WriteMessage(message, writer);
return writer.WrittenMemory;
}
public void WriteMessage(HubMessage message, IBufferWriter<byte> output)
{
// Reserve outer length prefix directly on the pipe (before BWO takes over)
var lengthSpan = output.GetSpan(LengthPrefixSize);
output.Advance(LengthPrefixSize);
var bw = new BufferWriterBinaryOutput(output, _options.BufferWriterChunkSize);
int externalBytes = 0;
switch (message)
{
case InvocationMessage m:
WriteInvocation(ref bw, output, m, ref externalBytes);
break;
case StreamInvocationMessage m:
WriteStreamInvocation(ref bw, output, m, ref externalBytes);
break;
case StreamItemMessage m:
WriteStreamItem(ref bw, output, m, ref externalBytes);
break;
case CompletionMessage m:
WriteCompletion(ref bw, output, m, ref externalBytes);
break;
case CancelInvocationMessage m:
WriteCancelInvocation(ref bw, m);
break;
case PingMessage:
bw.WriteByte(MsgPing);
break;
case CloseMessage m:
WriteClose(ref bw, m);
break;
case AckMessage m:
bw.WriteByte(MsgAck);
bw.WriteRaw(m.SequenceId);
break;
case SequenceMessage m:
bw.WriteByte(MsgSequence);
bw.WriteRaw(m.SequenceId);
break;
default:
throw new HubException($"Unexpected message type: {message.GetType().Name}");
}
var totalPayload = bw.Position + externalBytes;
bw.Flush();
Unsafe.WriteUnaligned(ref lengthSpan[0], totalPayload);
}
private void WriteInvocation(ref BufferWriterBinaryOutput bw, IBufferWriter<byte> output, InvocationMessage m, ref int externalBytes)
{
bw.WriteByte(MsgInvocation);
WriteNullableString(ref bw, m.InvocationId);
bw.WriteStringUtf8(m.Target);
WriteArguments(ref bw, output, m.Arguments, ref externalBytes);
WriteStringArray(ref bw, m.StreamIds);
WriteHeaders(ref bw, m.Headers);
}
private void WriteStreamInvocation(ref BufferWriterBinaryOutput bw, IBufferWriter<byte> output, StreamInvocationMessage m, ref int externalBytes)
{
bw.WriteByte(MsgStreamInvocation);
bw.WriteStringUtf8(m.InvocationId!);
bw.WriteStringUtf8(m.Target);
WriteArguments(ref bw, output, m.Arguments, ref externalBytes);
WriteStringArray(ref bw, m.StreamIds);
WriteHeaders(ref bw, m.Headers);
}
private void WriteStreamItem(ref BufferWriterBinaryOutput bw, IBufferWriter<byte> output, StreamItemMessage m, ref int externalBytes)
{
bw.WriteByte(MsgStreamItem);
bw.WriteStringUtf8(m.InvocationId!);
WriteArgument(ref bw, output, m.Item, ref externalBytes);
WriteHeaders(ref bw, m.Headers);
}
private void WriteCompletion(ref BufferWriterBinaryOutput bw, IBufferWriter<byte> output, CompletionMessage m, ref int externalBytes)
{
bw.WriteByte(MsgCompletion);
bw.WriteStringUtf8(m.InvocationId!);
WriteNullableString(ref bw, m.Error);
var hasResult = m.HasResult;
bw.WriteByte(hasResult ? (byte)1 : (byte)0);
if (hasResult)
WriteArgument(ref bw, output, m.Result, ref externalBytes);
WriteHeaders(ref bw, m.Headers);
}
private static void WriteCancelInvocation(ref BufferWriterBinaryOutput bw, CancelInvocationMessage m)
{
bw.WriteByte(MsgCancelInvocation);
bw.WriteStringUtf8(m.InvocationId!);
WriteHeaders(ref bw, m.Headers);
}
private static void WriteClose(ref BufferWriterBinaryOutput bw, CloseMessage m)
{
bw.WriteByte(MsgClose);
WriteNullableString(ref bw, m.Error);
bw.WriteByte(m.AllowReconnect ? (byte)1 : (byte)0);
}
#endregion
#region TryParseMessage
public bool TryParseMessage(ref ReadOnlySequence<byte> input, IInvocationBinder binder, [NotNullWhen(true)] out HubMessage? message)
{
message = null;
if (input.Length < LengthPrefixSize)
return false;
int payloadLength;
if (input.FirstSpan.Length >= LengthPrefixSize)
{
payloadLength = Unsafe.ReadUnaligned<int>(ref Unsafe.AsRef(in input.FirstSpan[0]));
}
else
{
Span<byte> lenBuf = stackalloc byte[LengthPrefixSize];
input.Slice(0, LengthPrefixSize).CopyTo(lenBuf);
payloadLength = Unsafe.ReadUnaligned<int>(ref lenBuf[0]);
}
var totalLength = LengthPrefixSize + payloadLength;
if (input.Length < totalLength)
return false;
var payload = input.Slice(LengthPrefixSize, payloadLength);
ReadOnlySpan<byte> span;
byte[]? rentedBuffer = null;
if (payload.IsSingleSegment)
{
span = payload.FirstSpan;
}
else
{
rentedBuffer = ArrayPool<byte>.Shared.Rent(payloadLength);
payload.CopyTo(rentedBuffer);
span = rentedBuffer.AsSpan(0, payloadLength);
}
try
{
message = ParseMessage(span, binder);
}
finally
{
if (rentedBuffer != null)
ArrayPool<byte>.Shared.Return(rentedBuffer);
}
input = input.Slice(totalLength);
return message != null;
}
private HubMessage? ParseMessage(ReadOnlySpan<byte> span, IInvocationBinder binder)
{
if (span.Length == 0)
return null;
var reader = new SpanReader(span);
var msgType = reader.ReadByte();
return msgType switch
{
MsgInvocation => ParseInvocation(ref reader, binder),
MsgStreamInvocation => ParseStreamInvocation(ref reader, binder),
MsgStreamItem => ParseStreamItem(ref reader, binder),
MsgCompletion => ParseCompletion(ref reader, binder),
MsgCancelInvocation => ParseCancelInvocation(ref reader),
MsgPing => PingMessage.Instance,
MsgClose => ParseClose(ref reader),
MsgAck => new AckMessage(reader.ReadInt64()),
MsgSequence => new SequenceMessage(reader.ReadInt64()),
_ => null
};
}
private HubMessage ParseInvocation(ref SpanReader r, IInvocationBinder binder)
{
var invocationId = r.ReadNullableString();
var target = r.ReadString();
var paramTypes = binder.GetParameterTypes(target);
var args = ReadArguments(ref r, paramTypes);
var streamIds = r.ReadStringArray();
var headers = ReadHeaders(ref r);
var msg = streamIds is { Length: > 0 }
? new InvocationMessage(invocationId, target, args, streamIds)
: ApplyInvocationId(new InvocationMessage(target, args), invocationId);
if (headers != null)
SetHeaders(msg, headers);
return msg;
}
private HubMessage ParseStreamInvocation(ref SpanReader r, IInvocationBinder binder)
{
var invocationId = r.ReadString();
var target = r.ReadString();
var paramTypes = binder.GetParameterTypes(target);
var args = ReadArguments(ref r, paramTypes);
var streamIds = r.ReadStringArray();
var headers = ReadHeaders(ref r);
var msg = new StreamInvocationMessage(invocationId, target, args, streamIds);
if (headers != null)
SetHeaders(msg, headers);
return msg;
}
private HubMessage ParseStreamItem(ref SpanReader r, IInvocationBinder binder)
{
var invocationId = r.ReadString();
var itemType = binder.GetStreamItemType(invocationId);
var item = ReadSingleArgument(ref r, itemType);
var headers = ReadHeaders(ref r);
var msg = new StreamItemMessage(invocationId, item);
if (headers != null)
SetHeaders(msg, headers);
return msg;
}
private HubMessage ParseCompletion(ref SpanReader r, IInvocationBinder binder)
{
var invocationId = r.ReadString();
var error = r.ReadNullableString();
var hasResult = r.ReadByte() == 1;
object? result = null;
if (hasResult)
{
var resultType = binder.GetReturnType(invocationId);
result = ReadSingleArgument(ref r, resultType);
}
var headers = ReadHeaders(ref r);
CompletionMessage msg;
if (error != null)
msg = CompletionMessage.WithError(invocationId, error);
else if (hasResult)
msg = CompletionMessage.WithResult(invocationId, result);
else
msg = CompletionMessage.Empty(invocationId);
if (headers != null)
SetHeaders(msg, headers);
return msg;
}
private static HubMessage ParseCancelInvocation(ref SpanReader r)
{
var invocationId = r.ReadString();
var headers = ReadHeaders(ref r);
var msg = new CancelInvocationMessage(invocationId);
if (headers != null)
SetHeaders(msg, headers);
return msg;
}
private static HubMessage ParseClose(ref SpanReader r)
{
var error = r.ReadNullableString();
var allowReconnect = r.Remaining > 0 && r.ReadByte() == 1;
return new CloseMessage(error, allowReconnect);
}
#endregion
#region Argument Serialization
private void WriteArguments(ref BufferWriterBinaryOutput bw, IBufferWriter<byte> output, object?[] arguments, ref int externalBytes)
{
bw.WriteVarUInt((uint)arguments.Length);
for (var i = 0; i < arguments.Length; i++)
WriteArgument(ref bw, output, arguments[i], ref externalBytes);
}
private void WriteArgument(ref BufferWriterBinaryOutput bw, IBufferWriter<byte> output, object? value, ref int externalBytes)
{
if (value is byte[] byteArray)
{
// byte[] fast-path: size known upfront, write entirely through BWO
var argPayload = 1 + VarUIntSize((uint)byteArray.Length) + byteArray.Length;
bw.WriteRaw(argPayload);
bw.WriteByte(BinaryTypeCode.ByteArray);
bw.WriteVarUInt((uint)byteArray.Length);
bw.WriteBytes(byteArray);
return;
}
// Flush BWO to pipe, then serialize directly to the pipe via AcBinarySerializer
bw.FlushAndReset();
// Reserve arg length prefix directly on the pipe
var argLenSpan = output.GetSpan(LengthPrefixSize);
output.Advance(LengthPrefixSize);
var argBytes = AcBinarySerializer.Serialize(value, output, _options);
Unsafe.WriteUnaligned(ref argLenSpan[0], argBytes);
externalBytes += LengthPrefixSize + argBytes;
}
private object?[] ReadArguments(ref SpanReader r, IReadOnlyList<Type> paramTypes)
{
var count = (int)r.ReadVarUInt();
var args = new object?[count];
for (var i = 0; i < count; i++)
{
var targetType = i < paramTypes.Count ? paramTypes[i] : typeof(object);
args[i] = ReadSingleArgument(ref r, targetType);
}
return args;
}
private object? ReadSingleArgument(ref SpanReader r, Type targetType)
{
var argLength = r.ReadInt32();
if (argLength == 0)
return null;
var argSpan = r.ReadSpan(argLength);
if (argLength == 1 && argSpan[0] == 0)
return null;
// byte[] fast-path: bypass deserializer engine
if (targetType == typeof(byte[]) && argSpan.Length > 0 && argSpan[0] == BinaryTypeCode.ByteArray)
{
var byteReader = new SpanReader(argSpan.Slice(1));
var len = (int)byteReader.ReadVarUInt();
return byteReader.ReadSpan(len).ToArray();
}
return AcBinaryDeserializer.Deserialize(argSpan, targetType, _options);
}
#endregion
#region Framing Helpers
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static void WriteNullableString(ref BufferWriterBinaryOutput bw, string? value)
{
if (value == null)
{
bw.WriteByte(0);
return;
}
bw.WriteByte(1);
bw.WriteStringUtf8(value);
}
private static void WriteStringArray(ref BufferWriterBinaryOutput bw, string[]? array)
{
if (array == null || array.Length == 0)
{
bw.WriteVarUInt(0);
return;
}
bw.WriteVarUInt((uint)array.Length);
for (var i = 0; i < array.Length; i++)
bw.WriteStringUtf8(array[i]);
}
private static void WriteHeaders(ref BufferWriterBinaryOutput bw, IDictionary<string, string>? headers)
{
if (headers == null || headers.Count == 0)
{
bw.WriteVarUInt(0);
return;
}
bw.WriteVarUInt((uint)headers.Count);
foreach (var kv in headers)
{
bw.WriteStringUtf8(kv.Key);
bw.WriteStringUtf8(kv.Value);
}
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static int VarUIntSize(uint value)
{
if (value < 0x80) return 1;
if (value < 0x4000) return 2;
if (value < 0x200000) return 3;
if (value < 0x10000000) return 4;
return 5;
}
#endregion
#region Helpers
private static InvocationMessage ApplyInvocationId(InvocationMessage msg, string? invocationId)
{
if (invocationId != null)
return new InvocationMessage(invocationId, msg.Target, msg.Arguments);
return msg;
}
private static void SetHeaders(HubMessage msg, Dictionary<string, string> headers)
{
if (msg is HubInvocationMessage invMsg)
invMsg.Headers = headers;
}
private static Dictionary<string, string>? ReadHeaders(ref SpanReader r)
{
if (r.Remaining == 0)
return null;
var count = (int)r.ReadVarUInt();
if (count == 0)
return null;
var headers = new Dictionary<string, string>(count, StringComparer.Ordinal);
for (var i = 0; i < count; i++)
{
var key = r.ReadString();
var value = r.ReadString();
headers[key] = value;
}
return headers;
}
#endregion
#region SpanReader
/// <summary>
/// Lightweight ref struct for sequential reading from a ReadOnlySpan.
/// </summary>
private ref struct SpanReader
{
private readonly ReadOnlySpan<byte> _span;
private int _pos;
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public SpanReader(ReadOnlySpan<byte> span)
{
_span = span;
_pos = 0;
}
public int Remaining
{
[MethodImpl(MethodImplOptions.AggressiveInlining)]
get => _span.Length - _pos;
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public byte ReadByte() => _span[_pos++];
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public int ReadInt32()
{
var value = Unsafe.ReadUnaligned<int>(ref Unsafe.AsRef(in _span[_pos]));
_pos += 4;
return value;
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public long ReadInt64()
{
var value = Unsafe.ReadUnaligned<long>(ref Unsafe.AsRef(in _span[_pos]));
_pos += 8;
return value;
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public uint ReadVarUInt()
{
uint value = 0;
var shift = 0;
while (true)
{
var b = _span[_pos++];
value |= (uint)(b & 0x7F) << shift;
if ((b & 0x80) == 0)
return value;
shift += 7;
}
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public ReadOnlySpan<byte> ReadSpan(int length)
{
var result = _span.Slice(_pos, length);
_pos += length;
return result;
}
public string ReadString()
{
var byteCount = (int)ReadVarUInt();
if (byteCount == 0)
return string.Empty;
var bytes = ReadSpan(byteCount);
return Encoding.UTF8.GetString(bytes);
}
public string? ReadNullableString()
{
var marker = ReadByte();
return marker == 0 ? null : ReadString();
}
public string[]? ReadStringArray()
{
var count = (int)ReadVarUInt();
if (count == 0)
return null;
var array = new string[count];
for (var i = 0; i < count; i++)
array[i] = ReadString();
return array;
}
}
#endregion
}