298 lines
12 KiB
C#
298 lines
12 KiB
C#
using System.Collections.Concurrent;
|
||
using Fleet.Core;
|
||
using Fleet.Rpc;
|
||
using MoShiro.Measurement;
|
||
using MoShiro.Protocol;
|
||
|
||
namespace MoShiro.Agent;
|
||
|
||
/// <summary>
|
||
/// 按需测试的编排（<c>moshiro.test.run</c>）：Agent 是时间的主人。每个阶段：
|
||
/// 登记数据面的流 → <c>moshiro.test.phase</c>（Probe 准备好自己一侧的角色）→ 运行（单向数据流由 Agent 发包，回显类等待 Probe 发包）
|
||
/// → <c>moshiro.test.phase_end</c>（Probe 收尾，回复它交出的列块个数）→ 等所有列块合并上传 → <c>phase_done</c>。
|
||
/// 同一时间只能运行 <c>max_concurrent_tests</c> 个测试；Probe 掉线或被取消时测试中止，已经上传的阶段数据保留。
|
||
/// </summary>
|
||
public sealed class TestRunner
|
||
{
|
||
private static readonly TimeSpan s_phaseSlack = TimeSpan.FromSeconds(15);
|
||
private static readonly TimeSpan s_fragmentWait = TimeSpan.FromSeconds(10);
|
||
|
||
private readonly ConcurrentDictionary<string, RunningTest> _running = new();
|
||
private readonly TaskManager _tasks;
|
||
private readonly SessionRegistry _sessions;
|
||
private readonly StreamRegistry _streams;
|
||
private readonly BlockEnv _env;
|
||
private readonly BackendUplink _uplink;
|
||
private readonly int _maxConcurrent;
|
||
private readonly Func<string, IReadOnlyList<int>> _ports;
|
||
private readonly TimeProvider _time;
|
||
private readonly INodeLog _log;
|
||
|
||
public TestRunner(TaskManager tasks, SessionRegistry sessions, StreamRegistry streams, BlockEnv env, BackendUplink uplink, int maxConcurrent, Func<string, IReadOnlyList<int>> ports, TimeProvider time, INodeLog log)
|
||
{
|
||
_tasks = tasks;
|
||
_sessions = sessions;
|
||
_streams = streams;
|
||
_env = env;
|
||
_uplink = uplink;
|
||
_maxConcurrent = Math.Max(1, maxConcurrent);
|
||
_ports = ports;
|
||
_time = time;
|
||
_log = log;
|
||
}
|
||
|
||
public int Running => _running.Count;
|
||
|
||
/// <summary>校验并接受一个测试（立即返回），实际运行在后台进行，状态通过 <c>moshiro.test.status</c> 回报。</summary>
|
||
public Result<TestRunResult> Start(TestRunParams run)
|
||
{
|
||
if (!IpVersions.IsValid(run.IpVersion))
|
||
{
|
||
return Result<TestRunResult>.Fail(MoShiroErrors.ValidationFailed, "ip_version must be ipv4 or ipv6");
|
||
}
|
||
|
||
if (_sessions.Get(run.ProbeId.Value, run.IpVersion) is not { } probe)
|
||
{
|
||
return Result<TestRunResult>.Fail(MoShiroErrors.ProbeOffline, $"probe {run.ProbeId.Value} is not connected to this agent");
|
||
}
|
||
|
||
if (run.Phases.Count == 0)
|
||
{
|
||
return Result<TestRunResult>.Fail(MoShiroErrors.ValidationFailed, "a test needs at least one phase");
|
||
}
|
||
|
||
var resolved = new List<(TestPhase Phase, TaskDecision Decision)>();
|
||
foreach (TestPhase phase in run.Phases)
|
||
{
|
||
WithDefaultPort(phase);
|
||
TaskDecision decision = _tasks.Guard.CheckPhase(phase.Kind, phase.Params);
|
||
if (!decision.Accepted)
|
||
{
|
||
return Result<TestRunResult>.Fail(
|
||
decision.Reason == RejectReasons.Unsupported ? FleetErrors.Unsupported : MoShiroErrors.BudgetExceeded,
|
||
$"phase {phase.Kind} refused: {decision.Reason}");
|
||
}
|
||
|
||
string? direction = decision.Resolved!.Values.ContainsKey("direction") ? decision.Resolved.GetString("direction") : null;
|
||
if (BlockCatalog.Find(phase.Kind) is not { Listener: not null } pair || !pair.Supports(direction))
|
||
{
|
||
return Result<TestRunResult>.Fail(FleetErrors.Unsupported, $"{phase.Kind} does not support direction={direction}");
|
||
}
|
||
}
|
||
|
||
foreach (TestPhase phase in run.Phases)
|
||
{
|
||
resolved.Add((phase, _tasks.Guard.CheckPhase(phase.Kind, phase.Params)));
|
||
}
|
||
|
||
var running = new RunningTest(probe);
|
||
if (_running.Count >= _maxConcurrent || !_running.TryAdd(run.TestId.Value, running))
|
||
{
|
||
running.Dispose();
|
||
return Result<TestRunResult>.Fail(MoShiroErrors.TestBusy, "this agent is already running a test");
|
||
}
|
||
|
||
_ = Task.Run(() => RunAsync(run.TestId, resolved, running), CancellationToken.None);
|
||
return Result<TestRunResult>.Ok(new TestRunResult { Accepted = true });
|
||
}
|
||
|
||
/// <summary>阶段没有指定端口时用本 Agent 自己监听的第一个端口（Backend 的模板不知道每个 Agent 的实际端口）。</summary>
|
||
private void WithDefaultPort(TestPhase phase)
|
||
{
|
||
KindSpec? kind = Kinds.Find(phase.Kind);
|
||
if (kind is null || !kind.Params.Any(p => p.Name == "port"))
|
||
{
|
||
return;
|
||
}
|
||
|
||
// 按需测试的数据端口由 Agent 决定：Backend 给的端口如果不是本机监听的端口（模板里是出厂默认值，每个 Agent 可以改，
|
||
// 而 Backend 并不知道 TCP 数据端口），就用本机监听的第一个端口。块头里记录的是实际用的端口。
|
||
IReadOnlyList<int> ports = _ports(kind.Blocks.Transport == "udp" ? "udp" : "tcp");
|
||
bool listening = phase.Params.TryGetValue("port", out System.Text.Json.JsonElement requested) && requested.TryGetInt32(out int port) && ports.Contains(port);
|
||
if (!listening)
|
||
{
|
||
phase.Params["port"] = System.Text.Json.JsonDocument.Parse(ports[0].ToString(System.Globalization.CultureInfo.InvariantCulture)).RootElement;
|
||
}
|
||
}
|
||
|
||
public bool Cancel(TestId test)
|
||
{
|
||
if (_running.TryGetValue(test.Value, out RunningTest? running))
|
||
{
|
||
running.Cancel.Cancel();
|
||
return true;
|
||
}
|
||
|
||
return false;
|
||
}
|
||
|
||
/// <summary>某个 Probe 掉线：它上面的测试中止，状态是 failed / probe_offline；别的 Probe 上的测试不受影响。</summary>
|
||
public void ProbeLost(ProbeSession probe)
|
||
{
|
||
foreach (RunningTest running in _running.Values)
|
||
{
|
||
if (ReferenceEquals(running.Probe, probe))
|
||
{
|
||
running.Lost = true;
|
||
running.Cancel.Cancel();
|
||
}
|
||
}
|
||
}
|
||
|
||
private async Task RunAsync(TestId test, List<(TestPhase Phase, TaskDecision Decision)> phases, RunningTest running)
|
||
{
|
||
CancellationToken ct = running.Cancel.Token;
|
||
string finalStatus = TestStatuses.Succeeded;
|
||
string? errorCode = null;
|
||
int? lastPhase = null;
|
||
try
|
||
{
|
||
ProbeSession probe = running.Probe;
|
||
for (int index = 0; index < phases.Count; index++)
|
||
{
|
||
lastPhase = index;
|
||
await _uplink.SendTestStatusAsync(test, TestStatuses.Running, index);
|
||
await RunPhaseAsync(probe, test, index, phases[index].Phase, phases[index].Decision, ct);
|
||
await _uplink.SendTestStatusAsync(test, TestStatuses.PhaseDone, index);
|
||
}
|
||
}
|
||
catch (OperationCanceledException)
|
||
{
|
||
finalStatus = running.Lost ? TestStatuses.Failed : TestStatuses.Aborted;
|
||
errorCode = running.Lost ? MoShiroErrors.ProbeOffline.Value : null;
|
||
await TryAbortProbeAsync(running.Probe, test);
|
||
}
|
||
catch (TestFailure failure)
|
||
{
|
||
finalStatus = TestStatuses.Failed;
|
||
errorCode = failure.Code;
|
||
await TryAbortProbeAsync(running.Probe, test);
|
||
}
|
||
finally
|
||
{
|
||
_running.TryRemove(test.Value, out _);
|
||
running.Dispose();
|
||
}
|
||
|
||
_log.Info("test finished", ("test", test.Value), ("status", finalStatus));
|
||
await _uplink.SendTestStatusAsync(test, finalStatus, lastPhase, errorCode);
|
||
}
|
||
|
||
private async Task RunPhaseAsync(ProbeSession probe, TestId test, int index, TestPhase phase, TaskDecision decision, CancellationToken ct)
|
||
{
|
||
ParamValues values = decision.Resolved!;
|
||
var stream = new StreamInfo(_streams.Allocate(), phase.Kind, values) { TestId = test, Phase = index, Owner = probe };
|
||
int durationSeconds = phase.DurationSeconds > 0 ? phase.DurationSeconds : (int)(values.Values.ContainsKey("duration_s") ? values.GetInt("duration_s") : 10);
|
||
var blockSpec = new BlockSpec(decision.Kind!, values, stream.Id, stream.Port, new FragmentScope(null, test, index));
|
||
stream.Running = RunningBlock.Start(BlockCatalog.Find(phase.Kind)!.Listener!(_env.WithSession(probe), blockSpec), blockSpec.Signature);
|
||
_streams.Add(stream);
|
||
try
|
||
{
|
||
RpcSession<PeerContext> rpc = probe.Rpc ?? throw new TestFailure(MoShiroErrors.ProbeOffline.Value);
|
||
var start = new TestPhaseStartParams
|
||
{
|
||
TestId = test,
|
||
Phase = index,
|
||
Kind = phase.Kind,
|
||
Params = new Dictionary<string, System.Text.Json.JsonElement>(values.Values),
|
||
StreamId = stream.Id,
|
||
DurationSeconds = durationSeconds,
|
||
Port = stream.Port,
|
||
};
|
||
Result<TestPhaseStartResult> accepted = await rpc.RequestAsync(MoShiroMethods.TestPhase, start, MoShiroJsonContext.Default.TestPhaseStartParams, MoShiroJsonContext.Default.TestPhaseStartResult, ct);
|
||
if (!accepted.TryGet(out TestPhaseStartResult? started, out FleetError? error) || !started.Accepted)
|
||
{
|
||
throw new TestFailure(error?.Code.Value ?? FleetErrors.Unsupported.Value);
|
||
}
|
||
|
||
// 单向数据流的发送端自己运行到结束（或失败）；其他积木是反应式的，等够时长即可
|
||
if (phase.Kind == "udp_train")
|
||
{
|
||
try
|
||
{
|
||
await stream.Running.Completion.WaitAsync(ct);
|
||
}
|
||
catch (BlockFaultException)
|
||
{
|
||
throw new TestFailure(MoShiroErrors.ProbeOffline.Value);
|
||
}
|
||
}
|
||
else
|
||
{
|
||
await Task.Delay(TimeSpan.FromSeconds(durationSeconds), _time, ct);
|
||
}
|
||
|
||
long sent = stream.Running.Block.Produced;
|
||
var end = new TestPhaseEndParams { TestId = test, Phase = index, StreamId = stream.Id, SentCount = sent };
|
||
Result<TestPhaseEndResult> ended = await rpc.RequestAsync(MoShiroMethods.TestPhaseEnd, end, MoShiroJsonContext.Default.TestPhaseEndParams, MoShiroJsonContext.Default.TestPhaseEndResult, ct, s_phaseSlack);
|
||
if (!ended.TryGet(out TestPhaseEndResult? finished, out FleetError? endError))
|
||
{
|
||
throw new TestFailure(endError.Code.Value);
|
||
}
|
||
|
||
// 列块是经 Probe 的发件箱至少一次送达的：全部合并上传之后才进入下一阶段（超时就带着已有的数据继续）
|
||
DateTimeOffset deadline = _time.GetUtcNow() + s_fragmentWait;
|
||
while (stream.FragmentsDone < finished.Fragments && _time.GetUtcNow() < deadline)
|
||
{
|
||
await Task.Delay(TimeSpan.FromMilliseconds(20), _time, ct);
|
||
}
|
||
|
||
if (stream.FragmentsDone < finished.Fragments)
|
||
{
|
||
_log.Warn("fragments missing", ("test", test.Value), ("phase", index), ("expected", finished.Fragments), ("got", stream.FragmentsDone));
|
||
}
|
||
}
|
||
finally
|
||
{
|
||
_streams.Remove(stream.Id);
|
||
await stream.Running.StopAsync();
|
||
probe.ForgetStream(stream.Id);
|
||
}
|
||
}
|
||
|
||
private async Task TryAbortProbeAsync(ProbeSession probe, TestId test)
|
||
{
|
||
if (probe.Rpc is not { } rpc)
|
||
{
|
||
return;
|
||
}
|
||
|
||
try
|
||
{
|
||
await rpc.NotifyAsync(MoShiroMethods.TestAbort, new TestAbortParams { TestId = test }, MoShiroJsonContext.Default.TestAbortParams, CancellationToken.None);
|
||
}
|
||
catch (Exception ex) when (ex is ObjectDisposedException or InvalidOperationException or IOException)
|
||
{
|
||
// 会话已经结束
|
||
}
|
||
}
|
||
|
||
/// <summary>一个正在运行的测试：它在哪个 Probe 上、怎样取消、是不是因为 Probe 掉线而中止。</summary>
|
||
private sealed class RunningTest : IDisposable
|
||
{
|
||
public RunningTest(ProbeSession probe)
|
||
{
|
||
Probe = probe;
|
||
}
|
||
|
||
public ProbeSession Probe { get; }
|
||
|
||
public CancellationTokenSource Cancel { get; } = new();
|
||
|
||
public volatile bool Lost;
|
||
|
||
public void Dispose() => Cancel.Dispose();
|
||
}
|
||
|
||
private sealed class TestFailure : Exception
|
||
{
|
||
public TestFailure(string code)
|
||
: base(code)
|
||
{
|
||
Code = code;
|
||
}
|
||
|
||
public string Code { get; }
|
||
}
|
||
}
|