forked from EllieBotDevs/elliebot
324 lines
No EOL
10 KiB
C#
324 lines
No EOL
10 KiB
C#
using LinqToDB;
|
|
using LinqToDB.Data;
|
|
using LinqToDB.EntityFrameworkCore;
|
|
using LinqToDB.Mapping;
|
|
using EllieBot.Common.ModuleBehaviors;
|
|
using EllieBot.Db.Models;
|
|
|
|
namespace EllieBot.Modules.Administration.DangerousCommands;
|
|
|
|
public sealed class CleanupService : ICleanupService, IReadyExecutor, IEService
|
|
{
|
|
private TypedKey<KeepReport> _cleanupReportKey = new("cleanup:report");
|
|
private TypedKey<bool> _cleanupTriggerKey = new("cleanup:trigger");
|
|
|
|
private TypedKey<int> _keepTriggerKey = new("keep:trigger");
|
|
|
|
private readonly IPubSub _pubSub;
|
|
private readonly DiscordSocketClient _client;
|
|
private ConcurrentDictionary<int, ulong[]> guildIds = new();
|
|
private readonly IBotCredsProvider _creds;
|
|
private readonly DbService _db;
|
|
|
|
public CleanupService(
|
|
IPubSub pubSub,
|
|
DiscordSocketClient client,
|
|
IBotCredsProvider creds,
|
|
DbService db)
|
|
{
|
|
_pubSub = pubSub;
|
|
_client = client;
|
|
_creds = creds;
|
|
_db = db;
|
|
}
|
|
|
|
public async Task OnReadyAsync()
|
|
{
|
|
await _pubSub.Sub(_cleanupTriggerKey, OnCleanupTrigger);
|
|
await _pubSub.Sub(_keepTriggerKey, InternalTriggerKeep);
|
|
|
|
_client.JoinedGuild += ClientOnJoinedGuild;
|
|
|
|
if (_client.ShardId == 0)
|
|
await _pubSub.Sub(_cleanupReportKey, OnKeepReport);
|
|
}
|
|
|
|
private bool keepTriggered = false;
|
|
|
|
private async ValueTask InternalTriggerKeep(int shardId)
|
|
{
|
|
if (_client.ShardId != shardId)
|
|
return;
|
|
|
|
if (keepTriggered)
|
|
return;
|
|
|
|
keepTriggered = true;
|
|
try
|
|
{
|
|
var allGuildIds = _client.Guilds.Select(x => x.Id).ToArray();
|
|
|
|
HashSet<ulong> dontDelete;
|
|
await using (var db = _db.GetDbContext())
|
|
{
|
|
await using var ctx = db.CreateLinqToDBContext();
|
|
var table = ctx.CreateTable<KeptGuilds>(tableOptions: TableOptions.CheckExistence);
|
|
|
|
var dontDeleteList = await table
|
|
.Where(x => allGuildIds.Contains(x.GuildId))
|
|
.Select(x => x.GuildId)
|
|
.ToListAsyncLinqToDB();
|
|
|
|
dontDelete = dontDeleteList.ToHashSet();
|
|
}
|
|
|
|
Log.Information("Leaving {RemainingCount} guilds, 1 every second. {DontDeleteCount} will remain",
|
|
allGuildIds.Length - dontDelete.Count,
|
|
dontDelete.Count);
|
|
|
|
foreach (var guildId in allGuildIds)
|
|
{
|
|
if (dontDelete.Contains(guildId))
|
|
continue;
|
|
|
|
await Task.Delay(1016);
|
|
|
|
SocketGuild? guild = null;
|
|
try
|
|
{
|
|
guild = _client.GetGuild(guildId);
|
|
|
|
if (guild is null)
|
|
{
|
|
Log.Warning("Unable to find guild {GuildId}", guildId);
|
|
continue;
|
|
}
|
|
|
|
await guild.LeaveAsync();
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
Log.Warning("Unable to leave guild {GuildName} [{GuildId}]: {ErrorMessage}",
|
|
guild?.Name,
|
|
guildId,
|
|
ex.Message);
|
|
}
|
|
}
|
|
}
|
|
finally
|
|
{
|
|
keepTriggered = false;
|
|
}
|
|
}
|
|
|
|
public async Task<KeepResult?> DeleteMissingGuildDataAsync()
|
|
{
|
|
guildIds = new();
|
|
var totalShards = _creds.GetCreds().TotalShards;
|
|
await _pubSub.Pub(_cleanupTriggerKey, true);
|
|
var counter = 0;
|
|
while (guildIds.Keys.Count < totalShards)
|
|
{
|
|
await Task.Delay(1000);
|
|
counter++;
|
|
|
|
if (counter >= 5)
|
|
break;
|
|
}
|
|
|
|
if (guildIds.Keys.Count < totalShards)
|
|
return default;
|
|
|
|
var allIds = guildIds.SelectMany(x => x.Value)
|
|
.ToArray();
|
|
|
|
await using var ctx = _db.GetDbContext();
|
|
await using var linqCtx = ctx.CreateLinqToDBContext();
|
|
await using var tempTable = linqCtx.CreateTempTable<CleanupId>();
|
|
|
|
foreach (var chunk in allIds.Chunk(10000))
|
|
{
|
|
await tempTable.BulkCopyAsync(chunk.Select(x => new CleanupId()
|
|
{
|
|
GuildId = x
|
|
}));
|
|
}
|
|
|
|
// delete guild configs
|
|
await ctx.GetTable<GuildConfig>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete guild xp
|
|
await ctx.GetTable<UserXpStats>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete expressions
|
|
await ctx.GetTable<EllieExpression>()
|
|
.Where(x => x.GuildId != null
|
|
&& !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId.Value))
|
|
.DeleteAsync();
|
|
|
|
// delete quotes
|
|
await ctx.GetTable<Quote>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete planted currencies
|
|
await ctx.GetTable<PlantedCurrency>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete image only channels
|
|
await ctx.GetTable<ImageOnlyChannel>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete reaction roles
|
|
await ctx.GetTable<ReactionRoleV2>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete perm overrides
|
|
await ctx.GetTable<DiscordPermOverride>()
|
|
.Where(x => x.GuildId != null
|
|
&& !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId.Value))
|
|
.DeleteAsync();
|
|
|
|
// delete repeaters
|
|
await ctx.GetTable<Repeater>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete autopublish channels
|
|
await ctx.GetTable<AutoPublishChannel>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete greet settings
|
|
await ctx.GetTable<GreetSettings>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete sar
|
|
await ctx.GetTable<SarGroup>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete warnings
|
|
await ctx.GetTable<Warning>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete warn punishments
|
|
await ctx.GetTable<WarningPunishment>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete sticky roles
|
|
await ctx.GetTable<StickyRole>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete at channels
|
|
await ctx.GetTable<AutoTranslateChannel>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete ban templates
|
|
await ctx.GetTable<BanTemplate>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
// delete reminders
|
|
await ctx.GetTable<Reminder>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.ServerId))
|
|
.DeleteAsync();
|
|
|
|
// delete button roles
|
|
await ctx.GetTable<ButtonRole>()
|
|
.Where(x => !tempTable.Select(x => x.GuildId)
|
|
.Contains(x.GuildId))
|
|
.DeleteAsync();
|
|
|
|
return new()
|
|
{
|
|
GuildCount = guildIds.Keys.Count,
|
|
};
|
|
}
|
|
|
|
public async Task<bool> KeepGuild(ulong guildId)
|
|
{
|
|
await using var db = _db.GetDbContext();
|
|
await using var ctx = db.CreateLinqToDBContext();
|
|
var table = ctx.CreateTable<KeptGuilds>(tableOptions: TableOptions.CheckExistence);
|
|
if (await table.AnyAsyncLinqToDB(x => x.GuildId == guildId))
|
|
return false;
|
|
|
|
await table.InsertAsync(() => new()
|
|
{
|
|
GuildId = guildId
|
|
});
|
|
|
|
return true;
|
|
}
|
|
|
|
public async Task<int> GetKeptGuildCount()
|
|
{
|
|
await using var db = _db.GetDbContext();
|
|
await using var ctx = db.CreateLinqToDBContext();
|
|
var table = ctx.CreateTable<KeptGuilds>(tableOptions: TableOptions.CheckExistence);
|
|
return await table.CountAsync();
|
|
}
|
|
|
|
public async Task StartLeavingUnkeptServers(int shardId)
|
|
=> await _pubSub.Pub(_keepTriggerKey, shardId);
|
|
|
|
private ValueTask OnKeepReport(KeepReport report)
|
|
{
|
|
guildIds[report.ShardId] = report.GuildIds;
|
|
return default;
|
|
}
|
|
|
|
private async Task ClientOnJoinedGuild(SocketGuild arg)
|
|
{
|
|
await KeepGuild(arg.Id);
|
|
}
|
|
|
|
private ValueTask OnCleanupTrigger(bool arg)
|
|
{
|
|
_pubSub.Pub(_cleanupReportKey,
|
|
new KeepReport()
|
|
{
|
|
ShardId = _client.ShardId,
|
|
GuildIds = _client.GetGuildIds(),
|
|
});
|
|
|
|
return default;
|
|
}
|
|
}
|
|
|
|
public class KeptGuilds
|
|
{
|
|
[PrimaryKey]
|
|
public ulong GuildId { get; set; }
|
|
} |