fix embed queue

This commit is contained in:
sam 2024-08-13 16:48:54 +02:00
parent ded4f4db26
commit 8d4a7b1729
Signed by: sam
GPG key ID: 5F3C3C1B3166639D
13 changed files with 188 additions and 32 deletions

View file

@ -1,10 +1,16 @@
using System.ComponentModel;
using System.Diagnostics;
using Catalogger.Backend.Cache;
using Catalogger.Backend.Database;
using Catalogger.Backend.Extensions;
using Humanizer;
using Microsoft.EntityFrameworkCore;
using NodaTime;
using Remora.Commands.Attributes;
using Remora.Commands.Groups;
using Remora.Discord.API.Abstractions.Rest;
using Remora.Discord.Commands.Feedback.Services;
using Remora.Discord.Extensions.Embeds;
using Remora.Discord.Gateway;
using Remora.Results;
using IResult = Remora.Results.IResult;
@ -14,7 +20,10 @@ namespace Catalogger.Backend.Bot.Commands;
[Group("catalogger")]
public class MetaCommands(
IClock clock,
DatabaseContext db,
DiscordGatewayClient client,
GuildCacheService guildCache,
ChannelCacheService channelCache,
IFeedbackService feedbackService,
IDiscordRestChannelAPI channelApi) : CommandGroup
{
@ -26,6 +35,24 @@ public class MetaCommands(
var msg = await feedbackService.SendContextualAsync("...").GetOrThrow();
var elapsed = clock.GetCurrentInstant() - t1;
var messageCount = await db.Messages.CountAsync();
var process = Process.GetCurrentProcess();
var memoryUsage = process.WorkingSet64;
var virtualMemory = process.VirtualMemorySize64;
var embed = new EmbedBuilder()
.WithColour(DiscordUtils.Purple)
.WithFooter("")
.WithCurrentTimestamp();
embed.AddField("Ping", $"Gateway: {client.Latency.Humanize()}\nAPI: {elapsed.ToTimeSpan().Humanize()}",
inline: true);
embed.AddField("Memory usage", $"{memoryUsage.Bytes().Humanize()} / {virtualMemory.Bytes().Humanize()}",
inline: true);
embed.AddField("Numbers",
$"{messageCount:N0} messages from {guildCache.Size:N0} servers\nCached {channelCache.Size:N0} channels",
inline: false);
return (Result)await channelApi.EditMessageAsync(msg.ChannelID, msg.ID,
content: $"Pong! API: {elapsed.TotalMilliseconds:N0}ms | Gateway: {client.Latency.TotalMilliseconds:N0}ms");
}

View file

@ -5,5 +5,5 @@ namespace Catalogger.Backend.Bot;
public static class DiscordUtils
{
public static readonly Color Red = Color.FromArgb(231, 76, 60);
public static readonly Color Blue = Color.FromArgb(155, 89, 182);
public static readonly Color Purple = Color.FromArgb(155, 89, 182);
}

View file

@ -3,14 +3,23 @@ using Catalogger.Backend.Cache;
using Catalogger.Backend.Database;
using Catalogger.Backend.Database.Models;
using Catalogger.Backend.Extensions;
using Catalogger.Backend.Services;
using Remora.Discord.API.Abstractions.Gateway.Events;
using Remora.Discord.API.Gateway.Events;
using Remora.Discord.Extensions.Embeds;
using Remora.Discord.Gateway.Responders;
using Remora.Results;
namespace Catalogger.Backend.Bot.Responders;
public class GuildCreateResponder(ILogger logger, DatabaseContext db, ChannelCacheService channelCache)
: IResponder<IGuildCreate>
public class GuildCreateResponder(
Config config,
ILogger logger,
DatabaseContext db,
GuildCacheService guildCache,
ChannelCacheService channelCache,
WebhookExecutorService webhookExecutor)
: IResponder<IGuildCreate>, IResponder<GuildDelete>
{
private readonly ILogger _logger = logger.ForContext<GuildCreateResponder>();
@ -20,11 +29,12 @@ public class GuildCreateResponder(ILogger logger, DatabaseContext db, ChannelCac
string? guildName = null;
if (evt.Guild.TryPickT0(out var guild, out _))
{
_logger.Verbose("Received guild create for available guild {GuildId} ({Name})", guild.ID, guild.Name);
_logger.Verbose("Received guild create for available guild {GuildName} / {GuildId})", guild.Name, guild.ID);
guildId = guild.ID.ToUlong();
guildName = guild.Name;
foreach (var c in guild.Channels) channelCache.AddChannel(c, guild.ID);
guildCache.Set(guild);
foreach (var c in guild.Channels) channelCache.Set(c, guild.ID);
}
else if (evt.Guild.TryPickT1(out var unavailableGuild, out _))
{
@ -43,7 +53,51 @@ public class GuildCreateResponder(ILogger logger, DatabaseContext db, ChannelCac
await db.SaveChangesAsync(ct);
await tx.CommitAsync(ct);
_logger.Information("Joined new guild {GuildId} ({Name})", guildId, guildName);
_logger.Information("Joined new guild {GuildName} / {GuildId}", guildName, guildId);
if (config.Discord.GuildLogId != null && evt.Guild.IsT0)
await webhookExecutor.QueueLogAsync(config.Discord.GuildLogId.Value, new EmbedBuilder()
.WithTitle("Joined new guild")
.WithDescription($"Joined new guild **{guild.Name}**")
.WithFooter($"ID: {guild.ID}")
.WithCurrentTimestamp()
#pragma warning disable CS8604 // Possible null reference argument.
.WithThumbnailUrl(guild.IconUrl())
#pragma warning restore CS8604 // Possible null reference argument.
.Build()
.GetOrThrow());
return Result.Success;
}
public async Task<Result> RespondAsync(GuildDelete evt, CancellationToken ct = default)
{
if (evt.IsUnavailable.OrDefault(false))
{
_logger.Debug("Guild {GuildId} became unavailable", evt.ID);
return Result.Success;
}
if (!guildCache.TryGet(evt.ID, out var guild))
{
_logger.Information("Left uncached guild {GuildId}", evt.ID);
return Result.Success;
}
_logger.Information("Left guild {GuildName} / {GuildId}", guild.Name, guild.ID);
if (config.Discord.GuildLogId != null)
await webhookExecutor.QueueLogAsync(config.Discord.GuildLogId.Value, new EmbedBuilder()
.WithTitle("Left guild")
.WithDescription($"Left guild **{guild.Name}**")
.WithFooter($"ID: {guild.ID}")
.WithCurrentTimestamp()
#pragma warning disable CS8604 // Possible null reference argument.
.WithThumbnailUrl(guild.IconUrl())
#pragma warning restore CS8604 // Possible null reference argument.
.Build()
.GetOrThrow());
return Result.Success;
}
}

View file

@ -16,6 +16,7 @@ namespace Catalogger.Backend.Bot.Responders;
public class MessageCreateResponder(
ILogger logger,
Config config,
DatabaseContext db,
MessageRepository messageRepository,
UserCacheService userCache,
@ -50,6 +51,12 @@ public class MessageCreateResponder(
_ = pkMessageHandler.HandlePkMessageAsync(msg);
if (msg.ApplicationID.IsDefined(out var appId) && appId == PkUserId)
_ = pkMessageHandler.HandleProxiedMessageAsync(msg.ID.Value);
else if (msg.ApplicationID.HasValue && appId == config.Discord.ApplicationId)
{
db.IgnoredMessages.Add(new IgnoredMessage(msg.ID.Value));
await db.SaveChangesAsync(ct);
return Result.Success;
}
await messageRepository.SaveMessageAsync(msg, ct);
return Result.Success;
@ -59,20 +66,20 @@ public class MessageCreateResponder(
public partial class PkMessageHandler(ILogger logger, IServiceProvider services)
{
private readonly ILogger _logger = logger.ForContext<PkMessageHandler>();
[GeneratedRegex(
@"^System ID: (\w{5,6}) \| Member ID: (\w{5,6}) \| Sender: .+ \((\d+)\) \| Message ID: (\d+) \| Original Message ID: (\d+)$")]
private static partial Regex FooterRegex();
[GeneratedRegex(@"^https:\/\/discord.com\/channels\/\d+\/(\d+)\/\d+$")]
private static partial Regex LinkRegex();
public async Task HandlePkMessageAsync(IMessageCreate msg)
{
_logger.Debug("Received PluralKit message");
await Task.Delay(500.Milliseconds());
_logger.Debug("Starting handling PluralKit message");
// Check if the content matches a Discord link--if not, it's not a log message (we already check if this is a PluralKit message earlier)
@ -122,11 +129,11 @@ public partial class PkMessageHandler(ILogger logger, IServiceProvider services)
db.IgnoredMessages.Add(new IgnoredMessage(originalId));
await db.SaveChangesAsync();
}
public async Task HandleProxiedMessageAsync(ulong msgId)
{
await Task.Delay(3.Seconds());
await using var scope = services.CreateAsyncScope();
await using var db = scope.ServiceProvider.GetRequiredService<DatabaseContext>();
var messageRepository = scope.ServiceProvider.GetRequiredService<MessageRepository>();

View file

@ -75,7 +75,7 @@ public class MessageDeleteResponder(
if (msg.Member != null) builder.WithTitle($"Message by {msg.Username} deleted");
string channelMention;
if (!channelCache.GetChannel(ev.ChannelID, out var channel))
if (!channelCache.TryGet(ev.ChannelID, out var channel))
channelMention = $"<#{msg.ChannelId}>";
else if (channel.Type is ChannelType.AnnouncementThread or ChannelType.PrivateThread
or ChannelType.PublicThread)

View file

@ -10,7 +10,9 @@ public class ChannelCacheService
private readonly ConcurrentDictionary<Snowflake, IChannel> _channels = new();
private readonly ConcurrentDictionary<Snowflake, HashSet<Snowflake>> _guildChannels = new();
public void AddChannel(IChannel channel, Snowflake? guildId = null)
public int Size => _channels.Count;
public void Set(IChannel channel, Snowflake? guildId = null)
{
_channels[channel.ID] = channel;
if (guildId == null)
@ -29,9 +31,9 @@ public class ChannelCacheService
});
}
public bool GetChannel(Snowflake id, [NotNullWhen(true)] out IChannel? channel) => _channels.TryGetValue(id, out channel);
public bool TryGet(Snowflake id, [NotNullWhen(true)] out IChannel? channel) => _channels.TryGetValue(id, out channel);
public void RemoveChannel(Snowflake? guildId, Snowflake id, out IChannel? channel)
public void Remove(Snowflake? guildId, Snowflake id, out IChannel? channel)
{
_channels.Remove(id, out channel);
if (guildId == null) return;

View file

@ -0,0 +1,17 @@
using System.Collections.Concurrent;
using System.Diagnostics.CodeAnalysis;
using Remora.Discord.API.Abstractions.Objects;
using Remora.Rest.Core;
namespace Catalogger.Backend.Cache;
public class GuildCacheService
{
private readonly ConcurrentDictionary<Snowflake, IGuild> _guilds = new();
public int Size => _guilds.Count;
public void Set(IGuild guild) => _guilds[guild.ID] = guild;
public bool Remove(Snowflake id, [NotNullWhen(true)] out IGuild? guild) => _guilds.Remove(id, out guild);
public bool TryGet(Snowflake id, [NotNullWhen(true)] out IGuild? guild) => _guilds.TryGetValue(id, out guild);
}

View file

@ -30,6 +30,7 @@ public class Config
public string Token { get; init; } = string.Empty;
public bool SyncCommands { get; init; }
public ulong? CommandsGuildId { get; init; }
public ulong? GuildLogId { get; init; }
}
public class WebConfig

View file

@ -21,7 +21,9 @@ public class MessageRepository(ILogger logger, DatabaseContext db, IEncryptionSe
ChannelId = msg.ChannelID.ToUlong(),
GuildId = msg.GuildID.ToUlong(),
EncryptedContent = await Task.Run(() => encryptionService.Encrypt(msg.Content), ct),
EncryptedContent =
await Task.Run(
() => encryptionService.Encrypt(string.IsNullOrWhiteSpace(msg.Content) ? "None" : msg.Content), ct),
EncryptedUsername = await Task.Run(() => encryptionService.Encrypt(msg.Author.Tag()), ct),
AttachmentSize = msg.Attachments.Select(a => a.Size).Sum()
};

View file

@ -16,7 +16,7 @@ public static class QueryExtensions
CancellationToken ct = default)
{
var guild = await db.Guilds.FindAsync(id);
if (guild == null) throw new Exception("oh");
if (guild == null) throw new CataloggerError("Guild not found, was not initialized during guild create");
return guild;
}
}

View file

@ -23,6 +23,15 @@ public static class DiscordExtensions
var avatarIndex = user.Discriminator == 0 ? (int)((user.ID.Value >> 22) % 6) : user.Discriminator % 5;
return $"https://cdn.discordapp.com/embed/avatars/{avatarIndex}.png?size={size}";
}
public static string? IconUrl(this IGuild guild, int size = 256)
{
if (guild.Icon == null) return null;
var ext = guild.Icon.HasGif ? ".gif" : ".webp";
return $"https://cdn.discordapp.com/icons/{guild.ID}/{guild.Icon.Value}{ext}?size={size}";
}
public static ulong ToUlong(this Snowflake snowflake) => snowflake.Value;

View file

@ -63,6 +63,7 @@ public static class StartupExtensions
public static IServiceCollection AddCustomServices(this IServiceCollection services) => services
.AddSingleton<IClock>(SystemClock.Instance)
.AddSingleton<GuildCacheService>()
.AddSingleton<ChannelCacheService>()
.AddSingleton<UserCacheService>()
.AddSingleton<PluralkitApiService>()

View file

@ -1,15 +1,15 @@
using System.Collections.Concurrent;
using Catalogger.Backend.Cache;
using Catalogger.Backend.Database.Models;
using Catalogger.Backend.Extensions;
using Humanizer;
using Remora.Discord.API;
using Remora.Discord.API.Abstractions.Objects;
using Remora.Discord.API.Abstractions.Rest;
using Remora.Rest.Core;
using Guild = Catalogger.Backend.Database.Models.Guild;
namespace Catalogger.Backend.Services;
// TODO: this entire class is a mess, clean it up
public class WebhookExecutorService(
Config config,
ILogger logger,
@ -20,6 +20,7 @@ public class WebhookExecutorService(
private readonly ILogger _logger = logger.ForContext<WebhookExecutorService>();
private readonly Snowflake _applicationId = DiscordSnowflake.New(config.Discord.ApplicationId);
private readonly ConcurrentDictionary<ulong, ConcurrentQueue<IEmbed>> _cache = new();
private readonly ConcurrentDictionary<ulong, object> _locks = new();
private readonly ConcurrentDictionary<ulong, Timer> _timers = new();
private IUser? _selfUser;
@ -33,29 +34,52 @@ public class WebhookExecutorService(
await QueueLogAsync(logChannel.Value, embed);
}
private List<IEmbed> TakeFromQueue(ulong channelId)
{
var queue = _cache.GetOrAdd(channelId, []);
var channelLock = _locks.GetOrAdd(channelId, channelId);
lock (channelLock)
{
var embeds = new List<IEmbed>();
for (var i = 0; i < 5; i++)
{
if (!queue.TryDequeue(out var embed)) break;
embeds.Add(embed);
}
return embeds;
}
}
public async Task QueueLogAsync(ulong channelId, IEmbed embed)
{
_logger.Debug("Queueing embed for channel {ChannelId}", channelId);
var webhook = await webhookCache.GetOrFetchWebhookAsync(channelId, id => FetchWebhookAsync(id));
var queue = _cache.GetOrAdd(channelId, []);
if (queue.Count >= 5)
await SendLogsAsync(channelId);
queue.Enqueue(embed);
_cache[channelId] = queue;
await SetTimer(channelId, queue);
}
private async Task SetTimer(ulong channelId, ConcurrentQueue<IEmbed> queue)
{
if (_timers.TryGetValue(channelId, out var existingTimer)) await existingTimer.DisposeAsync();
_timers[channelId] = new Timer(_ =>
{
var __ = SendLogsAsync(channelId);
_logger.Debug("Sending 5 queued embeds");
var __ = SendLogsAsync(channelId, TakeFromQueue(channelId));
if (!queue.IsEmpty)
{
if (_timers.TryGetValue(channelId, out var timer)) timer.Dispose();
var ___ = SetTimer(channelId, queue);
}
}, null, 3000, Timeout.Infinite);
}
private async Task SendLogsAsync(ulong channelId)
private async Task SendLogsAsync(ulong channelId, List<IEmbed> embeds)
{
var queue = _cache.GetValueOrDefault(channelId);
if (queue == null) return;
var embeds = queue.Take(5).ToList();
_logger.Debug("Sending {Count} embeds to channel {ChannelId}", embeds.Count, channelId);
if (embeds.Count == 0) return;
var webhook = await webhookCache.GetOrFetchWebhookAsync(channelId, id => FetchWebhookAsync(id));
@ -63,6 +87,18 @@ public class WebhookExecutorService(
embeds: embeds, username: _selfUser!.Username, avatarUrl: _selfUser.AvatarUrl());
}
public async Task SendLogWithAttachmentsAsync(ulong channelId, IEmbed embed, IEnumerable<FileData> files)
{
var attachments = files
.Select<FileData, OneOf.OneOf<FileData, IPartialAttachment>>(f => f)
.ToList();
var webhook = await webhookCache.GetOrFetchWebhookAsync(channelId, id => FetchWebhookAsync(id));
await webhookApi.ExecuteWebhookAsync(DiscordSnowflake.New(webhook.Id), webhook.Token, shouldWait: false,
embeds: new List<IEmbed>([embed]), attachments: attachments, username: _selfUser!.Username,
avatarUrl: _selfUser.AvatarUrl());
}
private async Task<IWebhook> FetchWebhookAsync(Snowflake channelId, CancellationToken ct = default)
{
var channelWebhooks =
@ -78,14 +114,14 @@ public class WebhookExecutorService(
ulong? userId = null)
{
if (channelId == null) return GetDefaultLogChannel(guild, logChannelType);
if (!channelCache.GetChannel(channelId.Value, out var channel)) return null;
if (!channelCache.TryGet(channelId.Value, out var channel)) return null;
Snowflake? categoryId;
if (channel.Type is ChannelType.AnnouncementThread or ChannelType.PrivateThread or ChannelType.PublicThread)
{
// parent_id should always have a value for threads
channelId = channel.ParentID.Value!.Value;
if (!channelCache.GetChannel(channelId.Value, out var parentChannel))
if (!channelCache.TryGet(channelId.Value, out var parentChannel))
return GetDefaultLogChannel(guild, logChannelType);
categoryId = parentChannel.ParentID.Value;
}