MoShiro/backend/fleet/src/Fleet.Server.Link/NodeConnection.Send.cs

using System.Net.WebSockets;
using System.Text.Json;
using Fleet.Core;
using Fleet.Rpc;

namespace Fleet.Server;

/// <summary>发给节点：关闭、吊销、发送一条消息、发起请求并等响应。</summary>
internal sealed partial class NodeConnection
{
    /// <summary>
    /// 服务端发起的关闭：发送 Close 帧，等对端确认之后接收循环自然结束。
    /// 注意不能直接取消接收的令牌——取消 <c>ReceiveAsync</c> 会中止整个 WebSocket，对端拿不到关闭原因；取消只作为超时兜底。
    /// </summary>
    public async ValueTask CloseAsync(string reason)
    {
        _closeReason ??= reason;
        try
        {
            await _sendLock.WaitAsync(CancellationToken.None);
        }
        catch (ObjectDisposedException)
        {
            return;
        }

        try
        {
            if (_disposed)
            {
                return;
            }

            _lifetime.CancelAfter(s_closeTimeout);
            if (_socket.State == WebSocketState.Open)
            {
                await _socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, reason, CancellationToken.None);
            }
        }
        catch (WebSocketException)
        {
            // 对端已经断开
        }
        finally
        {
            _sendLock.Release();
        }

    }

    public async ValueTask RevokeAsync(string reason, CancellationToken ct)
    {
        await SendAsync(RpcMessage.Notification(FleetMethods.Revoked, new RevokedParams { Reason = reason }, FleetRpcContext.Default.RevokedParams), ct);
        await CloseAsync(reason);
    }

    public async ValueTask<bool> SendAsync(RpcMessage message, CancellationToken ct)
    {
        if (message.Method is not null)
        {
            // 服务端只能发自己有资格发起的方法，并且对端的角色要允许（响应没有方法名，不检查）
            string? peerRole = _node is null ? null : _host.RoleOf?.Invoke(_node);
            _host.Router.CheckSend(message.Method, message.Id is null ? MethodKind.Notification : MethodKind.Request, peerRole);
        }

        byte[] bytes = JsonSerializer.SerializeToUtf8Bytes(message, FleetRpcContext.Default.RpcMessage);
        try
        {
            await _sendLock.WaitAsync(ct);
        }
        catch (ObjectDisposedException)
        {
            return false;
        }

        try
        {
            if (_disposed || _socket.State != WebSocketState.Open)
            {
                return false;
            }

            if (_cipher is { } cipher)
            {
                await _socket.SendAsync(cipher.Seal(0, bytes), WebSocketMessageType.Binary, endOfMessage: true, ct);
            }
            else
            {
                await _socket.SendAsync(bytes, WebSocketMessageType.Text, endOfMessage: true, ct);
            }

            return true;
        }
        catch (WebSocketException)
        {
            return false;
        }
        finally
        {
            _sendLock.Release();
        }
    }

    public async ValueTask<Result<JsonElement>> RequestAsync(string method, JsonElement parameters, CancellationToken ct)
    {
        long id = Interlocked.Increment(ref _nextRequestId);
        var waiter = new TaskCompletionSource<RpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
        _pending[id] = waiter;
        try
        {
            var request = new RpcMessage { Id = id, Method = method, Params = parameters };
            if (!await SendAsync(request, ct))
            {
                return Result<JsonElement>.Fail(FleetErrors.NodeOffline, "node is not connected");
            }

            using var timeout = new CancellationTokenSource(_host.Options.RequestTimeout, _host.Time);
            using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, timeout.Token);
            RpcMessage response = await waiter.Task.WaitAsync(linked.Token);
            return response.Error is not null
                ? Result<JsonElement>.Fail(response.Error.ToFleetError())
                : Result<JsonElement>.Ok(response.Result ?? default);
        }
        catch (OperationCanceledException) when (!ct.IsCancellationRequested)
        {
            return Result<JsonElement>.Fail(FleetErrors.Timeout, $"{method} timed out");
        }
        finally
        {
            _pending.TryRemove(id, out _);
        }
    }
}