MoShiro/backend/agent/src/MoShiro.Agent/Chunks/ChunkAssembler.cs

using Fleet.Core;
using Fleet.Link;
using MoShiro.Measurement;
using MoShiro.Protocol;

namespace MoShiro.Agent;

/// <summary>
/// 拨号方 的列块到达之后：找到对应的流 → 与 Agent 在数据面记录的列合并 → 盖章（块头、时钟锚点）→ 编码 → 放进发件箱（通道 16）。
/// 找不到对应流的列块（任务已经被撤销、测试已经结束）被丢弃并计数；这只会发生在配置变更的瞬间。
/// </summary>
public sealed class ChunkAssembler
{
    private readonly StreamRegistry _streams;
    private readonly MonoClock _clock;
    private readonly AgentCounters _counters;
    private readonly Outbox _outbox;
    private readonly NodeId _agent;
    private readonly TimeProvider _time;
    private readonly INodeLog _log;
    private readonly TaskTrafficReporter _traffic;

    public ChunkAssembler(StreamRegistry streams, MonoClock clock, AgentCounters counters, Outbox outbox, NodeId agent, TimeProvider time, INodeLog log)
    {
        _streams = streams;
        _clock = clock;
        _counters = counters;
        _outbox = outbox;
        _agent = agent;
        _time = time;
        _log = log;
        _traffic = new TaskTrafficReporter(log, time);
    }

    /// <summary>
    /// 处理某个拨号方交来的一个列块（已经解码）。返回是否被接受。
    /// 块头里的 <c>dialer_id</c> 以认证过的会话为准（拨号方 自报的值不可信），别的拨号方的流也不会被命中。
    /// </summary>
    public bool Accept(RawChunk fragment, DialerSession dialer)
    {
        RawChunkHeader header = fragment.Header;
        StreamInfo? stream = header switch
        {
            { TestId: { } test, Phase: { } phase } => _streams.FindByTest(test, phase) is { } found && ReferenceEquals(found.Owner, dialer) ? found : null,
            { TaskId: { } task } => _streams.FindByTask(task, dialer),
            _ => null,
        };
        header.DialerId = new NodeId(dialer.DialerId);
        header.IpVersion = dialer.IpVersion; // 地址族以认证过的控制会话为准，不接受列块自报值
        KindSpec? spec = stream is null ? null : Kinds.Find(stream.Kind);
        if (stream is null || spec is null || spec.Id != header.Kind)
        {
            _counters.FragmentDiscarded();
            return false;
        }

        Result<IReadOnlyList<RawStream>> merged = FragmentMerger.Merge(spec, fragment, stream.Id, _streams);
        if (!merged.TryGet(out IReadOnlyList<RawStream>? streams, out FleetError? error))
        {
            _log.Warn("fragment discarded", ("kind", spec.Id), ("err", error.Code));
            _counters.FragmentDiscarded();
            stream.FragmentDone();
            return false;
        }

        RawChunk chunk = ChunkFinalizer.Finalize(spec, fragment, streams, _agent, _clock.Anchor(), stream.Params.Values, stream.Port, ChunkId.New(_time));
        Emit(chunk);
        _counters.FragmentAccepted();
        stream.FragmentDone();
        return true;
    }

    public Task RunTrafficReportingAsync(CancellationToken ct) => _traffic.RunAsync(ct);

    public void FlushTraffic() => _traffic.Flush();

    /// <summary>Agent 自己采集的块（目标任务）直接发出。</summary>
    public void Emit(RawChunk chunk)
    {
        byte[] payload = RawChunkCodec.Encode(chunk, out int uncompressedBytes);
        _outbox.Enqueue(MoShiroChannels.RawChunk, chunk.Header.ChunkId.Value, payload);
        _traffic.Record(
            chunk.Header.TaskId?.Value ?? chunk.Header.TestId?.Value ?? "unknown",
            chunk.Header.Kind,
            chunk.Streams.Sum(stream => (long)stream.Count),
            uncompressedBytes,
            payload.Length);
    }
}