MoShiro/backend/shared/MoShiro.Measurement/Merging/FragmentMerger.cs

using System.Text.Json;
using Fleet.Core;
using MoShiro.Protocol;

namespace MoShiro.Measurement;

/// <summary>Agent 一侧为合并保留的按键记录（由正在运行的积木提供）。</summary>
public interface IRecordSource
{
    RowStore? Records(uint streamId, string streamName);
}

/// <summary>
/// 把 Probe 一侧的列块（一个只带 Probe 列的 RawChunk）与 Agent 一侧的记录合并成完整的行。合并完全由 Kind 的 <see cref="StreamSpec"/> 声明驱动：
/// <list type="bullet">
/// <item><description>声明了连接键、两侧都有列：按键连接。主导方是 Probe（回显）时行取自列块，Agent 的列按键补上，缺失就是缺失（上行丢失）；
/// 主导方是 Agent（单向数据流）时行取自 Agent 的记录（按列块头里的序号范围），Probe 的列按键补上，缺失就是缺失（丢失）。</description></item>
/// <item><description>没有连接键，或者只有主导方一侧有列：原样透传，只按声明的列顺序重排。</description></item>
/// </list>
/// 合并之后的流按声明的列顺序排列，所以 Backend 的分析层不需要知道这些列是谁记录的；新增一种记录流不需要写合并代码。
/// </summary>
public static class FragmentMerger
{
    public static Result<IReadOnlyList<RawStream>> Merge(KindSpec spec, RawChunk fragment, uint streamId, IRecordSource records)
    {
        var streams = new List<RawStream>();
        foreach (RawStream stream in fragment.Streams)
        {
            StreamSpec? streamSpec = spec.FindStream(stream.Name);
            if (streamSpec is null)
            {
                return Result<IReadOnlyList<RawStream>>.Fail(MoShiroErrors.ValidationFailed, $"stream {stream.Name} is not declared by {spec.Id}");
            }

            Result<RawStream> merged = Join(streamSpec, stream, fragment.Header, records.Records(streamId, stream.Name));
            if (!merged.TryGet(out RawStream? value, out FleetError? error))
            {
                return Result<IReadOnlyList<RawStream>>.Fail(error);
            }

            streams.Add(value);
        }

        // 两侧各自记录完整行的流：Probe 的列块里没有它（例如 Probe 在不支持内核采样的平台上）时，Agent 自己的行也要交出去
        foreach (StreamSpec union in spec.Streams.Where(s => s.Union && fragment.Find(s.Name) is null))
        {
            RawStream? own = records.Records(streamId, union.Name)?.Drain(includeKey: false);
            if (own is { Count: > 0 })
            {
                streams.Add(Conform(union, own));
            }
        }

        return Result<IReadOnlyList<RawStream>>.Ok(streams);
    }

    private static Result<RawStream> Join(StreamSpec spec, RawStream probe, RawChunkHeader header, RowStore? agent)
    {
        if (spec.Union)
        {
            return Result<RawStream>.Ok(Concatenate(spec, probe, agent));
        }

        IReadOnlyList<ColumnSpec> probeColumns = spec.ColumnsOf(Side.Probe);
        IReadOnlyList<ColumnSpec> agentColumns = spec.ColumnsOf(Side.Agent);
        if (spec.Key is null || (spec.Driver == Side.Probe && agentColumns.Count == 0))
        {
            return Result<RawStream>.Ok(Conform(spec, probe));
        }

        RawColumn? key = probe.Find(spec.Key);
        if (key is null || probeColumns.Any(c => probe.Find(c.Name) is null))
        {
            return Result<RawStream>.Fail(MoShiroErrors.ValidationFailed, $"a fragment of stream {spec.Name} needs the columns {string.Join(", ", new[] { spec.Key }.Concat(probeColumns.Select(c => c.Name)))}");
        }

        return spec.Driver == Side.Probe ? JoinOnProbeRows(spec, probe, key, agent) : JoinOnAgentRows(spec, probe, key, header, agent);
    }

    private static Result<RawStream> JoinOnProbeRows(StreamSpec spec, RawStream probe, RawColumn key, RowStore? agent)
    {
        int n = probe.Count;
        var columns = new List<RawColumn>(spec.Columns.Count);
        foreach (ColumnSpec column in spec.Columns)
        {
            long[] values = new long[n];
            bool[] valid = new bool[n];
            if (column.Name == spec.Key)
            {
                key.Values.AsSpan(0, n).CopyTo(values);
                Array.Fill(valid, true);
            }
            else if (spec.OwnerOf(column) == Side.Probe)
            {
                RawColumn source = probe.Find(column.Name)!;
                for (int i = 0; i < n; i++)
                {
                    values[i] = source.Values[i];
                    valid[i] = source.IsValid(i);
                }
            }
            else if (agent is not null)
            {
                int index = agent.ColumnIndex(column.Name);
                for (int i = 0; i < n; i++)
                {
                    valid[i] = agent.TryGet(key.Values[i], index, out values[i]);
                }
            }

            columns.Add(new RawColumn(column.Name, column.Type, column.Enc, column.Nullable, column.Role, values, column.Nullable ? valid : null));
        }

        return Result<RawStream>.Ok(new RawStream(spec.Name, n, columns));
    }

