Ananta-ps/Ananta.SDK/Network/TcpServer.cs

using System.Net;
using System.Net.Sockets;
using Ananta.SDK.Logging;
using Ananta.SDK.Rpc;

namespace Ananta.SDK.Network;

public sealed class TcpServer
{
    private readonly string _scope;
    private readonly IPEndPoint _endpoint;
    private readonly RpcRouter _router;
    private readonly Func<TcpSession, Frame, CancellationToken, Task>? _customFrameHandler;
    private readonly Func<TcpSession, Task>? _onClosed;
    private readonly int _heartbeatMs;
    private readonly int _writeTimeoutMs;
    private readonly ServerLogger _log;

    /// <param name="onClosed">
    /// 连接收尾钩子（纯可选）。在 <c>session.Close()</c> **之前** await —— 因为 Close 会
    /// 停掉写侧之后的再发幀，届时已无法通知其它会话。
    /// </param>
    public TcpServer(string scope, string host, int port, RpcRouter router,
        Func<TcpSession, Frame, CancellationToken, Task>? customFrameHandler = null,
        Func<TcpSession, Task>? onClosed = null,
        int heartbeatMs = 0,
        int writeTimeoutMs = 15000)
    {
        _scope = scope;
        _endpoint = new IPEndPoint(IPAddress.Parse(host), port);
        _router = router;
        _customFrameHandler = customFrameHandler;
        _onClosed = onClosed;
        _heartbeatMs = heartbeatMs;
        _writeTimeoutMs = writeTimeoutMs;
        _log = new ServerLogger(scope);
    }

    public async Task RunAsync(CancellationToken token)
    {
        var listener = new TcpListener(_endpoint);
        listener.Start();
        _log.Info($"listening at {_endpoint}");

        while (!token.IsCancellationRequested)
        {
            var client = await listener.AcceptTcpClientAsync(token);
            _ = Task.Run(() => HandleClientAsync(client, token), token);
        }
    }

    private async Task HandleClientAsync(TcpClient client, CancellationToken token)
    {
        var session = new TcpSession(client, _scope, _heartbeatMs, _writeTimeoutMs);
        session.Log.Info($"+ {client.Client.RemoteEndPoint}");

        try
        {
            while (!token.IsCancellationRequested && client.Connected)
            {
                var frame = await session.ReadFrameAsync(token);
                if (frame is null) break;

                if (frame.Value.Mode == 0x02)
                {
                    session.Log.Info($"HS {frame.Value.Payload.Length}b");
                    await session.SendHandshakeAsync(frame.Value.Payload, token);
                    continue;
                }

                if (_customFrameHandler is not null)
                {
                    await _customFrameHandler(session, frame.Value, token);
                    continue;
                }

                await _router.DispatchAsync(session, frame.Value, token);
            }
        }
        catch (Exception ex)
        {
            session.Log.Error("session error", ex);
        }
        finally
        {
            session.Log.Info("-");
            // 收尾钩子必须在 Close 之前跑 —— Close 之后这条会话就再也发不出东西了。
            if (_onClosed is not null)
            {
                try { await _onClosed(session); }
                catch (Exception ex) { session.Log.Warn($"onClosed handler failed: {ex.Message}"); }
            }
            session.Close();
        }
    }
}