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

using Fleet.Core;
using Fleet.Hosting;
using Fleet.Identity;
using Monitor.Measurement;

namespace Monitor.Probe;

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

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

    public ProbeApp(ProbeFile file, string configDirectory, TimeProvider time, INodeLog log)
    {
        _log = log;
        var directory = new AgentDirectory(new ConfigStore(configDirectory), configDirectory, file, log);
        NodeIdentity identity = NodeIdentity.FromPkcs8(Convert.FromBase64String(file.PrivateKey));
        _env = new ConnectionEnv(identity, file.ProbeId, file.Limits, directory, new MonoClock(time), time, log);
        directory.Changed += () => _changed.Release();
    }

    public AgentDirectory Directory => _env.Directory;

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

    public AgentConnection? Connection(string agentId)
    {
        lock (_connections)
        {
            return _connections.TryGetValue(agentId, out Running? running) ? running.Connection : null;
        }
    }

    public async Task RunAsync(CancellationToken ct)
    {
        if (_env.Directory.Agents.Count == 0)
        {
            _log.Warn("no agents known yet: probe.json has an empty bootstrap list");
        }

        try
        {
            while (!ct.IsCancellationRequested)
            {
                await ReconcileAsync(_env.Directory.Agents, ct);
                await _changed.WaitAsync(ct);
            }
        }
        catch (OperationCanceledException)
        {
            // 停止
        }
        finally
        {
            await ReconcileAsync([], CancellationToken.None);
        }
    }

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

    private async Task ReconcileAsync(IReadOnlyList<BootstrapAgent> agents, CancellationToken ct)
    {
        var wanted = agents.ToDictionary(a => a.AgentId);
        var stopped = new List<Running>();
        lock (_connections)
        {
            foreach ((string id, Running running) in _connections.ToList())
            {
                if (!wanted.TryGetValue(id, out BootstrapAgent? 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, BootstrapAgent 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();
    }
}