    private static Result<RawStream> JoinOnAgentRows(StreamSpec spec, RawStream probe, RawColumn key, RawChunkHeader header, RowStore? agent)
    {
        if (!TryGetRange(header, out uint from, out uint to))
        {
            return Result<RawStream>.Fail(MoShiroErrors.ValidationFailed, $"a fragment of stream {spec.Name} must declare its sequence range");
        }

        // 行取自 Agent 的记录：范围内发出过的每一个包都有一行，Probe 没有报告到达的就是丢失
        RawStream rows = agent?.Cut(from, to, includeKey: true) ?? new RawStream(spec.Name, 0, []);
        var arrived = new Dictionary<long, int>();
        for (int i = 0; i < probe.Count; i++)
        {
            arrived.TryAdd(key.Values[i], i); // 重复的包只算第一次
        }

        int n = rows.Count;
        var columns = new List<RawColumn>(spec.Columns.Count);
        foreach (ColumnSpec column in spec.Columns)
        {
            long[] values = new long[n];
            bool[] valid = new bool[n];
            RawColumn? own = rows.Find(column.Name);
            if (own is not null)
            {
                own.Values.AsSpan(0, n).CopyTo(values);
                for (int i = 0; i < n; i++)
                {
                    valid[i] = own.IsValid(i);
                }
            }
            else if (probe.Find(column.Name) is { } source)
            {
                RawColumn keys = rows.Find(spec.Key!)!;
                for (int i = 0; i < n; i++)
                {
                    if (arrived.TryGetValue(keys.Values[i], out int row) && source.IsValid(row))
                    {
                        values[i] = source.Values[row];
                        valid[i] = true;
                    }
                }
            }

            columns.Add(new RawColumn(column.Name, column.Type, column.Enc, column.Nullable, column.Role, values, column.Nullable ? valid : null));
        }

        return Result<RawStream>.Ok(new RawStream(spec.Name, n, columns));
    }

    /// <summary>
    /// 两侧各自记录完整行的流（例如各自套接字的内核采样）：Probe 的行（列块里）加上 Agent 自己新记录的行，按声明的列顺序拼接。
    /// 每一行里有一列标明它属于哪一侧，所以不需要连接。
    /// </summary>
    private static RawStream Concatenate(StreamSpec spec, RawStream probe, RowStore? agent)
    {
        RawStream own = Conform(spec, probe);
        RawStream? extra = agent?.Drain(includeKey: false);
        if (extra is null || extra.Count == 0)
        {
            return own;
        }

        int n = own.Count + extra.Count;
        var columns = new List<RawColumn>(spec.Columns.Count);
        foreach (ColumnSpec column in spec.Columns)
        {
            RawColumn a = own.Find(column.Name)!;
            RawColumn b = extra.Find(column.Name)!;
            long[] values = new long[n];
            bool[] valid = new bool[n];
            for (int i = 0; i < own.Count; i++)
            {
                values[i] = a.Values[i];
                valid[i] = a.IsValid(i);
            }

            for (int i = 0; i < extra.Count; i++)
            {
                values[own.Count + i] = b.Values[i];
                valid[own.Count + i] = b.IsValid(i);
            }

            columns.Add(new RawColumn(column.Name, column.Type, column.Enc, column.Nullable, column.Role, values, column.Nullable ? valid : null));
        }

        return new RawStream(spec.Name, n, columns);
    }

    /// <summary>按声明的列顺序重排；Probe 一侧已经是完整的行时不改变内容。</summary>
    private static RawStream Conform(StreamSpec spec, RawStream stream)
    {
        var columns = new List<RawColumn>(spec.Columns.Count);
        foreach (ColumnSpec column in spec.Columns)
        {
            RawColumn? existing = stream.Find(column.Name);
            columns.Add(existing ?? new RawColumn(column.Name, column.Type, column.Enc, column.Nullable, column.Role, new long[stream.Count], column.Nullable ? new bool[stream.Count] : null));
        }

        return new RawStream(stream.Name, stream.Count, columns);
    }

    public static bool TryGetRange(RawChunkHeader header, out uint from, out uint to)
    {
        from = 0;
        to = 0;
        return header.Params.TryGetValue(FragmentKeys.SeqFrom, out JsonElement f) && f.TryGetUInt32(out from)
            && header.Params.TryGetValue(FragmentKeys.SeqTo, out JsonElement t) && t.TryGetUInt32(out to) && to >= from;
    }
}