MoShiro/backend/probe/src/MoShiro.Probe/ProbeApp.cs

using Fleet.Core;
using Fleet.Hosting;
using System.Text.Json;
using Fleet.Identity;
using Fleet.Link;
using Fleet.Rpc;
using MoShiro.Measurement;
using MoShiro.Protocol;

namespace MoShiro.Probe;

/// <summary>
/// Probe 的组装点：到面板的长连接（身份、Agent 列表、状态）+ 每个 Agent 一个测量连接（认证、重连、角色）。没有全局状态，测试可以在同一进程里启动多个。
/// Agent 列表变化时只调和连接：新增的建立，移除的断开，地址或证书指纹变化的重建，其余不受影响。
/// </summary>
public sealed class ProbeApp : IDisposable
{
    public const string Version = MoShiro.Protocol.BuildInfo.Version;

    private readonly ConnectionEnv _env;
    private readonly Dictionary<string, Running> _connections = [];
    private readonly SemaphoreSlim _changed = new(0);
    private readonly INodeLog _log;
    private readonly FleetLink _link;

    /// <param name="identity">注册时由节点自己生成：<c>Identity</c> 用来连面板，<c>PeerIdentity</c> 是另一把独立的密钥，只用来连 Agent。</param>
    public ProbeApp(ProbeFile file, LoadedNodeIdentity identity, string configDirectory, TimeProvider time, INodeLog log, RestartController? restart = null)
    {
        _log = log;
        var trust = TrustAnchor.Load(file.Fleet.IdentityDir, log);
        var directory = new AgentDirectory(new ConfigStore(configDirectory), configDirectory, log, trust);
        _env = new ConnectionEnv(identity.PeerIdentity, identity.NodeId.Value, file.Limits, directory, new MonoClock(time), time, log);
        directory.Changed += () => _changed.Release();

        var update = new UpdateBridge(configDirectory, new ConfigStore(configDirectory).Write, log);
        var options = new FleetLinkOptions
        {
            Server = new Uri(file.Fleet.Server),
            Identity = identity.Identity,
            ServerVerifier = trust,
            NodeId = identity.NodeId,
            NodeVersion = Version,
            Capabilities = [MoShiroCapabilities.Base],
            HeartbeatExtra = () => JsonSerializer.SerializeToElement(new ProbeHeartbeatExtra { Update = update.ReadStatus(), Ipv6 = Ipv6Diagnostics.Sample(null, Environment.GetEnvironmentVariable("MOSHIRO_NETWORK_SCOPE"), Environment.GetEnvironmentVariable("HOST_PROC")) }, MoShiroJsonContext.Default.ProbeHeartbeatExtra),
            HelloMeta = () => JsonSerializer.SerializeToElement(new ProbeHelloMeta { Platform = PlatformDetector.Detect(), Ipv6 = Ipv6Diagnostics.Sample(null, Environment.GetEnvironmentVariable("MOSHIRO_NETWORK_SCOPE"), Environment.GetEnvironmentVariable("HOST_PROC")) }, MoShiroJsonContext.Default.ProbeHelloMeta),
        };
        var list = new StateHandler<ProbeAgentList>(
            ProbeAgentList.StateType,
            MoShiroJsonContext.Default.ProbeAgentList,
            (version, doc, _) =>
            {
                directory.Apply(doc);
                return ValueTask.FromResult(new StateSyncResult { AppliedVersion = version });
            });
        var router = new RpcRouter<LinkContext>(RpcSide.Node, MoShiroMethodTable.Table, NodeRoles.Probe)
            .OnNotify(MoShiroMethods.UpdateOffer, MoShiroJsonContext.Default.UpdateOfferParams, (_, offer, _) =>
            {
                update.WriteOffer(offer);
                return ValueTask.CompletedTask;
            })
            .OnNotify(MoShiroMethods.UpdateApply, MoShiroJsonContext.Default.UpdateOfferParams, (_, _, _) =>
            {
                update.RequestApply();
                return ValueTask.CompletedTask;
            })
            .OnNotify(MoShiroMethods.Schedule, MoShiroJsonContext.Default.ScheduleParams, (_, schedule, _) =>
            {
                if (restart is not null && !restart.SetRemote(schedule.RestartAt))
                {
                    log.Warn("ignoring an invalid restart schedule from the panel");
                }

                return ValueTask.CompletedTask;
            });
        _link = new FleetLink(options, router, [list], new Outbox(1 << 20), time, log);
    }

