PluralKit/Myriad/Gateway/Shard.cs

using System;
using System.Net.WebSockets;
using System.Text.Json;
using System.Threading.Tasks;

using Myriad.Gateway.Limit;
using Myriad.Gateway.State;
using Myriad.Serialization;
using Myriad.Types;

using Serilog;
using Serilog.Context;

namespace Myriad.Gateway
{
    public class Shard
    {
        private const string LibraryName = "Myriad (for PluralKit)";

        private readonly GatewaySettings _settings;
        private readonly ShardInfo _info;
        private readonly IGatewayRatelimiter _ratelimiter;
        private readonly string _url;
        private readonly ILogger _logger;
        private readonly ShardStateManager _stateManager;
        private readonly JsonSerializerOptions _jsonSerializerOptions;
        private readonly ShardConnection _conn;

        public int ShardId => _info.ShardId;
        public ShardState State => _stateManager.State;
        public TimeSpan? Latency => _stateManager.Latency;
        public User? User => _stateManager.User;
        public ApplicationPartial? Application => _stateManager.Application;

        // TODO: I wanna get rid of these or move them at some point
        public event Func<IGatewayEvent, Task>? OnEventReceived;
        public event Action<TimeSpan>? HeartbeatReceived;
        public event Action? SocketOpened;
        public event Action? Resumed;
        public event Action? Ready;
        public event Action<WebSocketCloseStatus?, string?>? SocketClosed;

        private TimeSpan _reconnectDelay = TimeSpan.Zero;
        private Task? _worker;

        public Shard(GatewaySettings settings, ShardInfo info, IGatewayRatelimiter ratelimiter, string url, ILogger logger)
        {
            _jsonSerializerOptions = new JsonSerializerOptions().ConfigureForMyriad();

            _settings = settings;
            _info = info;
            _ratelimiter = ratelimiter;
            _url = url;
            _logger = logger.ForContext<Shard>().ForContext("ShardId", info.ShardId);
            _stateManager = new ShardStateManager(info, _jsonSerializerOptions, logger)
            {
                HandleEvent = HandleEvent,
                SendHeartbeat = SendHeartbeat,
                SendIdentify = SendIdentify,
                SendResume = SendResume,
                Connect = ConnectInner,
                Reconnect = Reconnect,
            };
            _stateManager.OnHeartbeatReceived += latency =>
            {
                HeartbeatReceived?.Invoke(latency);
            };

            _conn = new ShardConnection(_jsonSerializerOptions, _logger);
        }
        
        private async Task ShardLoop()
        {
            // may be superfluous but this adds shard id to ambient context which is nice
            using var _ = LogContext.PushProperty("ShardId", _info.ShardId);
            
            while (true)
            {
                try
                {
                    await ConnectInner();
                    
                    await HandleConnectionOpened();

                    while (_conn.State == WebSocketState.Open)
                    {
                        var packet = await _conn.Read();
                        if (packet == null)
                            break;

                        await _stateManager.HandlePacketReceived(packet);
                    }
                    
                    await HandleConnectionClosed(_conn.CloseStatus, _conn.CloseStatusDescription);

                    _logger.Information("Shard {ShardId}: Reconnecting after delay {ReconnectDelay}",
                        _info.ShardId, _reconnectDelay);

                    if (_reconnectDelay > TimeSpan.Zero)
                        await Task.Delay(_reconnectDelay);
                    _reconnectDelay = TimeSpan.Zero;
                }
                catch (Exception e)
                {
                    _logger.Error(e, "Shard {ShardId}: Error in main shard loop, reconnecting in 5 seconds...", _info.ShardId);
                    
                    // todo: exponential backoff here? this should never happen, ideally...
                    await Task.Delay(TimeSpan.FromSeconds(5));
                }
            }
        }
        
        public async Task Start()
        {
            if (_worker == null)
                _worker = ShardLoop();

            // we can probably TCS this instead of spin loop but w/e
            while (State != ShardState.Connected) 
                await Task.Delay(100);
        }

        public async Task UpdateStatus(GatewayStatusUpdate payload)
        {
            await _conn.Send(new GatewayPacket
            {
                Opcode = GatewayOpcode.PresenceUpdate,
                Payload = payload
            });
        }
        
        private async Task ConnectInner()
        {
            while (true)
            {
                await _ratelimiter.Identify(_info.ShardId);

                _logger.Information("Shard {ShardId}: Connecting to WebSocket", _info.ShardId);
                try
                {
                    await _conn.Connect(_url, default);
                    break;
                }
                catch (WebSocketException e)
                {
                    _logger.Error(e, "Shard {ShardId}: Error connecting to WebSocket, retrying in 5 seconds...", _info.ShardId);
                    await Task.Delay(TimeSpan.FromSeconds(5));
                }
            }
        }
        
        private async Task DisconnectInner(WebSocketCloseStatus closeStatus)
        {
            await _conn.Disconnect(closeStatus, null);
        }
        
        private async Task SendIdentify()
        {
            await _conn.Send(new GatewayPacket
            {
                Opcode = GatewayOpcode.Identify,
                Payload = new GatewayIdentify
                {
                    Compress = false,
                    Intents = _settings.Intents,
                    Properties = new GatewayIdentify.ConnectionProperties
                    {
                        Browser = LibraryName, 
                        Device = LibraryName, 
                        Os = Environment.OSVersion.ToString()
                    },
                    Shard = _info,
                    Token = _settings.Token,
                    LargeThreshold = 50
                }
            });
        }
        
        private async Task SendResume((string SessionId, int? LastSeq) arg)
        {
            await _conn.Send(new GatewayPacket
            {
                Opcode = GatewayOpcode.Resume,
                Payload = new GatewayResume(_settings.Token, arg.SessionId, arg.LastSeq ?? 0)
            });
        }
        
        private async Task SendHeartbeat(int? lastSeq)
        {
            await _conn.Send(new GatewayPacket {Opcode = GatewayOpcode.Heartbeat, Payload = lastSeq});
        }
        
        private async Task Reconnect(WebSocketCloseStatus closeStatus, TimeSpan delay)
        {
            _reconnectDelay = delay;
            await DisconnectInner(closeStatus);
        }

        private async Task HandleEvent(IGatewayEvent arg)
        {
            if (arg is ReadyEvent)
                Ready?.Invoke();
            if (arg is ResumedEvent)
                Resumed?.Invoke();
            
            await (OnEventReceived?.Invoke(arg) ?? Task.CompletedTask);
        }

        private async Task HandleConnectionOpened()
        {
            _logger.Information("Shard {ShardId}: Connection opened", _info.ShardId);
            await _stateManager.HandleConnectionOpened();
            SocketOpened?.Invoke();
        }

        private async Task HandleConnectionClosed(WebSocketCloseStatus? closeStatus, string? description)
        {
            _logger.Information("Shard {ShardId}: Connection closed ({CloseStatus}/{Description})", 
                _info.ShardId, closeStatus, description ?? "<null>");
            await _stateManager.HandleConnectionClosed();
            SocketClosed?.Invoke(closeStatus, description);
        }
    }
}