Cache account lookup in memory when proxying

PluralKit.Bot/Bot.cs

@@ -106,6 +106,7 @@ namespace PluralKit.Bot
             .AddTransient<LogChannelService>()
             .AddTransient<DataFileService>()
 
+            .AddSingleton<ProxyCacheService>()
             .AddSingleton<WebhookCacheService>()
 
             .AddTransient<SystemStore>()
@@ -130,6 +131,8 @@ namespace PluralKit.Bot
             .AddSingleton(svc => new LoggerProvider(svc.GetRequiredService<CoreConfig>(), "bot"))
             .AddScoped(svc => svc.GetRequiredService<LoggerProvider>().RootLogger.ForContext("EventId", svc.GetRequiredService<EventIdProvider>().EventId))
             
+            .AddMemoryCache()
+            
             .BuildServiceProvider();
     }
     class Bot

PluralKit.Bot/Commands/MemberCommands.cs

@@ -16,6 +16,8 @@ namespace PluralKit.Bot.Commands
         public MemberStore Members { get; set; }
         public EmbedService Embeds { get; set;  }
         
+        public ProxyCacheService ProxyCache { get; set; }
+
         public override string Prefix => "member";
         public override string ContextNoun => "member";
 
@@ -170,6 +172,8 @@ namespace PluralKit.Bot.Commands
             ContextEntity.Suffix = prefixAndSuffix[1].Length > 0 ? prefixAndSuffix[1] : null;
             await Members.Save(ContextEntity);
             await Context.Channel.SendMessageAsync($"{Emojis.Success} Member proxy tags changed to `{ContextEntity.ProxyString.Sanitize()}`. Try proxying now!");
+            
+            ProxyCache.InvalidateResultsForSystem(Context.SenderSystem);
         }
 
         [Command("delete")]

PluralKit.Bot/PluralKit.Bot.csproj

@@ -14,8 +14,8 @@
       <PackageReference Include="Discord.Net.Webhook" Version="2.1.1" />
       <PackageReference Include="Discord.Net.WebSocket" Version="2.1.1" />
       <PackageReference Include="Humanizer.Core" Version="2.6.2" />
+      <PackageReference Include="Microsoft.Extensions.Caching.Memory" Version="2.2.0" />
       <PackageReference Include="Sentry" Version="2.0.0-beta2" />
       <PackageReference Include="SixLabors.ImageSharp" Version="1.0.0-beta0006" />
     </ItemGroup>
-
 </Project>

PluralKit.Bot/Services/ProxyCacheService.cs

@@ -0,0 +1,82 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using Dapper;
+using Microsoft.Extensions.Caching.Memory;
+using Serilog;
+
+namespace PluralKit.Bot
+{
+    public class ProxyCacheService
+    {
+        public class ProxyDatabaseResult
+        {
+            public PKSystem System;
+            public PKMember Member;
+        }
+        
+        private DbConnectionFactory _conn;
+        private IMemoryCache _cache;
+        private ILogger _logger;
+
+        public ProxyCacheService(DbConnectionFactory conn, IMemoryCache cache, ILogger logger)
+        {
+            _conn = conn;
+            _cache = cache;
+            _logger = logger;
+        }
+
+        public Task<IEnumerable<ProxyDatabaseResult>> GetResultsFor(ulong account)
+        {
+            _logger.Debug("Looking up members for account {Account} in cache...", account);
+            return _cache.GetOrCreateAsync(GetKey(account), (entry) => FetchResults(account, entry));
+        }
+
+        public void InvalidateResultsFor(ulong account)
+        {
+            _logger.Information("Invalidating proxy cache for account {Account}", account);
+            _cache.Remove(GetKey(account));
+        }
+        
+        public async Task InvalidateResultsForSystem(PKSystem system)
+        {
+            _logger.Information("Invalidating proxy cache for system {System}", system.Id);
+            using (var conn = await _conn.Obtain())
+                foreach (var accountId in await conn.QueryAsync<ulong>("select uid from accounts where system = @Id", system))
+                    _cache.Remove(GetKey(accountId));
+        }
+
+        private async Task<IEnumerable<ProxyDatabaseResult>> FetchResults(ulong account, ICacheEntry entry)
+        {
+            _logger.Information("Members for account {Account} not in cache, fetching", account);
+            using (var conn = await _conn.Obtain())
+            {
+                var results = (await conn.QueryAsync<PKMember, PKSystem, ProxyDatabaseResult>(
+                    "select members.*, systems.* from members, systems, accounts where members.system = systems.id and accounts.system = systems.id and accounts.uid = @Uid",
+                    (member, system) =>
+                        new ProxyDatabaseResult {Member = member, System = system}, new {Uid = account})).ToList();
+
+                if (results.Count == 0)
+                {
+                    // Long expiry for accounts with no system registered
+                    entry.SetSlidingExpiration(TimeSpan.FromMinutes(5));
+                    entry.SetAbsoluteExpiration(TimeSpan.FromHours(1));
+                }
+                else
+                {
+                    // Shorter expiry if they already have a system
+                    entry.SetSlidingExpiration(TimeSpan.FromMinutes(1));
+                    entry.SetAbsoluteExpiration(TimeSpan.FromMinutes(5));
+                }
+                
+                return results;
+            }
+        }
+
+        private object GetKey(ulong account)
+        {
+            return $"_proxy_account_{account}";
+        }
+    }
+}

