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

using System.Diagnostics.CodeAnalysis;
using Fleet.Rpc;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Routing;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;

namespace Fleet.Server;

/// <summary>Fleet 服务端的接入点：产品在这里挂上自己的状态、收件箱、事件和 RPC 方法。</summary>
public sealed class FleetServerBuilder
{
    private readonly List<Action<RpcRouter<NodeCallContext>, IServiceProvider>> _routeConfigurators = [];
    private MethodTable? _methods;
    private Func<NodeRecord, string?>? _roleOf;

    internal FleetServerBuilder(IServiceCollection services)
    {
        Services = services;
    }

    public IServiceCollection Services { get; }

    internal IReadOnlyList<Action<RpcRouter<NodeCallContext>, IServiceProvider>> RouteConfigurators => _routeConfigurators;

    internal MethodTable Methods => _methods ?? throw new InvalidOperationException("call WithMethods(...) when adding the Fleet server: the method table is required");

    internal Func<NodeRecord, string?>? RoleOf => _roleOf;

    /// <summary>
    /// 方法声明表（框架的加产品的，见 <see cref="MethodTable.Combine"/>）和节点的角色解析（例如按节点 ID 前缀区分 Agent 和 Probe）。必须设置：
    /// 路由器据此核对每个处理函数的方向、类型和角色，并拒绝没有声明或不该由节点发起的方法。
    /// </summary>
    public FleetServerBuilder WithMethods(MethodTable methods, Func<NodeRecord, string?> roleOf)
    {
        _methods = methods;
        _roleOf = roleOf;
        return this;
    }

    public FleetServerBuilder AddState<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors)] TProvider>()
        where TProvider : class, IStateProvider
    {
        Services.AddSingleton<IStateProvider, TProvider>();
        return this;
    }

    public FleetServerBuilder AddInbox<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors)] TInbox>()
        where TInbox : class, IInbox
    {
        Services.AddSingleton<IInbox, TInbox>();
        return this;
    }

    public FleetServerBuilder OnNodeEvent<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors)] TSink>()
        where TSink : class, INodeEventSink
    {
        Services.AddSingleton<INodeEventSink, TSink>();
        return this;
    }

    public FleetServerBuilder OnEnrolled<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors)] THook>()
        where THook : class, IEnrollmentHook
    {
        Services.AddSingleton<IEnrollmentHook, THook>();
        return this;
    }

    /// <summary>注册产品的 RPC 方法（<c>moshiro.*</c> 等）。处理函数通过 <see cref="IServiceProvider"/> 在调用时解析服务，避免构造期的循环依赖。</summary>
    public FleetServerBuilder MapMethods(Action<RpcRouter<NodeCallContext>, IServiceProvider> configure)
    {
        _routeConfigurators.Add(configure);
        return this;
    }
}

public static class FleetServerExtensions
{
    /// <summary>需要事先在 DI 中注册 <see cref="IFleetStore"/>。</summary>
    public static FleetServerBuilder AddFleetServer(this IServiceCollection services, Action<FleetServerOptions>? configure = null)
    {
        var options = new FleetServerOptions();
        configure?.Invoke(options);
        services.AddSingleton(options);

        var builder = new FleetServerBuilder(services);
        services.AddSingleton(provider =>
        {
            ILogger logger = provider.GetRequiredService<ILoggerFactory>().CreateLogger("Fleet.Rpc");
            Func<NodeRecord, string?>? roleOf = builder.RoleOf;
            var router = new RpcRouter<NodeCallContext>(
                RpcSide.Server,
                builder.Methods,
                callerRole: roleOf is null ? null : context => roleOf(context.Node),
                onUnhandled: (method, ex) => FleetLog.HandlerFailed(logger, ex, method));
            foreach (Action<RpcRouter<NodeCallContext>, IServiceProvider> configurator in builder.RouteConfigurators)
            {
                configurator(router, provider);
            }

            return new FleetRouter(router, roleOf);
        });
        services.AddSingleton(provider => new FleetLinkHost(
            provider.GetRequiredService<IFleetStore>(),
            provider.GetRequiredService<FleetServerOptions>(),
            provider.GetRequiredService<TimeProvider>(),
            provider.GetRequiredService<ILogger<FleetLinkHost>>(),
            provider.GetServices<IStateProvider>(),
            provider.GetServices<IInbox>(),
            provider.GetServices<INodeEventSink>(),
            provider.GetRequiredService<FleetRouter>(),
            provider.GetService<IServerIdentity>(),
            provider.GetService<IBinaryAttestor>()));
        services.AddSingleton<IFleetConnections>(provider => provider.GetRequiredService<FleetLinkHost>());
        services.AddSingleton<IStateNotifier>(provider => provider.GetRequiredService<FleetLinkHost>());
        return builder;
    }

    /// <summary>映射 WebSocket 端点。调用方需要先 <c>app.UseWebSockets()</c>。</summary>
    public static IEndpointRouteBuilder MapFleetLink(this IEndpointRouteBuilder app, string path)
    {
        app.Map(path, async (HttpContext http) =>
        {
            if (!http.WebSockets.IsWebSocketRequest)
            {
                http.Response.StatusCode = StatusCodes.Status400BadRequest;
                return;
            }

            using System.Net.WebSockets.WebSocket socket = await http.WebSockets.AcceptWebSocketAsync();
            await http.RequestServices.GetRequiredService<FleetLinkHost>().RunConnectionAsync(socket, http.RequestAborted);
        });
        return app;
    }
}