PluralKit/Myriad/Gateway/Shard.cs

224 lines
7.7 KiB
C#

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.Identifying)
await Task.Delay(50);
}
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);
}
}
}