MoShiro/backend/shared/MoShiro.Measurement/Net/TcpInfo.cs

using System.Buffers.Binary;
using System.Net.Sockets;
using MoShiro.Protocol;

namespace MoShiro.Measurement;

/// <summary>内核对一条 TCP 连接的采样（Linux <c>TCP_INFO</c>）。</summary>
public readonly record struct TcpInfoSample(long SrttUs, long RttVarUs, long Cwnd, long TotalRetrans, long SegsOut, long BytesAcked);

/// <summary>
/// 读 Linux 的 <c>TCP_INFO</c>：通过 <see cref="Socket.GetRawSocketOption(int, int, Span{byte})"/>，不需要 P/Invoke（NativeAOT、glibc 和 musl 都一样）。
/// 缓冲区先清零，并且只读内核实际返回的长度之内的字段：老内核返回的结构更短，缺的字段是 0，而不是垃圾值。
/// 其他平台返回 false（没有这一类采样，指标在对应窗口里就是没有数据）。
/// </summary>
public static class TcpInfo
{
    private const int SolTcp = 6;
    private const int TcpInfoOption = 11;

    // struct tcp_info 的偏移（include/uapi/linux/tcp.h）：8 字节的状态位之后是 u32 字段；total_retrans 之后是对齐到 8 字节的 u64 字段
    private const int RttOffset = 68;
    private const int RttVarOffset = 72;
    private const int SndCwndOffset = 80;
    private const int TotalRetransOffset = 100;
    private const int BytesAckedOffset = 120;
    private const int SegsOutOffset = 136;

    /// <summary>读到 <c>segs_out</c> 为止（内核 ≥ 4.2）。</summary>
    public const int Length = SegsOutOffset + 4;

    public static bool Supported => OperatingSystem.IsLinux();

    public static bool TrySample(Socket socket, out TcpInfoSample sample)
    {
        sample = default;
        if (!Supported)
        {
            return false;
        }

        try
        {
            Span<byte> buffer = stackalloc byte[Length];
            buffer.Clear();
            int returned = socket.GetRawSocketOption(SolTcp, TcpInfoOption, buffer);
            if (returned < RttVarOffset + 4)
            {
                return false;
            }

            ReadOnlySpan<byte> info = buffer[..Math.Min(returned, Length)];
            sample = new TcpInfoSample(
                U32(info, RttOffset),
                U32(info, RttVarOffset),
                U32(info, SndCwndOffset),
                U32(info, TotalRetransOffset),
                U32(info, SegsOutOffset),
                U64(info, BytesAckedOffset));
            return true;
        }
        catch (Exception ex) when (ex is SocketException or ObjectDisposedException or PlatformNotSupportedException)
        {
            return false;
        }
    }

    private static long U32(ReadOnlySpan<byte> info, int offset) => offset + 4 <= info.Length ? BinaryPrimitives.ReadUInt32LittleEndian(info[offset..]) : 0;

    private static long U64(ReadOnlySpan<byte> info, int offset) => offset + 8 <= info.Length ? (long)BinaryPrimitives.ReadUInt64LittleEndian(info[offset..]) : 0;
}

/// <summary>
/// 一条连接上周期性的内核采样（<c>tcp_info</c> 流）：每个采样一行 (ts, side, srtt, rttvar, cwnd, retrans_total, segs_out, bytes_acked, conn)。
/// 两端各采样自己的套接字：拨号方 的重传反映上行方向，Agent 的重传反映下行方向（代理场景里大流量走这个方向）。
/// 累计计数器在 Backend 按 (side, conn) 分别求差；conn 是当前记录器内分配的连接编号，不包含网络地址。
/// </summary>
public sealed class TcpInfoRecorder
{
    private readonly MonoClock _clock;
    private readonly long _side;
    private long _nextConnection;

    public TcpInfoRecorder(StreamSpec spec, Side side, MonoClock clock)
    {
        Rows = new RowStore(spec, side, 4096);
        _clock = clock;
        _side = side == Side.Dialer ? 0 : 1;
    }

    public RowStore Rows { get; }

    /// <summary>每隔 <paramref name="interval"/> 采样一次，直到取消或套接字不可用（连接断开）。</summary>
    public async Task RunAsync(Socket socket, TimeSpan interval, CancellationToken ct)
    {
        if (!TcpInfo.Supported)
        {
            return;
        }

        // A recorder survives reconnects, and listeners can briefly have more
        // than one socket. Kernel counters belong to a socket, not just a side.
        long connection = Interlocked.Increment(ref _nextConnection);
        try
        {
            using var timer = new PeriodicTimer(interval, _clock.Time);
            while (await timer.WaitForNextTickAsync(ct))
            {
                if (!TcpInfo.TrySample(socket, out TcpInfoSample s))
                {
                    return;
                }

                long key = Rows.Append();
                Rows.SetRow(key, [_clock.NowUs(), _side, s.SrttUs, s.RttVarUs, s.Cwnd, s.TotalRetrans, s.SegsOut, s.BytesAcked, connection]);
            }
        }
        catch (OperationCanceledException)
        {
            // 连接结束
        }
    }
}