PluralKit/Myriad/Gateway/ShardPacketSerializer.cs

using System.Buffers;
using System.Net.WebSockets;
using System.Text.Json;

namespace Myriad.Gateway;

public class ShardPacketSerializer
{
    private const int BufferSize = 64 * 1024;

    private readonly JsonSerializerOptions _jsonSerializerOptions;

    public ShardPacketSerializer(JsonSerializerOptions jsonSerializerOptions)
    {
        _jsonSerializerOptions = jsonSerializerOptions;
    }

    public async ValueTask<(WebSocketMessageType type, GatewayPacket? packet)> ReadPacket(ClientWebSocket socket)
    {
        using var buf = MemoryPool<byte>.Shared.Rent(BufferSize);

        var res = await socket.ReceiveAsync(buf.Memory, default);
        if (res.MessageType == WebSocketMessageType.Close)
            return (res.MessageType, null);

        if (res.EndOfMessage)
            // Entire packet fits within one buffer, deserialize directly
            return DeserializeSingleBuffer(buf, res);

        // Otherwise copy to stream buffer and deserialize from there
        return await DeserializeMultipleBuffer(socket, buf, res);
    }

    public async Task WritePacket(ClientWebSocket socket, GatewayPacket packet)
    {
        var bytes = JsonSerializer.SerializeToUtf8Bytes(packet, _jsonSerializerOptions);
        await socket.SendAsync(bytes.AsMemory(), WebSocketMessageType.Text, true, default);
    }

    private async Task<(WebSocketMessageType type, GatewayPacket packet)> DeserializeMultipleBuffer(
        ClientWebSocket socket, IMemoryOwner<byte> buf, ValueWebSocketReceiveResult res)
    {
        await using var stream = new MemoryStream(BufferSize * 4);
        stream.Write(buf.Memory.Span.Slice(0, res.Count));

        while (!res.EndOfMessage)
        {
            res = await socket.ReceiveAsync(buf.Memory, default);
            stream.Write(buf.Memory.Span.Slice(0, res.Count));
        }

        return DeserializeObject(res, stream.GetBuffer().AsSpan(0, (int)stream.Length));
    }

    private (WebSocketMessageType type, GatewayPacket packet) DeserializeSingleBuffer(
        IMemoryOwner<byte> buf, ValueWebSocketReceiveResult res)
    {
        var span = buf.Memory.Span.Slice(0, res.Count);
        return DeserializeObject(res, span);
    }

    private (WebSocketMessageType type, GatewayPacket packet) DeserializeObject(
        ValueWebSocketReceiveResult res, Span<byte> span)
    {
        var packet = JsonSerializer.Deserialize<GatewayPacket>(span, _jsonSerializerOptions)!;
        return (res.MessageType, packet);
    }
}