MoShiro/backend/shared/MoShiro.Measurement/Blocks/Echo.cs

using System.Net;
using System.Net.Sockets;
using Fleet.Core;
using MoShiro.Protocol;

namespace MoShiro.Measurement;

/// <summary>
/// 回显积木的行表（拨号方，行的主导方）：发出时记 size、t1，回复到达时按序号回填 t4，到点把已经稳定的行切成列块。
/// UDP 回显和 TCP 回显共用，列由 Kind 的 <c>msg</c> 流声明决定。
/// </summary>
internal sealed class EchoRows
{
    /// <summary>UDP：没有在这段时间内回来的包就是丢了。</summary>
    public const long DatagramGraceUs = 2_000_000;

    /// <summary>
    /// TCP 不会丢包，只会延迟（重传、队头阻塞）：一次几秒的中断之后，积压的回显会陆续回来，RTT 就是那几秒。
    /// 如果像 UDP 那样等 2 秒就把没有回复的行当作丢失交出去，最严重的长尾反而成了"丢包"，RTT 里一点也看不到。
    /// 所以 TCP 的行要等回复（或者连接真的断了——30 秒之后仍然没有回复才交出，t4 缺失）。
    /// </summary>
    public const long TcpGraceUs = 30_000_000;

    private readonly DriverLog _log;
    private readonly int _size;
    private readonly int _w;
    private readonly int _t1;
    private readonly int _t4;
    private readonly long _graceUs;
    private readonly bool _settleOnReply;
    private long _discarded;

    public EchoRows(KindSpec kind, int capacity, long graceUs, bool settleOnReply)
    {
        _graceUs = graceUs;
        _settleOnReply = settleOnReply;
        StreamSpec spec = kind.FindStream("msg")!;
        _log = new DriverLog(spec, Side.Dialer, capacity);
        _size = _log.Store.ColumnIndex("size");
        _w = _log.Store.ColumnIndex("w");
        _t1 = _log.Store.ColumnIndex("t1");
        _t4 = _log.Store.ColumnIndex("t4");
    }

    public long Count => _log.Count - Interlocked.Read(ref _discarded);

    /// <summary><paramref name="weight"/>：这个包代表多少个基础网格间隔（频率变化时分析按它累计，丢包率和分位数仍是按时间平均的）。</summary>
    public uint Add(int size, long t1Us, int weight = 1)
    {
        long key = _log.Store.Append();
        _log.Store.Set(key, _size, size);
        _log.Store.Set(key, _w, weight);
        _log.Store.Set(key, _t1, t1Us);
        return (uint)key;
    }

    public bool Reply(uint seq, long t4Us) => Reply(seq, t4Us, out _);

    /// <summary>回填 t4；第一次回复才算，<paramref name="t1Us"/> 是这个包的发出时间。</summary>
    public bool Reply(uint seq, long t4Us, out long t1Us)
    {
        t1Us = 0;
        if (_log.Store.Contains(seq) && !_log.Store.IsSet(seq, _t4) && _log.Store.TryGet(seq, _t1, out t1Us))
        {
            _log.Store.Set(seq, _t4, t4Us);
            return true;
        }

        return false;
    }

    public bool Discard(uint seq)
    {
        if (!_log.Store.Remove(seq))
        {
            return false;
        }

        Interlocked.Increment(ref _discarded);
        return true;
    }

    public RawStream? Take(long nowUs, bool all) =>
        _log.Take(
            key => (_settleOnReply && _log.Store.IsSet(key, _t4)) || (_log.Store.TryGet(key, _t1, out long t1) && t1 <= nowUs - _graceUs),
            all);
}

/// <summary>
/// UDP 回显，拨号方：按声明的速率发包，t1 在发出时记录，回复到达时回填 t4。回复丢失的行 t4 缺失（Agent 一侧再补 t2/t3）。
/// </summary>
public sealed class DatagramEchoDialer : BlockBase
{
    private readonly EchoRows _rows;
    private readonly object _gate = new();
    private long _nextSendFailureWarningUs;
    private AdaptiveRate? _rate;

