MoShiro/backend/fleet/src/Fleet.Rpc/RpcSession.cs

using System.Collections.Concurrent;
using System.Text.Json;
using System.Text.Json.Serialization.Metadata;
using System.Threading.Channels;
using Fleet.Core;

namespace Fleet.Rpc;

public enum FrameKind : byte
{
    Text = 0,
    Binary = 1,
}

public readonly record struct RpcFrame(FrameKind Kind, ReadOnlyMemory<byte> Data);

/// <summary>
/// 一条双向的、保序的消息通道：文本帧承载 JSON-RPC，二进制帧承载发件箱。
/// 节点到服务端用 WebSocket，Probe 到 Agent 用 TLS 帧；会话引擎对两者一视同仁。
/// </summary>
public interface IRpcTransport : IAsyncDisposable
{
    ValueTask SendAsync(RpcFrame frame, CancellationToken ct);

    /// <summary>对端关闭时返回 null。</summary>
    ValueTask<RpcFrame?> ReceiveAsync(CancellationToken ct);

    ValueTask CloseAsync(string reason, CancellationToken ct);
}

public sealed class RpcSessionOptions
{
    public TimeSpan RequestTimeout { get; set; } = TimeSpan.FromSeconds(15);

    /// <summary>待处理的入站请求上限；超过时接收循环等待（背压）。</summary>
    public int InboundQueue { get; set; } = 256;

    /// <summary>对端违反协议（方法没有声明、方向、类型或角色不对）的次数上限；超过就断开会话。</summary>
    public int MaxViolations { get; set; } = 20;
}

/// <summary>
/// 传输无关的会话：请求与响应的关联、方法路由、二进制帧。
/// 入站的请求和通知由独立的任务按顺序处理，接收循环只负责读帧和完成响应，
/// 所以处理函数里可以等待向对端发起的请求而不会死锁。
/// </summary>
public sealed class RpcSession<TContext> : IAsyncDisposable
{
    private readonly IRpcTransport _transport;
    private readonly RpcRouter<TContext> _router;
    private readonly TContext _context;
    private readonly TimeProvider _time;
    private readonly RpcSessionOptions _options;
    private readonly ConcurrentDictionary<long, TaskCompletionSource<RpcMessage>> _pending = new();
    private readonly Channel<RpcMessage> _inbound;
    private readonly CancellationTokenSource _stop = new();
    private long _nextId;
    private long _lastReceive;
    private int _violations;
    private bool _violated;

    public RpcSession(IRpcTransport transport, RpcRouter<TContext> router, TContext context, TimeProvider time, RpcSessionOptions? options = null)
    {
        _transport = transport;
        _router = router;
        _context = context;
        _time = time;
        _lastReceive = time.GetTimestamp();
        _options = options ?? new RpcSessionOptions();
        _inbound = Channel.CreateBounded<RpcMessage>(new BoundedChannelOptions(_options.InboundQueue) { SingleReader = true, FullMode = BoundedChannelFullMode.Wait });
    }

    /// <summary>收到二进制帧：<c>u8 channel | 正文</c>（发件箱帧再往里解析）。</summary>
    public Func<ReadOnlyMemory<byte>, CancellationToken, ValueTask>? OnBinary { get; set; }

    /// <summary>对端上一次发来任何帧到现在的时间。静默死亡（断电、进程被冻结）的对端不会关闭连接，调用方据此设置空闲超时。</summary>
    public TimeSpan IdleFor => _time.GetElapsedTime(Volatile.Read(ref _lastReceive));

