DiscordApi/Discord.API/Gateway/GatewayClient.cs

using System.Collections.Specialized;
using System.Security.Cryptography;
using System.Text.Json;
using Serilog;
using Websocket.Client;

namespace Discord.API;

public class GatewayClient {

    private const int ApiVersion = 10;

    private string ApiKey;
    private IWebsocketClient Websocket;
    private CancellationTokenSource? HeartbeatCts;
    private ulong? Sequence = null;
    private DateTime? HeartbeatAckReceived = null;
    private TimeProvider timeProvider;

    public GatewayClient(string url, string api_key, TimeProvider time_provider)
        : this(url, api_key, time_provider, uri => new WebsocketClient(uri))
    { }

    internal GatewayClient(string url, string api_key, TimeProvider time_provider, Func<Uri, IWebsocketClient> websocket_client_factory){
        this.timeProvider=time_provider;
        this.ApiKey=api_key;
        UriBuilder uriBuilder = new(url);
        NameValueCollection query = System.Web.HttpUtility.ParseQueryString("");
        query.Add("v", ApiVersion.ToString());
        query.Add("encoding", "json");
        uriBuilder.Query=query.ToString();
        Websocket = websocket_client_factory.Invoke(uriBuilder.Uri);
        Websocket.DisconnectionHappened.Subscribe(DisconnectionHandler);
        Websocket.ReconnectionHappened.Subscribe(ReconnectionHandler);
        Websocket.MessageReceived.Subscribe(MessageReceivedHandler);
        Websocket.Start();

        Log.Debug("GATEWAY: Created new gateway, with url: {url}", uriBuilder.ToString());
    }

    private void DisconnectionHandler(DisconnectionInfo info){
        StopHeartbeat();
        Log.Information("GATEWAY: Disconnected. Type: {DisconnectionType}", info.Type);
    }
    private void ReconnectionHandler(ReconnectionInfo info){
        Log.Information("GATEWAY: (Re)Connected to server. Url: {url}, Type: {Type}", Websocket.Url, info.Type);
    }
    private void MessageReceivedHandler(ResponseMessage msg){
        if(msg.MessageType != System.Net.WebSockets.WebSocketMessageType.Text) return;
        try{
            GatewayPacket packet = JsonSerializer.Deserialize(msg.Text!, SourceGenerationContext.Default.GatewayPacket)
                ?? throw new Exception("Failed to deserialize packet"); // This can be optimized //TODO

            Log.Debug("GATEWAY: Packet received");

            switch(packet){
                case HelloPacket helloPacket:
                    HelloPacketHandler(helloPacket);
                    break;
                case HeartbeatAckPacket:
                    HeartbeatAckHandler();
                    break;
                default:
                    Log.Debug("GATEWAY: Packet not handled");
                    break;
            }
        }catch(Exception ex){
            Log.Warning(ex, "GATEWAY: Error processing gateway event");
        }
    }

    private void HelloPacketHandler(HelloPacket packet){
        StartHeartbeat((int)packet.Data.HeartbeatInterval);
        Log.Debug("GATEWAY: Hello packet received, heartbeat interval: {heartbeat_ms}", packet.Data.HeartbeatInterval);
    }

    private void HeartbeatAckHandler(){
        HeartbeatAckReceived = timeProvider.GetUtcNow().DateTime;
        Log.Debug("GATEWAY: Heartbeat ACK received");
    }

    private void StartHeartbeat(int heartbeat_interval){
        HeartbeatCts?.Cancel();
        HeartbeatCts = new CancellationTokenSource();
        CancellationToken ct = HeartbeatCts.Token;
        Task.Run(async ()=>{
            await Task.Delay(Random.Shared.Next(1, heartbeat_interval));
            using PeriodicTimer pd = new(TimeSpan.FromMilliseconds(heartbeat_interval), timeProvider);
            HeartbeatAckReceived = timeProvider.GetUtcNow().DateTime;
            DateTime HeartbeatSent = timeProvider.GetUtcNow().DateTime;
            do{
                if(HeartbeatAckReceived == null){
                    Log.Debug("GATEWAY: Heartbeat ack not received. Reconnecting.");
                    _ = Websocket.Reconnect();
                    break;
                }
                Log.Information("GATEWAY: Heartbeat ping time is {time_ms} ms", (HeartbeatAckReceived.Value - HeartbeatSent).TotalMilliseconds);
                HeartbeatPacket packet = new HeartbeatPacket(){
                    Sequence=this.Sequence
                };
                if(Websocket.IsRunning)
                    if(!Websocket.Send(JsonSerializer.Serialize(packet, SourceGenerationContext.Default.HeartbeatPacket))){
                        Log.Warning("GATEWAY: Failed to queue heartbeat message");
                    }
                HeartbeatSent = timeProvider.GetUtcNow().DateTime;
                HeartbeatAckReceived = null;
                Log.Debug("GATEWAY: Heartbeat sent");
            }while(await pd.WaitForNextTickAsync(ct) && !ct.IsCancellationRequested);
        });
    }

    private void StopHeartbeat(){
        HeartbeatCts?.Cancel();
    }
}