    public DatagramEchoDialer(BlockEnv env, BlockSpec spec)
        : base(env, spec)
    {
        _rows = new EchoRows(spec.Kind, 16384, EchoRows.DatagramGraceUs, settleOnReply: false);
    }

    public override long Produced => _rows.Count;

    public override async Task RunAsync(CancellationToken ct)
    {
        QuicDatagramTransport datagrams = Env.Datagrams!;
        DataSession session = Env.Session() ?? throw new BlockFaultException("no data session");
        await datagrams.WaitUntilConnectedAsync(ct);
        IPEndPoint target = PeerEndpoint(Spec.Port);
        int size = (int)Spec.Int("size");
        byte[] packet = new byte[size];
        _rate = CreateAdaptiveRate();

        void SendOne()
        {
            long t1 = Env.Clock.NowUs();
            lock (_gate)
            {
                uint seq = _rows.Add(size, t1, _rate?.Weight ?? 1);
                MeasurementDatagram.Write(packet, new MeasurementDatagramHeader(MeasurementDatagramType.EchoRequest, session.SessionId, Spec.StreamId, seq, t1), session.Key);
                if (datagrams.Send(session, target, packet))
                {
                    _rate?.OnSend(t1);
                    return;
                }

                _rows.Discard(seq);
                if (t1 >= _nextSendFailureWarningUs)
                {
                    Env.Log.Warn("QUIC echo datagram was not queued; omitted from loss samples",
                        ("stream", Spec.StreamId), ("seq", seq), ("bytes", size));
                    _nextSendFailureWarningUs = t1 + 5_000_000;
                }
            }
        }

        datagrams.Register(Spec.StreamId, OnPacket);
        try
        {
            if (_rate is null)
            {
                using var timer = new PeriodicTimer(TimeSpan.FromSeconds(1.0 / Spec.Int("rate_hz")), Env.Time);
                while (await timer.WaitForNextTickAsync(ct))
                {
                    SendOne();
                }
            }
            else
            {
                await PaceAsync(_rate, () =>
                {
                    SendOne();
                    return ValueTask.CompletedTask;
                }, ct);
            }
        }
        catch (OperationCanceledException)
        {
            // 积木被停止
        }
        finally
        {
            datagrams.Unregister(Spec.StreamId);
        }
    }

    public override int Flush(FlushMode mode, long peerProduced = 0)
    {
        RawStream? stream;
        lock (_gate)
        {
            stream = _rows.Take(Env.Clock.NowUs(), mode == FlushMode.All);
        }

        if (stream is null)
        {
            return 0;
        }

        Ship(stream, "t1");
        return 1;
    }

    private void OnPacket(in QuicDatagramDelivery delivery)
    {
        if (delivery.Header.Type == MeasurementDatagramType.EchoReply)
        {
            lock (_gate)
            {
                if (_rows.Reply(delivery.Header.Seq, delivery.ReceiveUs, out long t1))
                {
                    _rate?.OnReply(t1, delivery.ReceiveUs, (delivery.ReceiveUs - t1) / 1000.0);
                }
            }
        }
    }
}

/// <summary>UDP 回显，监听方：校验后立即回复一个等长的包，记录 t2（到达）和 t3（回复发出），按序号留给合并。</summary>
public sealed class DatagramEchoListener : BlockBase
{
    private readonly RowStore _records;

    public DatagramEchoListener(BlockEnv env, BlockSpec spec)
        : base(env, spec)
    {
        _records = new RowStore(spec.Kind.FindStream("msg")!, Side.Listener, 32768);
    }

    public override RowStore? Records(string streamName) => streamName == "msg" ? _records : null;

    public override async Task RunAsync(CancellationToken ct)
    {
        Env.Datagrams!.Register(Spec.StreamId, OnPacket);
        try
        {
            await WaitUntilCancelledAsync(ct);
        }
        finally
        {
            Env.Datagrams.Unregister(Spec.StreamId);
        }
    }

