using BirthdayBot.BackgroundServices; using BirthdayBot.UserInterface; using Discord; using Discord.WebSocket; using System; using System.Collections.Generic; using System.Linq; using System.Net.Http; using System.Text; using System.Threading; using System.Threading.Tasks; using static BirthdayBot.UserInterface.CommandsCommon; namespace BirthdayBot { /// /// The highest level part of this bot: /// Starts up, looks over, and manages shard instances while containing common resources /// and providing common functions for all existing shards. /// class ShardManager : IDisposable { /// /// Number of seconds between each time the manager's watchdog task runs, in seconds. /// private const int WatchdogInterval = 90; /// /// 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; // Commonly used command handler instances private readonly Dictionary _dispatchCommands; private readonly UserCommands _cmdsUser; private readonly ListingCommands _cmdsListing; private readonly HelpInfoCommands _cmdsHelp; private readonly ManagerCommands _cmdsMods; // Watchdog stuff private readonly Task _watchdogTask; private readonly CancellationTokenSource _watchdogCancel; private int _destroyedShards = 0; 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; // Command handler setup _dispatchCommands = new Dictionary(StringComparer.OrdinalIgnoreCase); _cmdsUser = new UserCommands(cfg); foreach (var item in _cmdsUser.Commands) _dispatchCommands.Add(item.Item1, item.Item2); _cmdsListing = new ListingCommands(cfg); foreach (var item in _cmdsListing.Commands) _dispatchCommands.Add(item.Item1, item.Item2); _cmdsHelp = new HelpInfoCommands(cfg); foreach (var item in _cmdsHelp.Commands) _dispatchCommands.Add(item.Item1, item.Item2); _cmdsMods = new ManagerCommands(cfg, _cmdsUser.Commands); foreach (var item in _cmdsMods.Commands) _dispatchCommands.Add(item.Item1, item.Item2); _shards = new Dictionary(); // TODO implement more flexible sharding configuration here for (int i = 0; i < Config.ShardCount; i++) { _shards.Add(i, null); } // Start watchdog _watchdogCancel = new CancellationTokenSource(); _watchdogTask = Task.Factory.StartNew(WatchdogLoop, _watchdogCancel.Token, TaskCreationOptions.LongRunning, TaskScheduler.Default); } public void Dispose() { Log("Captured cancel key. Shutting down shard status watcher..."); _watchdogCancel.Cancel(); _watchdogTask.Wait(5000); if (!_watchdogTask.IsCompleted) Log("Warning: Shard status watcher has not ended 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(60000)) { Log("Warning: All shards did not properly stop after 60 seconds. Continuing..."); } Log($"Shutdown complete. Bot uptime: {Common.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) { ShardInstance newInstance; var clientConf = new DiscordSocketConfig() { ShardId = shardId, TotalShards = Config.ShardCount, LogLevel = LogSeverity.Info, DefaultRetryMode = RetryMode.RetryRatelimit, MessageCacheSize = 0, // not needed at all ExclusiveBulkDelete = true, // not relevant, but this is configured to skip the warning GatewayIntents = GatewayIntents.Guilds | GatewayIntents.GuildMembers | GatewayIntents.GuildMessages }; var newClient = new DiscordSocketClient(clientConf); newInstance = new ShardInstance(this, newClient, _dispatchCommands); await newInstance.StartAsync().ConfigureAwait(false); return newInstance; } private async Task WatchdogLoop() { try { while (!_watchdogCancel.IsCancellationRequested) { Log($"Bot uptime: {Common.BotUptime}"); // Iterate through shard list, extract data var guildInfo = new Dictionary(); var now = DateTimeOffset.UtcNow; ulong? botId = null; var nullShards = new List(); foreach (var item in _shards) { if (item.Value == null) { nullShards.Add(item.Key); continue; } var shard = item.Value; botId ??= shard.DiscordClient.CurrentUser?.Id; var guildCount = shard.DiscordClient.Guilds.Count; var connScore = shard.ConnectionScore; var lastRun = now - shard.LastBackgroundRun; guildInfo[item.Key] = (guildCount, connScore, lastRun); } // Process info var guildCounts = guildInfo.Select(i => i.Value.Item1); var guildTotal = guildCounts.Sum(); var guildAverage = guildCounts.Any() ? guildCounts.Average() : 0; Log($"Currently in {guildTotal} guilds. Average shard load: {guildAverage:0.0}."); if (nullShards.Count == 0 && botId.HasValue) await SendExternalStatistics(guildTotal, botId.Value, _watchdogCancel.Token).ConfigureAwait(false); // Health report var goodShards = new List(); var badShards = new List(); // shards with low connection score OR long time since last work var deadShards = new List(); // shards to destroy and reinitialize foreach (var item in guildInfo) { var connScore = item.Value.Item2; var lastRun = item.Value.Item3; if (lastRun > new TimeSpan(0, 10, 0) || connScore < ConnectionStatus.StableScore) { badShards.Add(item.Key); // Consider a shard dead after a long span without background activity if (lastRun > new TimeSpan(0, 30, 0)) deadShards.Add(item.Key); } else { goodShards.Add(item.Key); } } string statusDisplay(IEnumerable list, bool detailedInfo) { if (!list.Any()) return "--"; var result = new StringBuilder(); foreach (var item in list) { result.Append(item.ToString("00") + " "); if (detailedInfo) { result.Remove(result.Length - 1, 1); result.Append($"[{guildInfo[item].Item2:+000;-000}"); result.Append($" {Math.Floor(guildInfo[item].Item3.TotalMinutes):00}m"); result.Append($"{guildInfo[item].Item3.Seconds:00}s] "); } } if (result.Length > 0) result.Remove(result.Length - 1, 1); return result.ToString(); } Log("Stable shards: " + statusDisplay(goodShards, false)); if (badShards.Count > 0) Log("Unstable shards: " + statusDisplay(badShards, true)); if (deadShards.Count > 0) Log("Shards to be restarted: " + statusDisplay(deadShards, false)); if (nullShards.Count > 0) Log("Inactive shards: " + statusDisplay(nullShards, false)); // Remove dead shards foreach (var dead in deadShards) { _shards[dead].Dispose(); _shards[dead] = null; } // Start up any missing shards int startAllowance = 4; foreach (var id in nullShards) { // To avoid possible issues with resources strained over so many shards starting at once, // initialization is spread out by only starting a few at a time. if (startAllowance-- > 0) { _shards[id] = await InitializeShard(id).ConfigureAwait(false); } else break; } // All done for now await Task.Delay(WatchdogInterval * 1000, _watchdogCancel.Token).ConfigureAwait(false); } } catch (TaskCanceledException) { } } #region Statistical reporting private static readonly HttpClient _httpClient = new HttpClient(); /// /// Send statistical information to external services. /// private async Task SendExternalStatistics(int count, ulong botId, CancellationToken token) { var dbotsToken = Config.DBotsToken; if (dbotsToken != null) { try { const string dBotsApiUrl = "https://discord.bots.gg/api/v1/bots/{0}/stats"; const string Body = "{{ \"guildCount\": {0} }}"; var uri = new Uri(string.Format(dBotsApiUrl, botId)); var post = new HttpRequestMessage(HttpMethod.Post, uri); post.Headers.Add("Authorization", dbotsToken); post.Content = new StringContent(string.Format(Body, count), Encoding.UTF8, "application/json"); await _httpClient.SendAsync(post, token); Log("Discord Bots: Update successful."); } catch (Exception ex) { Log("Discord Bots: Exception encountered during update: " + ex.Message); } } } #endregion } }