MoShiro/backend/tests/MoShiro.TestSupport/ChunkFactory.cs

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

namespace MoShiro.TestSupport;

/// <summary>每一个回显包的设定：上行或下行丢失，以及两个方向的单向时延（微秒）。</summary>
public readonly record struct EchoPacket(bool LostUp, bool LostDown, long OwdUpUs, long OwdDownUs, long ProcessUs = 50);

/// <summary>合成原始块：用于测试、模拟器和离线工具。块头里的时钟锚点保证 Probe 与 Agent 时钟相差 <see cref="ProbeOffsetUs"/>。</summary>
public static class ChunkFactory
{
    /// <summary>合成块默认属于哪个 Probe（测试夹具在后端创建一个真实的 Probe 之后把它的 id 写在这里）。</summary>
    public static string DefaultProbe { get; set; } = "prb_test";

    public const long AgentMonoAnchorUs = 1_000_000_000;
    public const long ProbeOffsetUs = 5_000_000;

    public static RawChunkHeader Header(NodeId agent, string kind, long startUtcMs, int durationMs = 10_000, IReadOnlyDictionary<string, JsonElement>? parameters = null)
    {
        KindSpec spec = Kinds.Find(kind)!;
        ParamResolution resolved = ParamResolver.Resolve(spec.Params, parameters);
        return new RawChunkHeader
        {
            ChunkId = new ChunkId(Guid.CreateVersion7()),
            Kind = kind,
            AgentId = agent,
            ProbeId = new ProbeId(DefaultProbe),
            Params = new Dictionary<string, JsonElement>(resolved.Values!.Values),
            Port = 7443,
            StartUtcMs = startUtcMs,
            DurationMs = durationMs,
            Clock = new ChunkClock { AgentMonoUs = AgentMonoAnchorUs, UtcMs = startUtcMs, ProbeOffsetUs = ProbeOffsetUs },
        };
    }

    /// <summary>Probe 时钟下，距离块起点 <paramref name="offsetMs"/> 毫秒的时刻（微秒）。</summary>
    public static long ProbeUs(long offsetMs) => (offsetMs * 1000) + AgentMonoAnchorUs - ProbeOffsetUs;

    public static long AgentUs(long offsetMs) => (offsetMs * 1000) + AgentMonoAnchorUs;

    public static RawChunk Echo(
        NodeId agent,
        string kind,
        long startUtcMs,
        int count,
        int intervalMs,
        Func<int, EchoPacket> packet,
        int durationMs = 10_000,
        IReadOnlyDictionary<string, JsonElement>? parameters = null,
        IEnumerable<RawStream>? extraStreams = null)
    {
        var msg = new StreamBuilder(Kinds.Find(kind)!.FindStream("msg")!);
        for (int i = 0; i < count; i++)
        {
            EchoPacket p = packet(i);
            long t1 = ProbeUs((long)i * intervalMs);
            long t2 = t1 + ProbeOffsetUs + p.OwdUpUs;
            long t3 = t2 + p.ProcessUs;
            long t4 = t3 - ProbeOffsetUs + p.OwdDownUs;
            msg.Add(i, 64, t1, p.LostUp ? null : t2, p.LostUp ? null : t3, p.LostUp || p.LostDown ? null : t4);
        }

        var streams = new List<RawStream> { msg.Build() };
        if (extraStreams is not null)
        {
            streams.AddRange(extraStreams);
        }

        return new RawChunk(Header(agent, kind, startUtcMs, durationMs, parameters), streams);
    }

    /// <summary>10 Hz 的 udp_echo，没有丢包，RTT = 2 × owdMs。</summary>
    public static RawChunk UdpEchoClean(NodeId agent, long startUtcMs, int count = 100, double owdMs = 20) =>
        Echo(agent, "udp_echo", startUtcMs, count, 100, _ => new EchoPacket(false, false, (long)(owdMs * 1000), (long)(owdMs * 1000)));

    public static RawStream TcpInfo(IEnumerable<(long TsUs, byte Side, uint SrttUs, uint RetransTotal, uint SegsOut)> samples)
    {
        var builder = new StreamBuilder(Kinds.TcpEcho.FindStream("tcp_info")!);
        foreach ((long ts, byte side, uint srtt, uint retrans, uint segs) in samples)
        {
            builder.Add(ts, side, srtt, srtt / 4, 10, retrans, segs, segs * 100L);
        }

        return builder.Build();
    }

    public static RawChunk Connect(NodeId agent, long startUtcMs, IEnumerable<(long TcpUs, long? TlsUs, ConnectOutcomeCode Outcome)> attempts)
    {
        var builder = new StreamBuilder(Kinds.Connect.FindStream("conn")!);
        int i = 0;
        foreach ((long tcp, long? tls, ConnectOutcomeCode outcome) in attempts)
        {
            builder.Add(i, ProbeUs(i * 5000L), outcome == ConnectOutcomeCode.Timeout ? null : tcp, tls, (long)outcome);
            i++;
        }

        return new RawChunk(Header(agent, "connect", startUtcMs), [builder.Build()]);
    }

    /// <summary>udp_train：Agent 以固定间隔发包，Probe 记录收到的时刻；<paramref name="lossEvery"/> 为 N 时每 N 个包丢一个。</summary>
    public static RawChunk Train(NodeId agent, long startUtcMs, double mbps, int phase, int seconds = 10, int lossEvery = 0, long owdUs = 20_000)
    {
        const int size = 1250;
        int packets = (int)(mbps * 1_000_000 / 8 / size * seconds);
        double intervalUs = seconds * 1_000_000.0 / packets;
        var pkt = new StreamBuilder(Kinds.UdpTrain.FindStream("pkt")!);
        for (int i = 0; i < packets; i++)
        {
            long sendAgent = AgentMonoAnchorUs + (long)(i * intervalUs);
            bool lost = lossEvery > 0 && i % lossEvery == 0;
            long recvProbe = sendAgent + owdUs - ProbeOffsetUs;
            pkt.Add(i, size, sendAgent, lost ? null : recvProbe);
        }

        var parameters = new Dictionary<string, JsonElement>
        {
            ["mbps"] = JsonSerializer.SerializeToElement(mbps, FactoryJson.Default.Double),
            ["duration_s"] = JsonSerializer.SerializeToElement((double)seconds, FactoryJson.Default.Double),
        };
        RawChunkHeader header = Header(agent, "udp_train", startUtcMs, seconds * 1000, parameters);
        header.Phase = phase;
        header.TestId = new TestId("tst_test");
        return new RawChunk(header, [pkt.Build(), new StreamBuilder(Kinds.UdpTrain.FindStream("step")!).Build()]);
    }
}

public enum ConnectOutcomeCode
{
    Ok = 0,
    Timeout = 1,
    Reset = 2,
    TlsAlert = 3,
}

[System.Text.Json.Serialization.JsonSerializable(typeof(double))]
internal sealed partial class FactoryJson : System.Text.Json.Serialization.JsonSerializerContext;