233 lines
8.7 KiB
C#
233 lines
8.7 KiB
C#
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 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;
|
||
|
||
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)
|
||
{
|
||
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) =>
|
||
_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))
|
||
{
|
||
RpcMessage? response = await _router.DispatchAsync(_context, message, ct);
|
||
if (response is not null)
|
||
{
|
||
await SendAsync(response, ct);
|
||
}
|
||
}
|
||
}
|
||
catch (OperationCanceledException)
|
||
{
|
||
// 会话结束
|
||
}
|
||
}
|
||
|
||
private static RpcMessage? Parse(ReadOnlyMemory<byte> data)
|
||
{
|
||
try
|
||
{
|
||
return JsonSerializer.Deserialize(data.Span, FleetRpcContext.Default.RpcMessage);
|
||
}
|
||
catch (JsonException)
|
||
{
|
||
return null;
|
||
}
|
||
}
|
||
}
|