    private void OnPacket(in QuicDatagramDelivery delivery)
    {
        if (delivery.Header.Type != MeasurementDatagramType.EchoRequest)
        {
            return; // Data：探活包，只用来保持 NAT 映射和让准入策略学到源地址
        }

        long t2 = delivery.ReceiveUs;
        delivery.ReplyAs(MeasurementDatagramType.EchoReply);
        Span<long> row = [t2, Env.Clock.NowUs()];
        _records.SetRow(delivery.Header.Seq, row);
    }
}

/// <summary>
/// TCP 回显，拨号方：一条常驻连接（可选 TLS）上按速率发小消息，回复按序到达。连接断开期间不发消息（没有连接不是丢包），
/// 之后按退避重连；断开和重连本身由 connect/hold 记录。
/// </summary>
public sealed class TcpEchoDialer : BlockBase
{
    private readonly EchoRows _rows;
    private readonly TcpInfoRecorder _info;
    private AdaptiveRate? _rate;

    public TcpEchoDialer(BlockEnv env, BlockSpec spec)
        : base(env, spec)
    {
        _rows = new EchoRows(spec.Kind, 16384, EchoRows.TcpGraceUs, settleOnReply: true);
        _info = new TcpInfoRecorder(spec.Kind.FindStream("tcp_info")!, Side.Dialer, env.Clock);
    }

    public override long Produced => _rows.Count;

    public override async Task RunAsync(CancellationToken ct)
    {
        var backoff = new Fleet.Core.Backoff(TimeSpan.FromSeconds(1), TimeSpan.FromSeconds(30));
        _rate = CreateAdaptiveRate();
        while (!ct.IsCancellationRequested)
        {
            try
            {
                using TcpLink link = await TcpLink.OpenAsync(Env, Spec.Port, Spec.Bool("tls"), Spec.StreamId, sendAuth: true, ct);
                backoff.Reset();
                _rate?.Reconnected(Env.Clock.NowUs());
                await ExchangeAsync(link, ct);
            }
            catch (OperationCanceledException) when (ct.IsCancellationRequested)
            {
                return;
            }
            catch (Exception ex) when (ex is SocketException or IOException or System.Security.Authentication.AuthenticationException or OperationCanceledException or InvalidDataException or InvalidOperationException)
            {
                // 连接失败或断开：退避之后重连
            }

            try
            {
                await Task.Delay(backoff.Next(), Env.Time, ct);
            }
            catch (OperationCanceledException)
            {
                return;
            }
        }
    }

    public override int Flush(FlushMode mode, long peerProduced = 0)
    {
        RawStream? stream = _rows.Take(Env.Clock.NowUs(), mode == FlushMode.All);
        RawStream kernel = _info.Rows.Drain(includeKey: false);
        if (stream is null && kernel.Count == 0)
        {
            return 0;
        }

        // 内核采样（tcp_info）和回显行一起交出；只有采样没有回显行（例如连接刚断开）时单独交出
        if (stream is null)
        {
            Ship(kernel, "ts");
        }
        else
        {
            Ship(stream, "t1", kernel.Count > 0 ? [kernel] : null);
        }

        return 1;
    }

    private async Task ExchangeAsync(TcpLink link, CancellationToken ct)
    {
        using var cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        Task reader = ReadRepliesAsync(link.Stream, cts.Token);
        Task writer = WriteRequestsAsync(link.Stream, cts.Token);
        Task sampler = _info.RunAsync(link.Socket, TimeSpan.FromMilliseconds(Spec.Int("tcp_info_ms")), cts.Token);
        await Task.WhenAny(reader, writer);
        await cts.CancelAsync();
        foreach (Task task in new[] { reader, writer, sampler })
        {
            try
            {
                await task;
            }
            catch (Exception ex) when (ex is OperationCanceledException or IOException or SocketException or ObjectDisposedException)
            {
                // 连接已经断开
            }
        }
    }

