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

Fix database connection pool contention (maybe)

Browse Source 
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.
This commit is contained in:
Ske 2019-07-11 21:25:23 +02:00
parent ca56fd419b
commit d829630a35
8 changed files with 166 additions and 109 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
PluralKit.API/Controllers/SystemController.cs
View File

@ -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.");

7
PluralKit.API/Startup.cs
View File

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

20
PluralKit.Bot/Bot.cs
View File

@ -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;
await Schema.CreateTables(connection);
using (var conn = services.GetRequiredService<DbConnectionFactory>().Obtain())
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 =>
{
var conn = new NpgsqlConnection(svc.GetRequiredService<CoreConfig>().Database);
conn.Open();
return conn;
})
.AddTransient(svc => new DbConnectionFactory(svc.GetRequiredService<CoreConfig>().Database))
.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>();
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 });
await _commands.ExecuteAsync(new PKCommandContext(_client, arg, connection, system), argPos, serviceScope.ServiceProvider);
PKSystem system;
using (var conn = serviceScope.ServiceProvider.GetService<DbConnectionFactory>().Obtain())
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
{

20
PluralKit.Bot/Services/LogChannelService.cs
View File

@ -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);
}
}
}
}

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

@ -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) {
var results = await _connection.QueryAsync<PKMember, PKSystem, ProxyDatabaseResult>(
public async Task HandleMessageAsync(IMessage message)
{
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);

4
PluralKit.Bot/Utils.cs
View File

@ -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;
}

80
PluralKit.Core/Stores.cs
View File

@ -10,67 +10,80 @@ using NodaTime;
namespace PluralKit {
public class SystemStore {
private IDbConnection conn;
private DbConnectionFactory _conn;
public SystemStore(IDbConnection conn) {
this.conn = conn;
public SystemStore(DbConnectionFactory 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) {
this.conn = conn;
public MemberStore(DbConnectionFactory 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) {
this._connection = connection;
public MessageStore(DbConnectionFactory conn) {
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

15
PluralKit.Core/Utils.cs
View File

@ -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);
}
}
}