PluralKit.Bot/Services/ProxyService.cs

@@ -10,16 +10,11 @@ using Discord;
 using Discord.Net;
 using Discord.Webhook;
 using Discord.WebSocket;
+using Microsoft.Extensions.Caching.Memory;
 using Serilog;
 
 namespace PluralKit.Bot
 {
-    class ProxyDatabaseResult
-    {
-        public PKSystem System;
-        public PKMember Member;
-    }
-
     class ProxyMatch {
         public PKMember Member;
         public PKSystem System;
@@ -35,10 +30,11 @@ namespace PluralKit.Bot
         private EmbedService _embeds;
         private IMetrics _metrics;
         private ILogger _logger;
+        private ProxyCacheService _cache;
 
         private HttpClient _httpClient;
 
-        public ProxyService(IDiscordClient client, WebhookCacheService webhookCache, DbConnectionFactory conn, LogChannelService logChannel, MessageStore messageStorage, EmbedService embeds, IMetrics metrics, ILogger logger)
+        public ProxyService(IDiscordClient client, WebhookCacheService webhookCache, DbConnectionFactory conn, LogChannelService logChannel, MessageStore messageStorage, EmbedService embeds, IMetrics metrics, ILogger logger, ProxyCacheService cache)
         {
             _client = client;
             _webhookCache = webhookCache;
@@ -47,12 +43,13 @@ namespace PluralKit.Bot
             _messageStorage = messageStorage;
             _embeds = embeds;
             _metrics = metrics;
+            _cache = cache;
             _logger = logger.ForContext<ProxyService>();
 
             _httpClient = new HttpClient();
         }
 
-        private ProxyMatch GetProxyTagMatch(string message, IEnumerable<ProxyDatabaseResult> potentials)
+        private ProxyMatch GetProxyTagMatch(string message, IEnumerable<ProxyCacheService.ProxyDatabaseResult> potentials)
         {
             // If the message starts with a @mention, and then proceeds to have proxy tags,
             // extract the mention and place it inside the inner message
@@ -89,14 +86,7 @@ namespace PluralKit.Bot
             // Bail early if this isn't in a guild channel
             if (!(message.Channel is IGuildChannel)) return;
 
-            IEnumerable<ProxyDatabaseResult> results;
+            var results = await _cache.GetResultsFor(message.Author.Id);
-            using (var conn = await _conn.Obtain())
-            {
-                results = await conn.QueryAsync<PKMember, PKSystem, ProxyDatabaseResult>(
-                    "select members.*, systems.* from members, systems, accounts where members.system = systems.id and accounts.system = systems.id and accounts.uid = @Uid",
-                    (member, system) =>
-                        new ProxyDatabaseResult {Member = member, System = system}, new {Uid = message.Author.Id});
-            }
 
             // Find a member with proxy tags matching the message
             var match = GetProxyTagMatch(message.Content, results);
@@ -134,7 +124,6 @@ namespace PluralKit.Bot
                 await message.DeleteAsync();
             } catch (HttpException) {} // If it's already deleted, we just swallow the exception
         }
-
         private static string SanitizeEveryoneMaybe(IMessage message, string messageContents)
         {
             var senderPermissions = ((IGuildUser) message.Author).GetPermissions(message.Channel as IGuildChannel);