MoShiro/backend/fleet/src/Fleet.Server.Link/NodeConnection.Receive.cs

using System.Text.Json;
using Fleet.Core;
using Fleet.Rpc;

namespace Fleet.Server;

/// <summary>接收：读帧、分发消息、hello、心跳、二进制通道（收件箱）。</summary>
internal sealed partial class NodeConnection
{
    private async Task ReceiveLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            using var idle = new CancellationTokenSource(_host.Options.HeartbeatTimeout, _host.Time);
            using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, idle.Token);
            Frame? frame;
            try
            {
                frame = await ReceiveFrameAsync(linked.Token);
            }
            catch (OperationCanceledException) when (idle.IsCancellationRequested && !ct.IsCancellationRequested)
            {
                _closeReason ??= "heartbeat timeout";
                return;
            }

            if (frame is null)
            {
                _closeReason ??= "closed by peer";
                return;
            }

            if (frame.IsText)
            {
                if (TryParse(frame.Data, out RpcMessage? message))
                {
                    await HandleMessageAsync(message, ct);
                }
            }
            else
            {
                await HandleBinaryAsync(frame.Data, ct);
            }
        }
    }

    private async Task HandleMessageAsync(RpcMessage message, CancellationToken ct)
    {
        if (message.IsResponse)
        {
            if (_pending.TryGetValue(message.Id!.Value, out TaskCompletionSource<RpcMessage>? waiter))
            {
                waiter.TrySetResult(message);
            }

            return;
        }

        switch (message.Method)
        {
            case FleetMethods.Hello:
                await HandleHelloAsync(message, ct);
                return;
            case FleetMethods.Heartbeat:
                await HandleHeartbeatAsync(message, ct);
                return;
            case FleetMethods.Stopping:
                try
                {
                    string? reason = message.Params?.Deserialize(FleetRpcContext.Default.StoppingParams)?.Reason;
                    _closeReason = reason is "updating" ? "updating" : "stopping";
                }
                catch (JsonException)
                {
                    _closeReason = "stopping";
                }
                return;
        }

        if (_node!.Status != NodeStatus.Active)
        {
            if (message.Id is not null)
            {
                await SendAsync(RpcMessage.Failure(message.Id, new FleetError(FleetErrors.NodePending, "node is not approved yet")), ct);
            }

            return;
        }

        var context = new NodeCallContext(_node, _session!, _host);
        RpcDispatch outcome;
        try
        {
            outcome = await _host.Router.DispatchCheckedAsync(context, message, ct);
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            throw;
        }

        if (outcome.Response is not null)
        {
            await SendAsync(outcome.Response, ct);
        }

        await NoteViolationAsync(outcome.Violation);
    }

    /// <summary>记录节点的违规；超过上限就断开这条连接。</summary>
    private async Task NoteViolationAsync(RpcViolation? violation)
    {
        if (violation is not { } v)
        {
            return;
        }

        FleetLog.ProtocolViolation(_host.Logger, _node!.Id.Value, v.Method, v.Reason);
        if (++_violations > _host.Options.MaxViolations)
        {
            await CloseAsync("protocol violations");
        }
    }

    private async Task HandleHelloAsync(RpcMessage message, CancellationToken ct)
    {
        HelloParams? hello = null;
        try
        {
            hello = message.Params?.Deserialize(FleetRpcContext.Default.HelloParams);
        }
        catch (JsonException)
        {
            // 按缺失处理
        }

        if (hello is null)
        {
            if (message.Id is not null)
            {
                await SendAsync(RpcMessage.Failure(message.Id, new FleetError(FleetErrors.InvalidParams, "invalid hello")), ct);
            }

            return;
        }

        long nowMs = _host.Time.GetUtcNow().ToUnixTimeMilliseconds();
        _session = _session! with
        {
            NodeVersion = hello.NodeVersion,
            Capabilities = hello.Capabilities,
            HelloMetaJson = hello.Meta?.GetRawText(),
        };
        _node = _node! with { NodeVersion = hello.NodeVersion };
        _helloDone = true;

        await UpdatePresenceAsync(online: true, ct);
        await FireAsync(NodeEventKind.Hello, ct, hello);

        if (message.Id is not null)
        {
            await SendAsync(
                RpcMessage.Success(
                    message.Id.Value,
                    new HelloResult { ServerTimeMs = nowMs, ClockOffsetMs = hello.NowUtcMs - nowMs },
                    FleetRpcContext.Default.HelloResult),
                ct);
        }

        SignalStateChanged();
    }

    private async Task HandleHeartbeatAsync(RpcMessage message, CancellationToken ct)
    {
        JsonElement? extra = null;
        try
        {
            extra = message.Params?.Deserialize(FleetRpcContext.Default.HeartbeatParams)?.Extra;
        }
        catch (JsonException)
        {
            // 忽略格式不正确的附加字段
        }

        long now = _host.Time.GetTimestamp();
        if (_host.Time.GetElapsedTime(Interlocked.Read(ref _lastPresenceTicks), now) >= s_presenceRefresh)
        {
            await UpdatePresenceAsync(online: true, ct);
        }

        await FireAsync(NodeEventKind.Heartbeat, ct, hello: null, extra);
        if (message.Id is not null)
        {
            await SendAsync(new RpcMessage { Id = message.Id, Result = JsonSerializer.SerializeToElement(new EmptyObject(), FleetEmptyObject.Info) }, ct);
        }
    }

    private async Task HandleBinaryAsync(byte[] data, CancellationToken ct)
    {
        if (_node!.Status != NodeStatus.Active || !OutboxFrame.TryParse(data, out byte channel, out Guid msgId, out ReadOnlyMemory<byte> payload))
        {
            return;
        }

        // 通道也要声明：节点只能往声明过、并且允许它的角色使用的通道发
        ChannelSpec? spec = _host.Router.Table.FindChannel(RpcPlane.Node, channel);
        if (spec is null || spec.From != RpcSide.Node || !spec.AllowsRole(_host.RoleOf?.Invoke(_node)))
        {
            await NoteViolationAsync(new RpcViolation($"channel {channel}", $"channel {channel} is not declared for this node"));
            return;
        }

        if (!_host.Inboxes.TryGetValue(channel, out IInbox? inbox))
        {
            return;
        }

        InboxOutcome outcome;
        try
        {
            outcome = await inbox.HandleAsync(_node, msgId, payload, ct);
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            throw;
        }
        catch (Exception ex)
        {
            // 没有持久化就不确认，节点会重发。
            FleetLog.InboxFailed(_host.Logger, ex, _node.Id.Value, channel);
            return;
        }

        if (outcome == InboxOutcome.Rejected)
        {
            return;
        }

        var ack = new OutboxAckParams();
        string id = msgId.ToString();
        if (outcome == InboxOutcome.Duplicate)
        {
            ack.Duplicates.Add(id);
        }
        else
        {
            ack.Ids.Add(id);
        }

        await SendAsync(RpcMessage.Notification(FleetMethods.OutboxAck, ack, FleetRpcContext.Default.OutboxAckParams), ct);
    }
}