Fix database connection pool contention (maybe)

Instead of acquiring a connection per service per request, we
acquire connections more often but at a more granular level, meaning
they're also disposed of more quickly instead of staying for a long time
in case of long-running commands or leaks.

PluralKit.API/Controllers/SystemController.cs

@@ -36,10 +36,10 @@ namespace PluralKit.API.Controllers
         private SystemStore _systems;
         private MemberStore _members;
         private SwitchStore _switches;
-        private IDbConnection _conn;
+        private DbConnectionFactory _conn;
         private TokenAuthService _auth;
 
-        public SystemController(SystemStore systems, MemberStore members, SwitchStore switches, IDbConnection conn, TokenAuthService auth)
+        public SystemController(SystemStore systems, MemberStore members, SwitchStore switches, DbConnectionFactory conn, TokenAuthService auth)
         {
             _systems = systems;
             _members = members;
@@ -74,7 +74,9 @@ namespace PluralKit.API.Controllers
             var system = await _systems.GetByHid(hid);
             if (system == null) return NotFound("System not found.");
 
-            var res = await _conn.QueryAsync<SwitchesReturn>(
+            using (var conn = _conn.Obtain())
+            {
+                var res = await conn.QueryAsync<SwitchesReturn>(
                     @"select *, array(
                         select members.hid from switch_members, members
                         where switch_members.switch = switches.id and members.id = switch_members.member
@@ -84,6 +86,7 @@ namespace PluralKit.API.Controllers
                     limit 100;", new {System = system.Id, Before = before});
                 return Ok(res);
             }
+        }
 
         [HttpGet("{hid}/fronters")]
         public async Task<ActionResult<FrontersReturn>> GetFronters(string hid)
@@ -142,7 +145,10 @@ namespace PluralKit.API.Controllers
                 return BadRequest("New members identical to existing fronters.");
             
             // Resolve member objects for all given IDs
-            var membersList = (await _conn.QueryAsync<PKMember>("select * from members where hid = any(@Hids)", new {Hids = param.Members})).ToList();
+            IEnumerable<PKMember> membersList;
+            using (var conn = _conn.Obtain())
+                membersList = (await conn.QueryAsync<PKMember>("select * from members where hid = any(@Hids)", new {Hids = param.Members})).ToList();
+            
             foreach (var member in membersList)
                 if (member.System != _auth.CurrentSystem.Id)
                     return BadRequest($"Cannot switch to member '{member.Hid}' not in system.");

PluralKit.API/Startup.cs

@@ -43,12 +43,7 @@ namespace PluralKit.API
                 .AddScoped<TokenAuthService>()
 
                 .AddTransient(_ => Configuration.GetSection("PluralKit").Get<CoreConfig>() ?? new CoreConfig())
-                .AddScoped<IDbConnection>(svc =>
+                .AddSingleton(svc => new DbConnectionFactory(svc.GetRequiredService<CoreConfig>().Database));
-                {
-                    var conn = new NpgsqlConnection(svc.GetRequiredService<CoreConfig>().Database);
-                    conn.Open();
-                    return conn;
-                });
         }
 
         // This method gets called by the runtime. Use this method to configure the HTTP request pipeline.

PluralKit.Bot/Bot.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Data;
+using System.Data.Common;
 using System.Diagnostics;
 using System.IO;
 using System.Linq;
@@ -32,8 +33,8 @@ namespace PluralKit.Bot
             using (var services = BuildServiceProvider())
             {
                 Console.WriteLine("- Connecting to database...");
-                var connection = services.GetRequiredService<IDbConnection>() as NpgsqlConnection;
+                using (var conn = services.GetRequiredService<DbConnectionFactory>().Obtain())
-                await Schema.CreateTables(connection);
+                    await Schema.CreateTables(conn);
 
                 Console.WriteLine("- Connecting to Discord...");
                 var client = services.GetRequiredService<IDiscordClient>() as DiscordSocketClient;
@@ -51,13 +52,7 @@ namespace PluralKit.Bot
                 .AddTransient(_ => _config.GetSection("PluralKit").Get<CoreConfig>() ?? new CoreConfig())
                 .AddTransient(_ => _config.GetSection("PluralKit").GetSection("Bot").Get<BotConfig>() ?? new BotConfig())
                 
-                .AddTransient<IDbConnection>(svc =>
+                .AddTransient(svc => new DbConnectionFactory(svc.GetRequiredService<CoreConfig>().Database))
-                {
-                    
-                    var conn = new NpgsqlConnection(svc.GetRequiredService<CoreConfig>().Database);
-                    conn.Open();
-                    return conn;
-                })
                 
                 .AddSingleton<IDiscordClient, DiscordSocketClient>()
                 .AddSingleton<Bot>()
@@ -170,9 +165,10 @@ namespace PluralKit.Bot
                 // If it does, fetch the sender's system (because most commands need that) into the context,
                 // and start command execution
                 // Note system may be null if user has no system, hence `OrDefault`
-                var connection = serviceScope.ServiceProvider.GetService<IDbConnection>();
+                PKSystem system;
-                var system = await connection.QueryFirstOrDefaultAsync<PKSystem>("select systems.* from systems, accounts where accounts.uid = @Id and systems.id = accounts.system", new { Id = arg.Author.Id });
+                using (var conn = serviceScope.ServiceProvider.GetService<DbConnectionFactory>().Obtain())
-                await _commands.ExecuteAsync(new PKCommandContext(_client, arg, connection, system), argPos, serviceScope.ServiceProvider);
+                    system = await conn.QueryFirstOrDefaultAsync<PKSystem>("select systems.* from systems, accounts where accounts.uid = @Id and systems.id = accounts.system", new { Id = arg.Author.Id });
+                await _commands.ExecuteAsync(new PKCommandContext(_client, arg, system), argPos, serviceScope.ServiceProvider);
             }
             else
             {

PluralKit.Bot/Services/LogChannelService.cs

@@ -11,13 +11,13 @@ namespace PluralKit.Bot {
 
     public class LogChannelService {
         private IDiscordClient _client;
-        private IDbConnection _connection;
+        private DbConnectionFactory _conn;
         private EmbedService _embed;
 
-        public LogChannelService(IDiscordClient client, IDbConnection connection, EmbedService embed)
+        public LogChannelService(IDiscordClient client, DbConnectionFactory conn, EmbedService embed)
         {
             this._client = client;
-            this._connection = connection;
+            this._conn = conn;
             this._embed = embed;
         }
 
@@ -30,10 +30,15 @@ namespace PluralKit.Bot {
         }
 
         public async Task<ITextChannel> GetLogChannel(IGuild guild) {
-            var server = await _connection.QueryFirstOrDefaultAsync<ServerDefinition>("select * from servers where id = @Id", new { Id = guild.Id });
+            using (var conn = _conn.Obtain())
+            {
+                var server =
+                    await conn.QueryFirstOrDefaultAsync<ServerDefinition>("select * from servers where id = @Id",
+                        new {Id = guild.Id});
                 if (server?.LogChannel == null) return null;
                 return await _client.GetChannelAsync(server.LogChannel.Value) as ITextChannel;
             }
+        }
 
         public async Task SetLogChannel(IGuild guild, ITextChannel newLogChannel) {
             var def = new ServerDefinition {
@@ -41,7 +46,12 @@ namespace PluralKit.Bot {
                 LogChannel = newLogChannel?.Id
             };
 
-            await _connection.QueryAsync("insert into servers (id, log_channel) values (@Id, @LogChannel) on conflict (id) do update set log_channel = @LogChannel", def);
+            using (var conn = _conn.Obtain())
+            {
+                await conn.QueryAsync(
+                    "insert into servers (id, log_channel) values (@Id, @LogChannel) on conflict (id) do update set log_channel = @LogChannel",
+                    def);
+            }
         }
     }
 }

PluralKit.Bot/Services/ProxyService.cs

@@ -28,17 +28,17 @@ namespace PluralKit.Bot
 
     class ProxyService {
         private IDiscordClient _client;
-        private IDbConnection _connection;
+        private DbConnectionFactory _conn;
         private LogChannelService _logger;
         private WebhookCacheService _webhookCache;
         private MessageStore _messageStorage;
         private EmbedService _embeds;
         
-        public ProxyService(IDiscordClient client, WebhookCacheService webhookCache, IDbConnection connection, LogChannelService logger, MessageStore messageStorage, EmbedService embeds)
+        public ProxyService(IDiscordClient client, WebhookCacheService webhookCache, DbConnectionFactory conn, LogChannelService logger, MessageStore messageStorage, EmbedService embeds)
         {
             _client = client;
             _webhookCache = webhookCache;
-            _connection = connection;
+            _conn = conn;
             _logger = logger;
             _messageStorage = messageStorage;
             _embeds = embeds;
@@ -76,11 +76,16 @@ namespace PluralKit.Bot
             return null;
         }
 
-        public async Task HandleMessageAsync(IMessage message) {
+        public async Task HandleMessageAsync(IMessage message)
-            var results = await _connection.QueryAsync<PKMember, PKSystem, ProxyDatabaseResult>(
+        {
+            IEnumerable<ProxyDatabaseResult> results;
+            using (var conn = _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);

PluralKit.Bot/Utils.cs

@@ -152,14 +152,12 @@ namespace PluralKit.Bot
     /// Subclass of ICommandContext with PK-specific additional fields and functionality
     public class PKCommandContext : SocketCommandContext
     {
-        public IDbConnection Connection { get; }
         public PKSystem SenderSystem { get; }
         
         private object _entity;
 
-        public PKCommandContext(DiscordSocketClient client, SocketUserMessage msg, IDbConnection connection, PKSystem system) : base(client, msg)
+        public PKCommandContext(DiscordSocketClient client, SocketUserMessage msg, PKSystem system) : base(client, msg)
         {
-            Connection = connection;
             SenderSystem = system;
         }
 

PluralKit.Core/Stores.cs

@@ -10,67 +10,80 @@ using NodaTime;
 
 namespace PluralKit {
     public class SystemStore {
-        private IDbConnection conn;
+        private DbConnectionFactory _conn;
 
-        public SystemStore(IDbConnection conn) {
+        public SystemStore(DbConnectionFactory conn) {
-            this.conn = conn;
+            this._conn = conn;
         }
         
         public async Task<PKSystem> Create(string systemName = null) {
             // TODO: handle HID collision case
             var hid = Utils.GenerateHid();
+            
+            using (var conn = _conn.Obtain())
                 return await conn.QuerySingleAsync<PKSystem>("insert into systems (hid, name) values (@Hid, @Name) returning *", new { Hid = hid, Name = systemName });
         }
 
         public async Task Link(PKSystem system, ulong accountId) {
+            using (var conn = _conn.Obtain())
                 await conn.ExecuteAsync("insert into accounts (uid, system) values (@Id, @SystemId)", new { Id = accountId, SystemId = system.Id });
         }
         
         public async Task Unlink(PKSystem system, ulong accountId) {
+            using (var conn = _conn.Obtain())
                 await conn.ExecuteAsync("delete from accounts where uid = @Id and system = @SystemId", new { Id = accountId, SystemId = system.Id });
         }
 
         public async Task<PKSystem> GetByAccount(ulong accountId) {
+            using (var conn = _conn.Obtain())
                 return await conn.QuerySingleOrDefaultAsync<PKSystem>("select systems.* from systems, accounts where accounts.system = systems.id and accounts.uid = @Id", new { Id = accountId });
         }
 
         public async Task<PKSystem> GetByHid(string hid) {
+            using (var conn = _conn.Obtain())
                 return await conn.QuerySingleOrDefaultAsync<PKSystem>("select * from systems where systems.hid = @Hid", new { Hid = hid.ToLower() });
         }
 
         public async Task<PKSystem> GetByToken(string token) {
+            using (var conn = _conn.Obtain())
                 return await conn.QuerySingleOrDefaultAsync<PKSystem>("select * from systems where token = @Token", new { Token = token });
         }
         
         public async Task<PKSystem> GetById(int id)
         {
+            using (var conn = _conn.Obtain())
                 return await conn.QuerySingleOrDefaultAsync<PKSystem>("select * from systems where id = @Id", new { Id = id });
         }
 
         public async Task Save(PKSystem system) {
+            using (var conn = _conn.Obtain())
                 await conn.ExecuteAsync("update systems set name = @Name, description = @Description, tag = @Tag, avatar_url = @AvatarUrl, token = @Token, ui_tz = @UiTz where id = @Id", system);
         }
 
         public async Task Delete(PKSystem system) {
+            using (var conn = _conn.Obtain())
                 await conn.ExecuteAsync("delete from systems where id = @Id", system);
         }        
 
         public async Task<IEnumerable<ulong>> GetLinkedAccountIds(PKSystem system)
         {
+            using (var conn = _conn.Obtain())
                 return await conn.QueryAsync<ulong>("select uid from accounts where system = @Id", new { Id = system.Id });
         }
     }
 
     public class MemberStore {
-        private IDbConnection conn;
+        private DbConnectionFactory _conn;
 
-        public MemberStore(IDbConnection conn) {
+        public MemberStore(DbConnectionFactory conn) {
-            this.conn = conn;
+            this._conn = conn;
         }
 
         public async Task<PKMember> Create(PKSystem system, string name) {
             // TODO: handle collision
             var hid = Utils.GenerateHid();
+            
+            using (var conn = _conn.Obtain())
                 return await conn.QuerySingleAsync<PKMember>("insert into members (hid, system, name) values (@Hid, @SystemId, @Name) returning *", new {
                     Hid = hid,
                     SystemID = system.Id,
@@ -79,11 +92,13 @@ namespace PluralKit {
         }
 
         public async Task<PKMember> GetByHid(string hid) {
+            using (var conn = _conn.Obtain())
                 return await conn.QuerySingleOrDefaultAsync<PKMember>("select * from members where hid = @Hid", new { Hid = hid.ToLower() });
         }
 
         public async Task<PKMember> GetByName(PKSystem system, string name) {
             // QueryFirst, since members can (in rare cases) share names
+            using (var conn = _conn.Obtain())
                 return await conn.QueryFirstOrDefaultAsync<PKMember>("select * from members where lower(name) = lower(@Name) and system = @SystemID", new { Name = name, SystemID = system.Id });
         }
 
@@ -96,19 +111,23 @@ namespace PluralKit {
         }
 
         public async Task<IEnumerable<PKMember>> GetBySystem(PKSystem system) {
+            using (var conn = _conn.Obtain())
                 return await conn.QueryAsync<PKMember>("select * from members where system = @SystemID", new { SystemID = system.Id });
         }
 
         public async Task Save(PKMember member) {
+            using (var conn = _conn.Obtain())
                 await conn.ExecuteAsync("update members set name = @Name, description = @Description, color = @Color, avatar_url = @AvatarUrl, birthday = @Birthday, pronouns = @Pronouns, prefix = @Prefix, suffix = @Suffix where id = @Id", member);
         }
 
         public async Task Delete(PKMember member) {
+            using (var conn = _conn.Obtain())
                 await conn.ExecuteAsync("delete from members where id = @Id", member);
         }
 
         public async Task<int> MessageCount(PKMember member)
         {
+            using (var conn = _conn.Obtain())
                 return await conn.QuerySingleAsync<int>("select count(*) from messages where member = @Id", member);
         }
     }
@@ -127,14 +146,15 @@ namespace PluralKit {
             public PKSystem System;
         }
 
-        private IDbConnection _connection;
+        private DbConnectionFactory _conn;
 
-        public MessageStore(IDbConnection connection) {
+        public MessageStore(DbConnectionFactory conn) {
-            this._connection = connection;
+            this._conn = conn;
         }
 
         public async Task Store(ulong senderId, ulong messageId, ulong channelId, PKMember member) {
-            await _connection.ExecuteAsync("insert into messages(mid, channel, member, sender) values(@MessageId, @ChannelId, @MemberId, @SenderId)", new {
+            using (var conn = _conn.Obtain())
+                await conn.ExecuteAsync("insert into messages(mid, channel, member, sender) values(@MessageId, @ChannelId, @MemberId, @SenderId)", new {
                     MessageId = messageId,
                     ChannelId = channelId,
                     MemberId = member.Id,
@@ -144,7 +164,8 @@ namespace PluralKit {
 
         public async Task<StoredMessage> Get(ulong id)
         {
-            return (await _connection.QueryAsync<PKMessage, PKMember, PKSystem, StoredMessage>("select messages.*, members.*, systems.* from messages, members, systems where mid = @Id and messages.member = members.id and systems.id = members.system", (msg, member, system) => new StoredMessage
+            using (var conn = _conn.Obtain())
+                return (await conn.QueryAsync<PKMessage, PKMember, PKSystem, StoredMessage>("select messages.*, members.*, systems.* from messages, members, systems where mid = @Id and messages.member = members.id and systems.id = members.system", (msg, member, system) => new StoredMessage
                 {
                     Message = msg,
                     System = system,
@@ -153,33 +174,35 @@ namespace PluralKit {
         }
         
         public async Task Delete(ulong id) {
-            await _connection.ExecuteAsync("delete from messages where mid = @Id", new { Id = id });
+            using (var conn = _conn.Obtain())
+                await conn.ExecuteAsync("delete from messages where mid = @Id", new { Id = id });
         }
     }
 
     public class SwitchStore
     {
-        private IDbConnection _connection;
+        private DbConnectionFactory _conn;
 
-        public SwitchStore(IDbConnection connection)
+        public SwitchStore(DbConnectionFactory conn)
         {
-            _connection = connection;
+            _conn = conn;
         }
 
         public async Task RegisterSwitch(PKSystem system, IEnumerable<PKMember> members)
         {
             // Use a transaction here since we're doing multiple executed commands in one
-            using (var tx = _connection.BeginTransaction())
+            using (var conn = _conn.Obtain())
+            using (var tx = conn.BeginTransaction())
             {
                 // First, we insert the switch itself
-                var sw = await _connection.QuerySingleAsync<PKSwitch>("insert into switches(system) values (@System) returning *",
+                var sw = await conn.QuerySingleAsync<PKSwitch>("insert into switches(system) values (@System) returning *",
                     new {System = system.Id});
                 
                 // Then we insert each member in the switch in the switch_members table
                 // TODO: can we parallelize this or send it in bulk somehow?
                 foreach (var member in members)
                 {
-                    await _connection.ExecuteAsync(
+                    await conn.ExecuteAsync(
                         "insert into switch_members(switch, member) values(@Switch, @Member)",
                         new {Switch = sw.Id, Member = member.Id});
                 }
@@ -193,18 +216,21 @@ namespace PluralKit {
         {
             // TODO: refactor the PKSwitch data structure to somehow include a hydrated member list
             // (maybe when we get caching in?)
-            return await _connection.QueryAsync<PKSwitch>("select * from switches where system = @System order by timestamp desc limit @Count", new {System = system.Id, Count = count});
+            using (var conn = _conn.Obtain())
+                return await conn.QueryAsync<PKSwitch>("select * from switches where system = @System order by timestamp desc limit @Count", new {System = system.Id, Count = count});
         }
 
         public async Task<IEnumerable<int>> GetSwitchMemberIds(PKSwitch sw)
         {
-            return await _connection.QueryAsync<int>("select member from switch_members where switch = @Switch",
+            using (var conn = _conn.Obtain())
+                return await conn.QueryAsync<int>("select member from switch_members where switch = @Switch",
                     new {Switch = sw.Id});
         }
         
         public async Task<IEnumerable<PKMember>> GetSwitchMembers(PKSwitch sw)
         {
-            return await _connection.QueryAsync<PKMember>(
+            using (var conn = _conn.Obtain())
+                return await conn.QueryAsync<PKMember>(
                     "select * from switch_members, members where switch_members.member = members.id and switch_members.switch = @Switch",
                     new {Switch = sw.Id});
         }
@@ -213,13 +239,15 @@ namespace PluralKit {
 
         public async Task MoveSwitch(PKSwitch sw, Instant time)
         {
-            await _connection.ExecuteAsync("update switches set timestamp = @Time where id = @Id",
+            using (var conn = _conn.Obtain())
+                await conn.ExecuteAsync("update switches set timestamp = @Time where id = @Id",
                     new {Time = time, Id = sw.Id});
         }
 
         public async Task DeleteSwitch(PKSwitch sw)
         {
-            await _connection.ExecuteAsync("delete from switches where id = @Id", new {Id = sw.Id});
+            using (var conn = _conn.Obtain())
+                await conn.ExecuteAsync("delete from switches where id = @Id", new {Id = sw.Id});
         }
         
         public struct SwitchListEntry
@@ -242,10 +270,14 @@ namespace PluralKit {
             // query DB for all members involved in any of the switches above and collect into a dictionary for future use
             // this makes sure the return list has the same instances of PKMember throughout, which is important for the dictionary
             // key used in GetPerMemberSwitchDuration below
-            var memberObjects = (await _connection.QueryAsync<PKMember>(
+            Dictionary<int, PKMember> memberObjects;
+            using (var conn = _conn.Obtain())
+            {
+                memberObjects = (await conn.QueryAsync<PKMember>(
                         "select distinct members.* from members, switch_members where switch_members.switch = any(@Switches) and switch_members.member = members.id", // lol postgres specific `= any()` syntax
                         new {Switches = switchesInRange.Select(sw => sw.Id).ToList()}))
                     .ToDictionary(m => m.Id);
+            }
 
 
             // we create the entry objects

PluralKit.Core/Utils.cs

@@ -333,4 +333,19 @@ namespace PluralKit
             return (T) value;
         }
     }
+    
+    public class DbConnectionFactory
+    {
+        private string _connectionString;
+
+        public DbConnectionFactory(string connectionString)
+        {
+            _connectionString = connectionString;
+        }
+
+        public IDbConnection Obtain()
+        {
+            return new NpgsqlConnection(_connectionString);
+        }
+    }
 }