liz/PluralKit
liz/PluralKit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Cache the entire webhook client rather than just the token

Browse Source
This commit is contained in:
Ske 2019-08-12 04:32:01 +02:00
parent 1b32cd8b6d
commit 5b13c1b100
2 changed files with 78 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
PluralKit.Bot/Services/ProxyService.cs
View File

@ -104,13 +104,13 @@ namespace PluralKit.Bot
var messageContents = SanitizeEveryoneMaybe(message, match.InnerText);
// Fetch a webhook for this channel, and send the proxied message
var webhook = await _webhookCache.GetWebhook(message.Channel as ITextChannel);
var webhookCacheEntry = await _webhookCache.GetWebhook(message.Channel as ITextChannel);
var avatarUrl = match.Member.AvatarUrl ?? match.System.AvatarUrl;
var proxyName = match.Member.ProxyName(match.System.Tag);
ulong hookMessageId;
using (_metrics.Measure.Timer.Time(BotMetrics.WebhookResponseTime))
hookMessageId = await ExecuteWebhook(webhook, messageContents, proxyName, avatarUrl, message.Attachments.FirstOrDefault());
hookMessageId = await ExecuteWebhook(webhookCacheEntry, messageContents, proxyName, avatarUrl, message.Attachments.FirstOrDefault());
// Store the message in the database, and log it in the log channel (if applicable)
await _messageStorage.Store(message.Author.Id, hookMessageId, message.Channel.Id, message.Id, match.Member);
@ -153,37 +153,21 @@ namespace PluralKit.Bot
return true;
}
private async Task<ulong> ExecuteWebhook(IWebhook webhook, string text, string username, string avatarUrl, IAttachment attachment)
private async Task<ulong> ExecuteWebhook(WebhookCacheService.WebhookCacheEntry cacheEntry, string text, string username, string avatarUrl, IAttachment attachment)
{
username = FixClyde(username);
// TODO: DiscordWebhookClient's ctor does a call to GetWebhook that may be unnecessary, see if there's a way to do this The Hard Way :tm:
// TODO: this will probably crash if there are multiple consecutive failures, perhaps have a loop instead?
_logger.Debug("Instantiating webhook client");
DiscordWebhookClient client;
try
{
client = new DiscordWebhookClient(webhook);
}
catch (InvalidOperationException)
{
// TODO: does this leak internal stuff in the (now-invalid) client?
// webhook was deleted or invalid
webhook = await _webhookCache.InvalidateAndRefreshWebhook(webhook);
client = new DiscordWebhookClient(webhook);
}
_logger.Debug("Invoking webhook");
// TODO: clean this entire block up
using (client)
{
ulong messageId;
try
{
var client = cacheEntry.Client;
if (attachment != null)
{
using (var stream = await _httpClient.GetStreamAsync(attachment.Url))
@ -197,16 +181,16 @@ namespace PluralKit.Bot
messageId = await client.SendMessageAsync(text, username: username, avatarUrl: avatarUrl);
}
_logger.Information("Invoked webhook {Webhook} in channel {Channel}", webhook.Id,
webhook.ChannelId);
_logger.Information("Invoked webhook {Webhook} in channel {Channel}", cacheEntry.Webhook.Id,
cacheEntry.Webhook.ChannelId);
// Log it in the metrics
_metrics.Measure.Meter.Mark(BotMetrics.MessagesProxied, "success");
}
catch (HttpException e)
{
_logger.Warning(e, "Error invoking webhook {Webhook} in channel {Channel}", webhook.Id,
webhook.ChannelId);
_logger.Warning(e, "Error invoking webhook {Webhook} in channel {Channel}", cacheEntry.Webhook.Id,
cacheEntry.Webhook.ChannelId);
// Log failure in metrics and rethrow (we still need to cancel everything else)
_metrics.Measure.Meter.Mark(BotMetrics.MessagesProxied, "failure");
@ -217,7 +201,6 @@ namespace PluralKit.Bot
// doesn't work if there's no permission to)
return messageId;
}
}
public Task HandleReactionAddedAsync(Cacheable<IUserMessage, ulong> message, ISocketMessageChannel channel, SocketReaction reaction)
{

62
PluralKit.Bot/Services/WebhookCacheService.cs
View File

@ -3,16 +3,23 @@ using System.Collections.Concurrent;
using System.Linq;
using System.Threading.Tasks;
using Discord;
using Discord.Webhook;
using Serilog;
namespace PluralKit.Bot
{
public class WebhookCacheService
public class WebhookCacheService: IDisposable
{
public class WebhookCacheEntry
{
internal DiscordWebhookClient Client;
internal IWebhook Webhook;
}
public static readonly string WebhookName = "PluralKit Proxy Webhook";
private IDiscordClient _client;
private ConcurrentDictionary<ulong, Lazy<Task<IWebhook>>> _webhooks;
private ConcurrentDictionary<ulong, Lazy<Task<WebhookCacheEntry>>> _webhooks;
private ILogger _logger;
@ -20,41 +27,44 @@ namespace PluralKit.Bot
{
_client = client;
_logger = logger.ForContext<WebhookCacheService>();
_webhooks = new ConcurrentDictionary<ulong, Lazy<Task<IWebhook>>>();
_webhooks = new ConcurrentDictionary<ulong, Lazy<Task<WebhookCacheEntry>>>();
}
public async Task<IWebhook> GetWebhook(ulong channelId)
public async Task<WebhookCacheEntry> GetWebhook(ulong channelId)
{
var channel = await _client.GetChannelAsync(channelId) as ITextChannel;
if (channel == null) return null;
return await GetWebhook(channel);
}
public async Task<IWebhook> GetWebhook(ITextChannel channel)
public async Task<WebhookCacheEntry> GetWebhook(ITextChannel channel)
{
// We cache the webhook through a Lazy<Task<T>>, this way we make sure to only create one webhook per channel
// If the webhook is requested twice before it's actually been found, the Lazy<T> wrapper will stop the
// webhook from being created twice.
var lazyWebhookValue =
_webhooks.GetOrAdd(channel.Id, new Lazy<Task<IWebhook>>(() => GetOrCreateWebhook(channel)));
_webhooks.GetOrAdd(channel.Id, new Lazy<Task<WebhookCacheEntry>>(() => GetOrCreateWebhook(channel)));
// It's possible to "move" a webhook to a different channel after creation
// Here, we ensure it's actually still pointing towards the proper channel, and if not, wipe and refetch one.
var webhook = await lazyWebhookValue.Value;
if (webhook.Channel.Id != channel.Id) return await InvalidateAndRefreshWebhook(webhook);
if (webhook.Webhook.ChannelId != channel.Id) return await InvalidateAndRefreshWebhook(webhook);
return webhook;
}
public async Task<IWebhook> InvalidateAndRefreshWebhook(IWebhook webhook)
public async Task<WebhookCacheEntry> InvalidateAndRefreshWebhook(WebhookCacheEntry webhook)
{
_logger.Information("Refreshing webhook for channel {Channel}", webhook.ChannelId);
_logger.Information("Refreshing webhook for channel {Channel}", webhook.Webhook.ChannelId);
_webhooks.TryRemove(webhook.ChannelId, out _);
return await GetWebhook(webhook.ChannelId);
_webhooks.TryRemove(webhook.Webhook.ChannelId, out _);
return await GetWebhook(webhook.Webhook.Channel);
}
private async Task<IWebhook> GetOrCreateWebhook(ITextChannel channel) =>
await FindExistingWebhook(channel) ?? await DoCreateWebhook(channel);
private async Task<WebhookCacheEntry> GetOrCreateWebhook(ITextChannel channel)
{
var webhook = await FindExistingWebhook(channel) ?? await DoCreateWebhook(channel);
return await DoCreateWebhookClient(webhook);
}
private async Task<IWebhook> FindExistingWebhook(ITextChannel channel)
{
@ -62,14 +72,34 @@ namespace PluralKit.Bot
return (await channel.GetWebhooksAsync()).FirstOrDefault(IsWebhookMine);
}
private async Task<IWebhook> DoCreateWebhook(ITextChannel channel)
private Task<IWebhook> DoCreateWebhook(ITextChannel channel)
{
_logger.Information("Creating new webhook for channel {Channel}", channel.Id);
return await channel.CreateWebhookAsync(WebhookName);
return channel.CreateWebhookAsync(WebhookName);
}
private bool IsWebhookMine(IWebhook arg) => arg.Creator.Id == _client.CurrentUser.Id && arg.Name == WebhookName;
private Task<WebhookCacheEntry> DoCreateWebhookClient(IWebhook webhook)
{
// DiscordWebhookClient's ctor is synchronous despite doing web calls, so we wrap it in a Task
return Task.Run(() =>
{
return new WebhookCacheEntry
{
Client = new DiscordWebhookClient(webhook),
Webhook = webhook
};
});
}
private bool IsWebhookMine(IWebhook arg) => arg.Id == _client.CurrentUser.Id && arg.Name == WebhookName;
public int CacheSize => _webhooks.Count;
public void Dispose()
{
foreach (var entry in _webhooks.Values)
if (entry.IsValueCreated)
entry.Value.Result.Client.Dispose();
}
}
}