global using Discord; global using Discord.WebSocket; using Discord.Interactions; using Microsoft.Extensions.DependencyInjection; using System.Text; namespace BirthdayBot; /// /// More or less the main class for the program. Handles individual shards and provides frequent /// status reports regarding the overall health of the application. /// class ShardManager : IDisposable { /// /// Number of seconds between each time the status task runs, in seconds. /// private const int StatusInterval = 90; /// /// Number of concurrent shard startups to happen on each check. /// This value also determines the maximum amount of concurrent background database operations. /// public const int MaxConcurrentOperations = 4; /// /// Amount of time without a completed background service run before a shard instance /// is considered "dead" and tasked to be removed. /// public static readonly TimeSpan DeadShardThreshold = new(0, 20, 0); /// /// A dictionary with shard IDs as its keys and shard instances as its values. /// When initialized, all keys will be created as configured. If an instance is removed, /// a key's corresponding value will temporarily become null instead of the key/value /// pair being removed. /// private readonly Dictionary _shards; private readonly Task _statusTask; private readonly CancellationTokenSource _mainCancel; internal Configuration Config { get; } public ShardManager(Configuration cfg) { var ver = System.Reflection.Assembly.GetExecutingAssembly().GetName().Version; Log($"Birthday Bot v{ver!.ToString(3)} is starting..."); Config = cfg; // Allocate shards based on configuration _shards = new Dictionary(); for (var i = Config.ShardStart; i < (Config.ShardStart + Config.ShardAmount); i++) { _shards.Add(i, null); } // Start status reporting thread _mainCancel = new CancellationTokenSource(); _statusTask = Task.Factory.StartNew(StatusLoop, _mainCancel.Token, TaskCreationOptions.LongRunning, TaskScheduler.Default); } public void Dispose() { _mainCancel.Cancel(); _statusTask.Wait(10000); if (!_statusTask.IsCompleted) Log("Warning: Main thread did not cleanly finish up in time. Continuing..."); Log("Shutting down all shards..."); var shardDisposes = new List(); foreach (var item in _shards) { if (item.Value == null) continue; shardDisposes.Add(Task.Run(item.Value.Dispose)); } if (!Task.WhenAll(shardDisposes).Wait(30000)) { Log("Warning: Not all shards terminated cleanly after 30 seconds. Continuing..."); } Log($"Uptime: {Program.BotUptime}"); } private void Log(string message) => Program.Log(nameof(ShardManager), message); /// /// Creates and sets up a new shard instance. /// private async Task InitializeShard(int shardId) { var clientConf = new DiscordSocketConfig() { ShardId = shardId, TotalShards = Config.ShardTotal, LogLevel = LogSeverity.Info, DefaultRetryMode = RetryMode.Retry502 | RetryMode.RetryTimeouts, GatewayIntents = GatewayIntents.Guilds | GatewayIntents.GuildMembers, SuppressUnknownDispatchWarnings = true, LogGatewayIntentWarnings = false }; var services = new ServiceCollection() .AddSingleton(s => new ShardInstance(this, s)) .AddSingleton(s => new DiscordSocketClient(clientConf)) .AddSingleton(s => new InteractionService(s.GetRequiredService())) .BuildServiceProvider(); var newInstance = services.GetRequiredService(); await newInstance.StartAsync(); return newInstance; } public int? GetShardIdFor(ulong guildId) { foreach (var sh in _shards.Values) { if (sh == null) continue; if (sh.DiscordClient.GetGuild(guildId) != null) return sh.ShardId; } return null; } private async Task StatusLoop() { try { while (!_mainCancel.IsCancellationRequested) { Log($"Uptime: {Program.BotUptime}"); // Iterate through shards, create report on each var shardStatuses = new StringBuilder(); var nullShards = new List(); var deadShards = new List(); foreach (var i in _shards.Keys) { shardStatuses.Append($"Shard {i:00}: "); if (_shards[i] == null) { shardStatuses.AppendLine("Inactive."); nullShards.Add(i); continue; } var shard = _shards[i]!; var client = shard.DiscordClient; shardStatuses.Append($"{Enum.GetName(typeof(ConnectionState), client.ConnectionState)} ({client.Latency:000}ms)."); shardStatuses.Append($" Guilds: {client.Guilds.Count}."); shardStatuses.Append($" Background: {shard.CurrentExecutingService ?? "Idle"}"); var lastRun = DateTimeOffset.UtcNow - shard.LastBackgroundRun; if (lastRun > DeadShardThreshold / 3) { // Formerly known as a 'slow' shard shardStatuses.Append($", heartbeat {lastRun.TotalMinutes:00.0}m ago."); } else { shardStatuses.Append('.'); } shardStatuses.AppendLine(); if (lastRun > DeadShardThreshold) { shardStatuses.AppendLine($"Shard {i:00} marked for disposal."); deadShards.Add(i); } } Log(shardStatuses.ToString().TrimEnd()); // Remove dead shards foreach (var dead in deadShards) { _shards[dead]!.Dispose(); _shards[dead] = null; } // Start null shards, a few at at time var startAllowance = MaxConcurrentOperations; foreach (var id in nullShards) { if (startAllowance-- > 0) { _shards[id] = await InitializeShard(id); } else break; } await Task.Delay(StatusInterval * 1000, _mainCancel.Token); } } catch (TaskCanceledException) { } } }