MoShiro/backend/agent/src/MoShiro.Agent/Tasks/TargetRunners.cs

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

namespace MoShiro.Agent;

/// <summary>
/// 目标任务（Agent 自己探测一个地址，没有对端）：用和线路任务同一套积木（<see cref="BlockCatalog"/> 里的拨号方，例如 icmp_echo、connect），
/// 只是列块不交给 Agent 合并，而是自己盖章、10 秒一块直接上传。目标地址必须在本地 <c>allowed_target_cidrs</c> 内
/// （由 <see cref="BudgetGuard"/> 在接受任务之前检查）。
/// </summary>
public sealed class TargetRunners
{
    private static readonly TimeSpan s_chunkSpan = TimeSpan.FromSeconds(ProtocolLimits.ChunkSeconds);

    private readonly Dictionary<string, Running> _running = [];
    private readonly object _gate = new();
    private readonly ChunkAssembler _chunks;
    private readonly NodeId _agent;
    private readonly MonoClock _clock;
    private readonly TimeProvider _time;
    private readonly INodeLog _log;

    public TargetRunners(ChunkAssembler chunks, NodeId agent, MonoClock clock, TimeProvider time, INodeLog log)
    {
        _chunks = chunks;
        _agent = agent;
        _clock = clock;
        _time = time;
        _log = log;
    }

    /// <summary>让运行中的任务与期望集合一致：参数没有变化的保持，其余启动或停止。</summary>
    public void Reconcile(IReadOnlyList<(TaskSpec Spec, TaskDecision Decision)> desired)
    {
        lock (_gate)
        {
            var wanted = desired.ToDictionary(d => d.Spec.Id.Value, d => d);
            foreach ((string id, Running running) in _running.ToList())
            {
                if (!wanted.TryGetValue(id, out (TaskSpec Spec, TaskDecision Decision) entry) || Signature(entry.Spec) != running.Signature)
                {
                    running.Cancel.Cancel();
                    _running.Remove(id);
                }
            }

            foreach ((string id, (TaskSpec spec, TaskDecision decision)) in wanted)
            {
                if (_running.ContainsKey(id))
                {
                    continue;
                }

                var cts = new CancellationTokenSource();
                _running[id] = new Running(Signature(spec), cts);
                _ = Task.Run(() => RunAsync(spec, decision, cts.Token), CancellationToken.None);
            }
        }
    }

    private static string Signature(TaskSpec spec) =>
        spec.Kind + "|" + spec.Target + "|" + string.Join(',', spec.Params.OrderBy(p => p.Key, StringComparer.Ordinal).Select(p => p.Key + "=" + p.Value.GetRawText()));

    private async Task RunAsync(TaskSpec spec, TaskDecision decision, CancellationToken ct)
    {
        try
        {
            if (!TryResolve(spec.Target!, decision.Resolved!, spec.Kind, out IPAddress? address, out int port))
            {
                _log.Warn("target is not an address", ("task", spec.Id.Value));
                return;
            }

            var env = new BlockEnv
            {
                Clock = _clock,
                Log = _log,
                Session = () => null,
                PeerAddress = address,
                Fragments = new LocalSink(this, spec, address.AddressFamily == System.Net.Sockets.AddressFamily.InterNetworkV6 ? IpVersions.V6 : IpVersions.V4),
            };
            IBlock block = BlockCatalog.Find(spec.Kind)!.Dialer(env, new BlockSpec(decision.Kind!, decision.Resolved!, 0, port, new FragmentScope(spec.Id, null, null)));
            using var cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
            Task run = block.RunAsync(cts.Token);
            try
            {
                using var timer = new PeriodicTimer(s_chunkSpan, _time);
                while (await timer.WaitForNextTickAsync(ct))
                {
                    block.Flush(FlushMode.Settled);
                }
            }
            finally
            {
                await cts.CancelAsync();
                await run;
                block.Flush(FlushMode.All);
            }
        }
        catch (OperationCanceledException)
        {
            // 任务被撤销
        }
        catch (Exception ex) when (ex is System.Net.Sockets.SocketException or InvalidOperationException or BlockFaultException)
        {
            _log.Warn("target task stopped", ("task", spec.Id.Value), ("err", ex.Message));
        }
    }

    private static bool TryResolve(string target, ParamValues values, string kind, out IPAddress address, out int port)
    {
        port = kind == "connect" ? (int)values.GetInt("port") : 0;
        if (IPEndPoint.TryParse(target, out IPEndPoint? endpoint))
        {
            address = endpoint.Address;
            if (endpoint.Port != 0)
            {
                port = endpoint.Port;
            }

            return true;
        }

        return IPAddress.TryParse(target, out address!);
    }

    /// <summary>目标任务的行没有 Probe 一侧可以合并：积木交出的就是完整的行，Agent 自己盖章（目标、时钟锚点）并放进发件箱。</summary>
    private sealed class LocalSink : IFragmentSink
    {
        private readonly TargetRunners _owner;
        private readonly TaskSpec _spec;
        private readonly string _ipVersion;

        public LocalSink(TargetRunners owner, TaskSpec spec, string ipVersion)
        {
            _owner = owner;
            _spec = spec;
            _ipVersion = ipVersion;
        }

        public void Ship(string kind, FragmentScope scope, IReadOnlyList<RawStream> streams, IReadOnlyDictionary<string, JsonElement>? parameters, long startUtcMs, int durationMs, int port)
        {
            ClockAnchor anchor = _owner._clock.Anchor();
            var header = new RawChunkHeader
            {
                ChunkId = ChunkId.New(_owner._time),
                Kind = kind,
                TaskId = scope.Task,
                AgentId = _owner._agent,
                Target = _spec.Target,
                IpVersion = _ipVersion,
                Params = parameters is null ? [] : new Dictionary<string, JsonElement>(parameters),
                Port = port,
                StartUtcMs = startUtcMs,
                DurationMs = Math.Clamp(durationMs, 1, ChunkFinalizer.MaxDurationMs),
                Clock = new ChunkClock { AgentMonoUs = anchor.AgentMonoUs, UtcMs = anchor.UtcMs, ProbeOffsetUs = 0 },
            };
            _owner._chunks.Emit(new RawChunk(header, streams));
        }
    }

    private sealed record Running(string Signature, CancellationTokenSource Cancel);
}