    private async Task WriteRequestsAsync(Stream stream, CancellationToken ct)
    {
        int min = (int)Spec.Int("min");
        int max = (int)Spec.Int("max");
        byte[] body = new byte[Math.Max(max, min)];

        async ValueTask SendOneAsync()
        {
            int size = Math.Max(TcpWire.MinEchoBody, max > min ? Random.Shared.Next(min, max + 1) : min);
            long t1 = Env.Clock.NowUs();
            uint seq = _rows.Add(size, t1, _rate?.Weight ?? 1);
            body.AsSpan(0, size).Clear();
            TcpWire.WriteEchoBody(body, seq, t1);
            await TcpWire.WriteMessageAsync(stream, body.AsMemory(0, size), ct);
            await stream.FlushAsync(ct);
            _rate?.OnSend(t1);
        }

        if (_rate is not null)
        {
            await PaceAsync(_rate, SendOneAsync, ct);
            return;
        }

        using var timer = new PeriodicTimer(TimeSpan.FromSeconds(1.0 / Spec.Int("rate_hz")), Env.Time);
        while (await timer.WaitForNextTickAsync(ct))
        {
            await SendOneAsync();
        }
    }

    private async Task ReadRepliesAsync(Stream stream, CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            byte[]? reply = await TcpWire.ReadMessageAsync(stream, ct);
            if (reply is null)
            {
                return; // 对端关闭
            }

            long now = Env.Clock.NowUs();
            if (TcpWire.TryReadEchoBody(reply, out uint seq, out _) && _rows.Reply(seq, now, out long t1))
            {
                _rate?.OnReply(t1, now, (now - t1) / 1000.0);
            }
        }
    }
}

/// <summary>TCP 回显，监听方：逐条回显并记录 t2（读到）和 t3（回复发出）。保持（hold）的监听方也用它，只是不记录。</summary>
public sealed class TcpEchoListener : BlockBase
{
    private readonly RowStore? _records;
    private readonly TcpInfoRecorder? _info;

    public TcpEchoListener(BlockEnv env, BlockSpec spec, bool record = true)
        : base(env, spec)
    {
        _records = record ? new RowStore(spec.Kind.FindStream("msg")!, Side.Listener, 32768) : null;
        _info = record ? new TcpInfoRecorder(spec.Kind.FindStream("tcp_info")!, Side.Listener, env.Clock) : null;
    }

    public override RowStore? Records(string streamName) => streamName switch
    {
        "msg" => _records,
        "tcp_info" => _info?.Rows,
        _ => null,
    };

    public override async Task RunAsync(CancellationToken ct)
    {
        Env.TcpListen!.Register(Spec.StreamId, EchoAsync);
        try
        {
            await WaitUntilCancelledAsync(ct);
        }
        finally
        {
            Env.TcpListen.Unregister(Spec.StreamId);
        }
    }

    private async Task EchoAsync(TcpConnection connection, CancellationToken ct)
    {
        // 这条连接上 Agent 一侧套接字的内核采样（下行方向的重传、cwnd、srtt）
        using var samplerCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        Task sampler = _info is null ? Task.CompletedTask : _info.RunAsync(connection.Socket, TimeSpan.FromMilliseconds(Spec.Int("tcp_info_ms")), samplerCts.Token);
        try
        {
            await EchoLoopAsync(connection, ct);
        }
        finally
        {
            await samplerCts.CancelAsync();
            await sampler;
        }
    }

    private async Task EchoLoopAsync(TcpConnection connection, CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            byte[]? body = await TcpWire.ReadMessageAsync(connection.Stream, ct);
            if (body is null)
            {
                return;
            }

            long t2 = Env.Clock.NowUs();
            Env.Counters.PacketIn();
            if (!TcpWire.TryReadEchoBody(body, out uint seq, out _))
            {
                Env.Counters.DropInvalid();
                return; // 协议错乱：关闭连接
            }

            if (Env.Admission is { } admission && !admission.AdmitMessage(connection.Session, connection.StreamId, t2))
            {
                Env.Counters.DropRateLimited();
                continue;
            }

            long t3 = Env.Clock.NowUs();
            await TcpWire.WriteMessageAsync(connection.Stream, body, ct);
            await connection.Stream.FlushAsync(ct);
            Env.Counters.PacketOut();
            if (_records is not null)
            {
                Span<long> row = [t2, t3];
                _records.SetRow(seq, row);
            }
        }
    }
}