    /// <summary>读帧并分发，直到连接关闭或被取消。</summary>
    public async Task RunAsync(CancellationToken ct)
    {
        using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, _stop.Token);
        Task dispatcher = Task.Run(() => DispatchAsync(linked.Token), CancellationToken.None);
        try
        {
            while (!linked.IsCancellationRequested)
            {
                RpcFrame? frame = await _transport.ReceiveAsync(linked.Token);
                if (frame is null || Volatile.Read(ref _violated))
                {
                    return;
                }

                Volatile.Write(ref _lastReceive, _time.GetTimestamp());
                if (frame.Value.Kind == FrameKind.Binary)
                {
                    if (OnBinary is not null)
                    {
                        await OnBinary(frame.Value.Data, linked.Token);
                    }

                    continue;
                }

                RpcMessage? message = Parse(frame.Value.Data);
                if (message is null)
                {
                    continue; // 格式不对的帧忽略，连接保持可用
                }

                if (message.IsResponse)
                {
                    if (_pending.TryGetValue(message.Id!.Value, out TaskCompletionSource<RpcMessage>? waiter))
                    {
                        waiter.TrySetResult(message);
                    }
                }
                else
                {
                    await _inbound.Writer.WriteAsync(message, linked.Token);
                }
            }
        }
        finally
        {
            _inbound.Writer.TryComplete();
            foreach (TaskCompletionSource<RpcMessage> waiter in _pending.Values)
            {
                waiter.TrySetCanceled(CancellationToken.None);
            }

            await linked.CancelAsync();
            try
            {
                await dispatcher;
            }
            catch (OperationCanceledException)
            {
                // 正常结束
            }
        }
    }

    public async ValueTask<Result<JsonElement>> RequestAsync(string method, JsonElement? parameters, CancellationToken ct, TimeSpan? timeout = null)
    {
        long id = Interlocked.Increment(ref _nextId);
        var waiter = new TaskCompletionSource<RpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
        _pending[id] = waiter;
        try
        {
            await SendAsync(new RpcMessage { Id = id, Method = method, Params = parameters }, ct);
            using var timer = new CancellationTokenSource(timeout ?? _options.RequestTimeout, _time);
            using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, timer.Token);
            RpcMessage reply = await waiter.Task.WaitAsync(linked.Token);
            return reply.Error is not null
                ? Result<JsonElement>.Fail(reply.Error.ToFleetError())
                : Result<JsonElement>.Ok(reply.Result ?? default);
        }
        catch (OperationCanceledException) when (!ct.IsCancellationRequested)
        {
            return Result<JsonElement>.Fail(FleetErrors.Timeout, $"{method} timed out or the session closed");
        }
        finally
        {
            _pending.TryRemove(id, out _);
        }
    }

    public async ValueTask<Result<TResult>> RequestAsync<TParams, TResult>(
        string method, TParams parameters, JsonTypeInfo<TParams> paramsInfo, JsonTypeInfo<TResult> resultInfo, CancellationToken ct, TimeSpan? timeout = null)
    {
        Result<JsonElement> raw = await RequestAsync(method, JsonSerializer.SerializeToElement(parameters, paramsInfo), ct, timeout);
        if (!raw.TryGet(out JsonElement value, out FleetError? error))
        {
            return Result<TResult>.Fail(error);
        }

        try
        {
            TResult? result = value.Deserialize(resultInfo);
            return result is null ? Result<TResult>.Fail(FleetErrors.Internal, "empty result") : Result<TResult>.Ok(result);
        }
        catch (JsonException)
        {
            return Result<TResult>.Fail(FleetErrors.Internal, "result could not be parsed");
        }
    }

    public ValueTask NotifyAsync<TParams>(string method, TParams parameters, JsonTypeInfo<TParams> paramsInfo, CancellationToken ct) =>
        SendAsync(RpcMessage.Notification(method, parameters, paramsInfo), ct);

    public ValueTask SendAsync(RpcMessage message, CancellationToken ct)
    {
        if (message.Method is not null)
        {
            // 请求和通知只能发本侧有资格发起的方法（响应没有方法名，不检查）
            _router.CheckSend(message.Method, message.Id is null ? MethodKind.Notification : MethodKind.Request);
        }

        return _transport.SendAsync(new RpcFrame(FrameKind.Text, JsonSerializer.SerializeToUtf8Bytes(message, FleetRpcContext.Default.RpcMessage)), ct);
    }

    public ValueTask SendBinaryAsync(ReadOnlyMemory<byte> frame, CancellationToken ct) =>
        _transport.SendAsync(new RpcFrame(FrameKind.Binary, frame), ct);

    public ValueTask CloseAsync(string reason, CancellationToken ct) => _transport.CloseAsync(reason, ct);

    public async ValueTask DisposeAsync()
    {
        await _stop.CancelAsync();
        await _transport.DisposeAsync();
        _stop.Dispose();
    }

    private async Task DispatchAsync(CancellationToken ct)
    {
        try
        {
            await foreach (RpcMessage message in _inbound.Reader.ReadAllAsync(ct))
            {
                RpcDispatch outcome = await _router.DispatchCheckedAsync(_context, message, ct);
                if (outcome.Response is not null)
                {
                    await SendAsync(outcome.Response, ct);
                }

                if (outcome.Violation is not null && ++_violations > _options.MaxViolations)
                {
                    // 只关闭传输，不取消接收：接收循环在下一次读到帧或对端确认关闭时正常结束
                    Volatile.Write(ref _violated, true);
                    await _transport.CloseAsync("protocol violations", CancellationToken.None);
                    return;
                }
            }
        }
        catch (OperationCanceledException)
        {
            // 会话结束
        }
    }

    private static RpcMessage? Parse(ReadOnlyMemory<byte> data)
    {
        try
        {
            return JsonSerializer.Deserialize(data.Span, FleetRpcContext.Default.RpcMessage);
        }
        catch (JsonException)
        {
            return null;
        }
    }
}