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

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

@ -104,13 +104,13 @@ namespace PluralKit.Bot
var messageContents = SanitizeEveryoneMaybe(message, match.InnerText); var messageContents = SanitizeEveryoneMaybe(message, match.InnerText);
// Fetch a webhook for this channel, and send the proxied message // 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 avatarUrl = match.Member.AvatarUrl ?? match.System.AvatarUrl;
var proxyName = match.Member.ProxyName(match.System.Tag); var proxyName = match.Member.ProxyName(match.System.Tag);
ulong hookMessageId; ulong hookMessageId;
using (_metrics.Measure.Timer.Time(BotMetrics.WebhookResponseTime)) 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) // 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); await _messageStorage.Store(message.Author.Id, hookMessageId, message.Channel.Id, message.Id, match.Member);
@ -153,70 +153,53 @@ namespace PluralKit.Bot
return true; 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); 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: 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? // 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"); _logger.Debug("Invoking webhook");
// TODO: clean this entire block up // TODO: clean this entire block up
using (client) ulong messageId;
try
{ {
ulong messageId; var client = cacheEntry.Client;
if (attachment != null)
try
{ {
if (attachment != null) using (var stream = await _httpClient.GetStreamAsync(attachment.Url))
{ {
using (var stream = await _httpClient.GetStreamAsync(attachment.Url)) messageId = await client.SendFileAsync(stream, filename: attachment.Filename, text: text,
{ username: username, avatarUrl: avatarUrl);
messageId = await client.SendFileAsync(stream, filename: attachment.Filename, text: text,
username: username, avatarUrl: avatarUrl);
}
} }
else
{
messageId = await client.SendMessageAsync(text, username: username, avatarUrl: avatarUrl);
}
_logger.Information("Invoked webhook {Webhook} in channel {Channel}", webhook.Id,
webhook.ChannelId);
// Log it in the metrics
_metrics.Measure.Meter.Mark(BotMetrics.MessagesProxied, "success");
} }
catch (HttpException e) else
{ {
_logger.Warning(e, "Error invoking webhook {Webhook} in channel {Channel}", webhook.Id, messageId = await client.SendMessageAsync(text, username: username, avatarUrl: avatarUrl);
webhook.ChannelId);
// Log failure in metrics and rethrow (we still need to cancel everything else)
_metrics.Measure.Meter.Mark(BotMetrics.MessagesProxied, "failure");
throw;
} }
// TODO: figure out a way to return the full message object (without doing a GetMessageAsync call, which _logger.Information("Invoked webhook {Webhook} in channel {Channel}", cacheEntry.Webhook.Id,
// doesn't work if there's no permission to) cacheEntry.Webhook.ChannelId);
return messageId;
// 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}", 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");
throw;
}
// TODO: figure out a way to return the full message object (without doing a GetMessageAsync call, which
// doesn't work if there's no permission to)
return messageId;
} }
public Task HandleReactionAddedAsync(Cacheable<IUserMessage, ulong> message, ISocketMessageChannel channel, SocketReaction reaction) public Task HandleReactionAddedAsync(Cacheable<IUserMessage, ulong> message, ISocketMessageChannel channel, SocketReaction reaction)

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

@ -3,16 +3,23 @@ using System.Collections.Concurrent;
using System.Linq; using System.Linq;
using System.Threading.Tasks; using System.Threading.Tasks;
using Discord; using Discord;
using Discord.Webhook;
using Serilog; using Serilog;
namespace PluralKit.Bot 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"; public static readonly string WebhookName = "PluralKit Proxy Webhook";
private IDiscordClient _client; private IDiscordClient _client;
private ConcurrentDictionary<ulong, Lazy<Task<IWebhook>>> _webhooks; private ConcurrentDictionary<ulong, Lazy<Task<WebhookCacheEntry>>> _webhooks;
private ILogger _logger; private ILogger _logger;
@ -20,56 +27,79 @@ namespace PluralKit.Bot
{ {
_client = client; _client = client;
_logger = logger.ForContext<WebhookCacheService>(); _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; var channel = await _client.GetChannelAsync(channelId) as ITextChannel;
if (channel == null) return null; if (channel == null) return null;
return await GetWebhook(channel); 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 // 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 // If the webhook is requested twice before it's actually been found, the Lazy<T> wrapper will stop the
// webhook from being created twice. // webhook from being created twice.
var lazyWebhookValue = 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 // 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. // Here, we ensure it's actually still pointing towards the proper channel, and if not, wipe and refetch one.
var webhook = await lazyWebhookValue.Value; 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; 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 _); _webhooks.TryRemove(webhook.Webhook.ChannelId, out _);
return await GetWebhook(webhook.ChannelId); return await GetWebhook(webhook.Webhook.Channel);
} }
private async Task<IWebhook> GetOrCreateWebhook(ITextChannel channel) => private async Task<WebhookCacheEntry> GetOrCreateWebhook(ITextChannel channel)
await FindExistingWebhook(channel) ?? await DoCreateWebhook(channel); {
var webhook = await FindExistingWebhook(channel) ?? await DoCreateWebhook(channel);
return await DoCreateWebhookClient(webhook);
}
private async Task<IWebhook> FindExistingWebhook(ITextChannel channel) private async Task<IWebhook> FindExistingWebhook(ITextChannel channel)
{ {
_logger.Debug("Finding webhook for channel {Channel}", channel.Id); _logger.Debug("Finding webhook for channel {Channel}", channel.Id);
return (await channel.GetWebhooksAsync()).FirstOrDefault(IsWebhookMine); 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); _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 int CacheSize => _webhooks.Count;
public void Dispose()
{
foreach (var entry in _webhooks.Values)
if (entry.IsValueCreated)
entry.Value.Result.Client.Dispose();
}
} }
} }