MoShiro/backend/fleet/src/Fleet.Server.Link/FleetLinkHost.cs

using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Text.Json;
using System.Text.Json.Serialization.Metadata;
using Fleet.Core;
using Fleet.Rpc;
using Microsoft.Extensions.Logging;

namespace Fleet.Server;

/// <summary>
/// 服务端连接宿主：在线表 + 服务端发起的调用 + 状态推送通知。每个 WebSocket 对应一个 <see cref="NodeConnection"/>。
/// </summary>
public sealed class FleetLinkHost : IFleetConnections, IStateNotifier
{
    private readonly ConcurrentDictionary<NodeId, NodeConnection> _connections = new();

    internal FleetLinkHost(
        IFleetStore store,
        FleetServerOptions options,
        TimeProvider time,
        ILogger<FleetLinkHost> logger,
        IEnumerable<IStateProvider> stateProviders,
        IEnumerable<IInbox> inboxes,
        IEnumerable<INodeEventSink> sinks,
        FleetRouter router)
    {
        Store = store;
        Options = options;
        Time = time;
        Logger = logger;
        StateProviders = [.. stateProviders];
        Inboxes = inboxes.ToDictionary(i => i.Channel);
        Sinks = [.. sinks];
        Router = router.Inner;
        RoleOf = router.RoleOf;
    }

    internal IFleetStore Store { get; }

    internal FleetServerOptions Options { get; }

    internal TimeProvider Time { get; }

    internal ILogger<FleetLinkHost> Logger { get; }

    internal IReadOnlyList<IStateProvider> StateProviders { get; }

    internal IReadOnlyDictionary<byte, IInbox> Inboxes { get; }

    internal IReadOnlyList<INodeEventSink> Sinks { get; }

    internal RpcRouter<NodeCallContext> Router { get; }

    internal Func<NodeRecord, string?>? RoleOf { get; }

    public IReadOnlyCollection<NodeId> OnlineNodes => [.. _connections.Keys];

    public bool IsOnline(NodeId node) => _connections.ContainsKey(node);

    public NodeSessionInfo? GetSession(NodeId node) =>
        _connections.TryGetValue(node, out NodeConnection? c) ? c.Session : null;

    /// <summary>处理一条 WebSocket 连接，直到连接结束。</summary>
    public async Task RunConnectionAsync(WebSocket socket, CancellationToken ct)
    {
        using var connection = new NodeConnection(this, socket);
        await connection.RunAsync(ct);
    }

    /// <summary>返回被替换掉的旧连接（如果有）。</summary>
    internal NodeConnection? Register(NodeConnection connection)
    {
        NodeConnection? replaced = null;
        _connections.AddOrUpdate(
            connection.NodeId,
            connection,
            (_, existing) =>
            {
                replaced = existing;
                return connection;
            });
        return replaced;
    }

    internal void Unregister(NodeConnection connection) =>
        _connections.TryRemove(new KeyValuePair<NodeId, NodeConnection>(connection.NodeId, connection));

    public async ValueTask<Result<TResult>> RequestAsync<TParams, TResult>(
        NodeId node,
        string method,
        TParams parameters,
        JsonTypeInfo<TParams> paramsInfo,
        JsonTypeInfo<TResult> resultInfo,
        CancellationToken ct)
    {
        if (!_connections.TryGetValue(node, out NodeConnection? connection))
        {
            return Result<TResult>.Fail(FleetErrors.NodeOffline, "node is not connected");
        }

        Result<JsonElement> raw = await connection.RequestAsync(method, JsonSerializer.SerializeToElement(parameters, paramsInfo), ct);
        if (!raw.TryGet(out JsonElement value, out FleetError? error))
        {
            return Result<TResult>.Fail(error);
        }

        try
        {
            TResult? result = value.Deserialize(resultInfo);
            return result is null
                ? Result<TResult>.Fail(FleetErrors.Internal, "empty result")
                : Result<TResult>.Ok(result);
        }
        catch (JsonException)
        {
            return Result<TResult>.Fail(FleetErrors.Internal, "result could not be parsed");
        }
    }

    public async ValueTask<bool> NotifyAsync<TParams>(NodeId node, string method, TParams parameters, JsonTypeInfo<TParams> paramsInfo, CancellationToken ct)
    {
        if (!_connections.TryGetValue(node, out NodeConnection? connection))
        {
            return false;
        }

        return await connection.SendAsync(RpcMessage.Notification(method, parameters, paramsInfo), ct);
    }

    public async ValueTask CloseAllAsync(string reason)
    {
        foreach (NodeConnection connection in _connections.Values.ToArray())
        {
            await connection.CloseAsync(reason);
        }
    }

    public async ValueTask DisconnectAsync(NodeId node, string reason, CancellationToken ct)
    {
        if (_connections.TryGetValue(node, out NodeConnection? connection))
        {
            await connection.RevokeAsync(reason, ct);
        }
    }

    public void NotifyChanged(NodeId node)
    {
        if (_connections.TryGetValue(node, out NodeConnection? connection))
        {
            connection.SignalStateChanged();
        }
    }

    public void NotifyAllChanged()
    {
        foreach (NodeConnection connection in _connections.Values)
        {
            connection.SignalStateChanged();
        }
    }
}

/// <summary>产品注册的 RPC 方法路由。</summary>
public sealed class FleetRouter
{
    public FleetRouter(RpcRouter<NodeCallContext> inner, Func<NodeRecord, string?>? roleOf)
    {
        Inner = inner;
        RoleOf = roleOf;
    }

    internal RpcRouter<NodeCallContext> Inner { get; }

    internal Func<NodeRecord, string?>? RoleOf { get; }

    /// <summary>已注册的入站方法名（只读，测试用它核对声明表）。</summary>
    public IReadOnlyCollection<string> Methods => Inner.Methods;

    public MethodTable Table => Inner.Table;
}