diff --git a/PluralKit.Bot/Bot.cs b/PluralKit.Bot/Bot.cs index ad41ac5e..9ceab074 100644 --- a/PluralKit.Bot/Bot.cs +++ b/PluralKit.Bot/Bot.cs @@ -62,7 +62,7 @@ namespace PluralKit.Bot _client.Resumed += args => UpdateBotStatus(args.Client); // Init the shard stuff - _services.Resolve().Init(_client); + _services.Resolve().Init(); // Not awaited, just needs to run in the background // Trying our best to run it at whole minute boundaries (xx:00), with ~250ms buffer @@ -175,7 +175,7 @@ namespace PluralKit.Bot catch (WebSocketException) { } } - private void FrameworkLog(object sender, DebugLogMessageEventArgs args) + public void FrameworkLog(object sender, DebugLogMessageEventArgs args) { // Bridge D#+ logging to Serilog LogEventLevel level = LogEventLevel.Verbose; diff --git a/PluralKit.Bot/Init.cs b/PluralKit.Bot/Init.cs index a9f9363b..f7d4aacf 100644 --- a/PluralKit.Bot/Init.cs +++ b/PluralKit.Bot/Init.cs @@ -37,14 +37,16 @@ namespace PluralKit.Bot logger.Information("Connecting to database"); await services.Resolve().ApplyMigrations(); - // Start the Discord client; StartAsync returns once shard instances are *created* (not necessarily connected) - logger.Information("Connecting to Discord"); - await services.Resolve().StartAsync(); - - // Start the bot stuff and let it register things + // Init the bot instance itself, register handlers and such to the client before beginning to connect + logger.Information("Initializing bot"); var bot = services.Resolve(); bot.Init(); - + + // Start the Discord shards themselves (handlers already set up) + logger.Information("Connecting to Discord"); + await services.Resolve().StartAsync(); + logger.Information("Connected! All is good (probably)."); + // Lastly, we just... wait. Everything else is handled in the DiscordClient event loop try { diff --git a/PluralKit.Bot/Services/ShardInfoService.cs b/PluralKit.Bot/Services/ShardInfoService.cs index 125fcf30..08cd392e 100644 --- a/PluralKit.Bot/Services/ShardInfoService.cs +++ b/PluralKit.Bot/Services/ShardInfoService.cs @@ -13,6 +13,7 @@ namespace PluralKit.Bot { public class ShardInfoService { + public class ShardInfo { public Instant LastConnectionTime; @@ -23,24 +24,38 @@ namespace PluralKit.Bot } private ILogger _logger; + private DiscordShardedClient _client; private Dictionary _shardInfo = new Dictionary(); - public ShardInfoService(ILogger logger) + public ShardInfoService(ILogger logger, DiscordShardedClient client) { + _client = client; _logger = logger.ForContext(); } - public void Init(DiscordShardedClient client) + public void Init() { - foreach (var (shardId, shard) in client.ShardClients) - { - _shardInfo[shardId] = new ShardInfo(); + // We initialize this before any shards are actually created and connected + // This means the client won't know the shard count, so we attach a listener every time a shard gets connected + _client.SocketOpened += RefreshShardList; + } - shard.Heartbeated += Heartbeated; - shard.SocketClosed += SocketClosed; - shard.Ready += Ready; - shard.Resumed += Resumed; + private async Task RefreshShardList() + { + // This callback doesn't actually receive the shard that was opening, so we just try to check we have 'em all (so far) + foreach (var (id, shard) in _client.ShardClients) + { + if (_shardInfo.ContainsKey(id)) continue; + // Call our own SocketOpened listener manually (and then attach the listener properly) + await SocketOpened(shard); shard.SocketOpened += () => SocketOpened(shard); + + // Register listeners for new shards + _logger.Information("Attaching listeners to new shard #{Shard}", shard.ShardId); + shard.Resumed += Resumed; + shard.Ready += Ready; + shard.SocketClosed += SocketClosed; + shard.Heartbeated += Heartbeated; } } @@ -52,9 +67,11 @@ namespace PluralKit.Bot return Task.CompletedTask; } - private ShardInfo UpdateShard(DiscordClient shard) + private ShardInfo TryGetShard(DiscordClient shard) { // If we haven't seen this shard before, add it to the dict! + // I don't think this will ever occur since the shard number is constant up-front and we handle those + // in the RefreshShardList handler above but you never know, I guess~ if (!_shardInfo.TryGetValue(shard.ShardId, out var info)) _shardInfo[shard.ShardId] = info = new ShardInfo(); return info; @@ -64,7 +81,7 @@ namespace PluralKit.Bot { _logger.Information("Shard #{Shard} resumed connection", e.Client.ShardId); - var info = UpdateShard(e.Client); + var info = TryGetShard(e.Client); // info.LastConnectionTime = SystemClock.Instance.GetCurrentInstant(); info.Connected = true; return Task.CompletedTask; @@ -74,7 +91,7 @@ namespace PluralKit.Bot { _logger.Information("Shard #{Shard} sent Ready event", e.Client.ShardId); - var info = UpdateShard(e.Client); + var info = TryGetShard(e.Client); info.LastConnectionTime = SystemClock.Instance.GetCurrentInstant(); info.Connected = true; return Task.CompletedTask; @@ -84,7 +101,7 @@ namespace PluralKit.Bot { _logger.Warning("Shard #{Shard} disconnected ({CloseCode}: {CloseMessage})", e.Client.ShardId, e.CloseCode, e.CloseMessage); - var info = UpdateShard(e.Client); + var info = TryGetShard(e.Client); info.DisconnectionCount++; info.Connected = false; return Task.CompletedTask; @@ -95,7 +112,7 @@ namespace PluralKit.Bot var latency = Duration.FromMilliseconds(e.Ping); _logger.Information("Shard #{Shard} received heartbeat (latency: {Latency} ms)", e.Client.ShardId, latency.Milliseconds); - var info = UpdateShard(e.Client); + var info = TryGetShard(e.Client); info.LastHeartbeatTime = e.Timestamp.ToInstant(); info.Connected = true; info.ShardLatency = latency;