mirror of
https://github.com/NoiTheCat/BirthdayBot.git
synced 2024-10-16 07:59:57 +00:00
206 lines
8.5 KiB
C#
206 lines
8.5 KiB
C#
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
|
|
{
|
|
/// <summary>
|
|
/// 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.
|
|
/// </summary>
|
|
class ShardManager : IDisposable
|
|
{
|
|
/// <summary>
|
|
/// Array indexes correspond to shard IDs. Lock on itself when modifying.
|
|
/// </summary>
|
|
private readonly ShardInstance[] _shards;
|
|
|
|
// Commonly used command handler instances
|
|
private readonly Dictionary<string, CommandHandler> _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;
|
|
|
|
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<string, CommandHandler>(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);
|
|
|
|
// Start shards
|
|
_shards = new ShardInstance[Config.ShardCount];
|
|
for (int i = 0; i < _shards.Length; i++)
|
|
InitializeShard(i).Wait();
|
|
|
|
// 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<Task>();
|
|
foreach (var shard in _shards)
|
|
{
|
|
if (shard == null) continue;
|
|
shardDisposes.Add(Task.Run(shard.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);
|
|
|
|
/// <summary>
|
|
/// Creates and sets up a new shard instance.
|
|
/// Shuts down and removes an instance with equivalent ID if already exists.
|
|
/// </summary>
|
|
private async Task InitializeShard(int shardId)
|
|
{
|
|
ShardInstance newInstance;
|
|
lock (_shards)
|
|
{
|
|
Task disposeOldShard;
|
|
if (_shards[shardId] != null)
|
|
disposeOldShard = Task.Run(_shards[shardId].Dispose);
|
|
else
|
|
disposeOldShard = Task.CompletedTask;
|
|
|
|
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
|
|
AlwaysDownloadUsers = true, // TODO set to false when more stable to do so
|
|
GatewayIntents = GatewayIntents.Guilds | GatewayIntents.GuildMembers | GatewayIntents.GuildMessages
|
|
};
|
|
var newClient = new DiscordSocketClient(clientConf);
|
|
newInstance = new ShardInstance(this, newClient, _dispatchCommands);
|
|
|
|
disposeOldShard.Wait();
|
|
_shards[shardId] = newInstance;
|
|
}
|
|
await newInstance.Start();
|
|
}
|
|
|
|
private async Task WatchdogLoop()
|
|
{
|
|
try
|
|
{
|
|
while (!_watchdogCancel.IsCancellationRequested)
|
|
{
|
|
Log($"Bot uptime: {Common.BotUptime}");
|
|
|
|
// Gather statistical information within the lock
|
|
var guildCounts = new int[_shards.Length];
|
|
var connScores = new int[_shards.Length];
|
|
var lastRuns = new DateTimeOffset[_shards.Length];
|
|
ulong? botId = null;
|
|
lock (_shards)
|
|
{
|
|
for (int i = 0; i < _shards.Length; i++)
|
|
{
|
|
var shard = _shards[i];
|
|
if (shard == null) continue;
|
|
|
|
guildCounts[i] = shard.DiscordClient.Guilds.Count;
|
|
connScores[i] = shard.ConnectionScore;
|
|
lastRuns[i] = shard.LastBackgroundRun;
|
|
botId ??= shard.DiscordClient.CurrentUser?.Id;
|
|
}
|
|
}
|
|
|
|
// Guild count
|
|
var guildCountSum = guildCounts.Sum();
|
|
Log($"Currently in {guildCountSum} guilds.");
|
|
if (botId.HasValue)
|
|
await SendExternalStatistics(guildCountSum, botId.Value, _watchdogCancel.Token);
|
|
|
|
// Connection scores and worker health display
|
|
var now = DateTimeOffset.UtcNow;
|
|
for (int i = 0; i < connScores.Length; i++)
|
|
{
|
|
var dur = now - lastRuns[i];
|
|
var lastRunDuration = $"Last run: {Math.Floor(dur.TotalMinutes):00}m{dur.Seconds:00}s ago";
|
|
|
|
Log($"Shard {i:00}: Score {connScores[i]:+0000;-0000} - " + lastRunDuration);
|
|
}
|
|
|
|
// 120 second delay
|
|
await Task.Delay(120 * 1000, _watchdogCancel.Token);
|
|
}
|
|
}
|
|
catch (TaskCanceledException) { }
|
|
}
|
|
|
|
#region Statistical reporting
|
|
private static readonly HttpClient _httpClient = new HttpClient();
|
|
|
|
/// <summary>
|
|
/// Send statistical information to external services.
|
|
/// </summary>
|
|
private async Task SendExternalStatistics(int count, ulong botId, CancellationToken token)
|
|
{
|
|
// TODO protect against exceptions
|
|
var dbotsToken = Config.DBotsToken;
|
|
if (dbotsToken != null)
|
|
{
|
|
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 Task.Delay(80); // Discord Bots rate limit for this endpoint is 20 per second
|
|
await _httpClient.SendAsync(post, token);
|
|
Log("Discord Bots: Count sent successfully.");
|
|
}
|
|
}
|
|
#endregion
|
|
}
|
|
}
|