    /// <summary>到面板的连接（测试和状态查看用）。</summary>
    public FleetLink Link => _link;

    public AgentDirectory Directory => _env.Directory;

    /// <summary>某个 Agent 的会话是否已经建立（测试和状态查看用）。</summary>
    public bool IsConnected(string agentId, string ipVersion = IpVersions.V4)
    {
        lock (_connections)
        {
            return _connections.TryGetValue(agentId + "~" + ipVersion, out Running? running) && running.Connection.IsConnected;
        }
    }

    public AgentConnection? Connection(string agentId, string ipVersion = IpVersions.V4)
    {
        lock (_connections)
        {
            return _connections.TryGetValue(agentId + "~" + ipVersion, out Running? running) ? running.Connection : null;
        }
    }

    public async Task RunAsync(CancellationToken ct)
    {
        if (_env.Directory.Agents.Count == 0)
        {
            _log.Info("no agents known yet, waiting for the panel to send the list");
        }

        Task link = _link.RunAsync(ct);
        try
        {
            while (!ct.IsCancellationRequested)
            {
                await ReconcileAsync(_env.Directory.Agents, ct);
                await _changed.WaitAsync(ct);
            }
        }
        catch (OperationCanceledException)
        {
            // 停止
        }
        finally
        {
            await ReconcileAsync([], CancellationToken.None);
            await Task.WhenAny(link, Task.Delay(TimeSpan.FromSeconds(5), CancellationToken.None));
        }
    }

    public void Dispose()
    {
        _env.Identity.Dispose();
        _changed.Dispose();
    }

    private async Task ReconcileAsync(IReadOnlyList<KnownAgent> agents, CancellationToken ct)
    {
        var wanted = agents.ToDictionary(a => a.PathKey);
        var stopped = new List<Running>();
        lock (_connections)
        {
            foreach ((string id, Running running) in _connections.ToList())
            {
                if (!wanted.TryGetValue(id, out KnownAgent? endpoint) || !AgentDirectory.Equal(endpoint, running.Connection.Endpoint))
                {
                    stopped.Add(running);
                    _connections.Remove(id);
                }
            }
        }

        foreach (Running running in stopped)
        {
            await running.StopAsync();
        }

        if (ct.IsCancellationRequested)
        {
            return;
        }

        lock (_connections)
        {
            foreach ((string id, KnownAgent endpoint) in wanted)
            {
                if (!_connections.ContainsKey(id))
                {
                    _connections[id] = Running.Start(new AgentConnection(endpoint, _env));
                }
            }
        }
    }

    private sealed class Running : IDisposable
    {
        private readonly CancellationTokenSource _cts = new();
        private Task _run = Task.CompletedTask;

        private Running(AgentConnection connection)
        {
            Connection = connection;
        }

        public AgentConnection Connection { get; }

        public static Running Start(AgentConnection connection)
        {
            var running = new Running(connection);
            running._run = Task.Run(() => connection.RunAsync(running._cts.Token), CancellationToken.None);
            return running;
        }

        public async Task StopAsync()
        {
            await _cts.CancelAsync();
            try
            {
                await _run.WaitAsync(TimeSpan.FromSeconds(10));
            }
            catch (Exception ex) when (ex is OperationCanceledException or TimeoutException)
            {
                // 没有及时退出：不再等
            }

            Dispose();
        }

        public void Dispose() => _cts.Dispose();
    }
}