MoShiro/backend/agent/src/MoShiro.Agent/Dial/DialSession.cs

using System.Collections.Concurrent;
using Fleet.Link;
using Fleet.Rpc;
using MoShiro.Measurement;
using MoShiro.Protocol;

namespace MoShiro.Agent;

/// <summary>一次会话内的运行状态：线路任务的积木、按需测试的阶段。方法由路由器在单个入站队列里顺序调用。</summary>
public sealed class DialSession
{
    private const int MaxBlocks = 16;

    private readonly object _gate = new();
    private readonly Dictionary<string, RunningBlock> _lineBlocks = [];
    private readonly Dictionary<string, string> _lineKinds = [];
    private readonly ConcurrentDictionary<uint, PhaseState> _phases = new();

    public DialSession(BlockEnv blockEnv, ClockEstimator clock, Outbox outbox, DialEnv env)
    {
        BlockEnv = blockEnv;
        Clock = clock;
        Outbox = outbox;
        Env = env;
    }

    public BlockEnv BlockEnv { get; }

    public ClockEstimator Clock { get; }

    public Outbox Outbox { get; }

    public DialEnv Env { get; }

    public int RoleCount
    {
        get
        {
            lock (_gate)
            {
                return _lineBlocks.Count;
            }
        }
    }

    /// <summary>正在运行的 TCP 类线路积木数（它们自带套接字和重拨，不依赖 QUIC 数据通道）。</summary>
    public int TcpBlockCount
    {
        get
        {
            lock (_gate)
            {
                return _lineKinds.Values.Count(IsTcpKind);
            }
        }
    }

    private static bool IsTcpKind(string kind) => kind is "tcp_ping" or "tcp_connect" or "tls_hold" or "tcp_transfer";

    /// <summary>调和线路任务：参数没有变化的积木继续运行，其余启动或停止。返回被拒绝的任务（不认识的 Kind、参数非法、超过上限）。</summary>
    public async Task<List<RejectedItem>> ApplyAsync(LineTaskSet set)
    {
        var rejected = new List<RejectedItem>();
        var wanted = new Dictionary<string, BlockSpec>();
        foreach (LineTask task in set.Tasks)
        {
            BlockSpec? spec = BlockCatalog.DialerKinds.Contains(task.Kind)
                ? BlockSpec.Create(task.Kind, task.Params, task.StreamId, task.Port, new FragmentScope(task.Id, null, null))
                : null;
            if (spec is null)
            {
                rejected.Add(new RejectedItem { Id = task.Id.Value, Reason = BlockCatalog.DialerKinds.Contains(task.Kind) ? "invalid_params" : "unsupported" });
            }
            else if (wanted.Count >= MaxBlocks)
            {
                rejected.Add(new RejectedItem { Id = task.Id.Value, Reason = "budget_exceeded" });
            }
            else
            {
                wanted[task.Id.Value] = spec;
            }
        }

        var stopped = new List<RunningBlock>();
        lock (_gate)
        {
            foreach ((string id, RunningBlock running) in _lineBlocks.ToList())
            {
                if (!wanted.TryGetValue(id, out BlockSpec? spec) || spec.Signature != running.Signature)
                {
                    stopped.Add(running);
                    _lineBlocks.Remove(id);
                    _lineKinds.Remove(id);
                }
            }

            foreach ((string id, BlockSpec spec) in wanted)
            {
                if (!_lineBlocks.ContainsKey(id))
                {
                    _lineBlocks[id] = RunningBlock.Start(BlockCatalog.Find(spec.KindId)!.Dialer(BlockEnv, spec), spec.Signature);
                    _lineKinds[id] = spec.KindId;
                }
            }
        }

        // 被停止的积木把已经稳定的行交出去（刚发出、还在等回复的行不交：它们不是丢失，只是没有等到）
        foreach (RunningBlock running in stopped)
        {
            await running.StopAsync();
            running.Block.Flush(FlushMode.Settled);
        }

        return rejected;
    }

    public void FlushLineBlocks()
    {
        List<RunningBlock> blocks;
        lock (_gate)
        {
            blocks = [.. _lineBlocks.Values];
        }

        foreach (RunningBlock running in blocks)
        {
            running.Block.Flush(FlushMode.Settled);
        }
    }

    /// <summary>按需测试的一个阶段：起对应的拨号方积木（单向数据流是接收端）。超过本地上限（或不认识）则拒绝。</summary>
    public bool StartPhase(TestPhaseStartParams start)
    {
        string? direction = start.Params.TryGetValue("direction", out System.Text.Json.JsonElement d) ? d.GetString() : null;
        if (_phases.ContainsKey(start.StreamId) || !BlockCatalog.DialerKinds.Contains(start.Kind) || BlockCatalog.Find(start.Kind)!.Supports(direction) is false)
        {
            return false;
        }

        if (start.Kind == "quic_stream")
        {
            double mbps = start.Params["mbps"].GetDouble();
            long size = start.Params["size"].GetInt64();
            long pps = (long)(mbps * 1_000_000 / 8 / Math.Max(1, size));
            if (mbps > Env.Limits.MaxTotalMbps || pps > Env.Limits.MaxTotalPps)
            {
                return false;
            }
        }

        BlockSpec? spec = BlockSpec.Create(start.Kind, start.Params, start.StreamId, start.Port, new FragmentScope(null, start.TestId, start.Phase));
        if (spec is null)
        {
            return false;
        }

        _phases[start.StreamId] = new PhaseState(start.TestId, start.Kind, RunningBlock.Start(BlockCatalog.Find(start.Kind)!.Dialer(BlockEnv, spec), spec.Signature));
        return true;
    }

    /// <summary>阶段结束：收尾、交出列块，返回块数。</summary>
    public async Task<int> EndPhaseAsync(TestPhaseEndParams end, CancellationToken ct)
    {
        if (!_phases.TryRemove(end.StreamId, out PhaseState? phase))
        {
            return 0;
        }

        // 最后一批包（单向数据流）、最后几个回显的回复还在途中
        await Task.Delay(TimeSpan.FromMilliseconds(phase.Kind == "quic_stream" ? 200 : 300), BlockEnv.Time, ct);
        await phase.Running.StopAsync();
        return phase.Running.Block.Flush(FlushMode.All, end.SentCount);
    }

    public async Task AbortAsync(TestId test)
    {
        foreach ((uint stream, PhaseState phase) in _phases.ToList())
        {
            if (phase.Test == test && _phases.TryRemove(stream, out _))
            {
                await phase.Running.StopAsync();
            }
        }
    }

    public async Task StopAllAsync()
    {
        List<RunningBlock> blocks;
        lock (_gate)
        {
            blocks = [.. _lineBlocks.Values];
            _lineBlocks.Clear();
            _lineKinds.Clear();
        }

        blocks.AddRange(_phases.Values.Select(p => p.Running));
        _phases.Clear();
        foreach (RunningBlock running in blocks)
        {
            await running.StopAsync();
        }
    }

    private sealed record PhaseState(TestId Test, string Kind, RunningBlock Running);
}