From 397da2e1fa6481b6bb406eb75ec6d1258c8aaa99 Mon Sep 17 00:00:00 2001 From: Noko Date: Sun, 20 Oct 2019 02:16:57 -0500 Subject: [PATCH 1/4] Added max member count to limits A given system can now have up to 1000 members. Within 50 members of that limit, a warning will display whenever a new member is created via the bot. Once the limit is reached, a final warning will appear indicating that no additional members can be created unless members are first deleted. Attempting to create a new member at that point by any method will result in an error message indicating that the limit has been reached. Respecting this in pk;import required some restructuring to tease apart which members already exist and which ones need to be created prior to creating any members as it seems preferable to fail early and give the user the ability to intervene rather than pushing the system to the member cap and requiring manual deletion of "lower priority" members before others can be created. One consequence of the restructure is that existing members are being read in bulk which is a performance improvement of 25-70% depending on how many switches need to be imported (the more members you have, the more noticeable this is). --- PluralKit.API/Controllers/MemberController.cs | 5 + .../Commands/ImportExportCommands.cs | 5 +- PluralKit.Bot/Commands/MemberCommands.cs | 15 ++- PluralKit.Bot/Errors.cs | 3 +- PluralKit.Core/DataFiles.cs | 110 +++++++++--------- PluralKit.Core/Limits.cs | 2 + PluralKit.Core/Stores.cs | 31 +++++ 7 files changed, 111 insertions(+), 60 deletions(-) diff --git a/PluralKit.API/Controllers/MemberController.cs b/PluralKit.API/Controllers/MemberController.cs index 3bf30976..f420b7f9 100644 --- a/PluralKit.API/Controllers/MemberController.cs +++ b/PluralKit.API/Controllers/MemberController.cs @@ -38,6 +38,11 @@ namespace PluralKit.API.Controllers if (newMember.Name == null) return BadRequest("Member name cannot be null."); + // Enforce per-system member limit + var memberCount = await _members.MemberCount(system); + if (memberCount >= Limits.MaxMemberCount) + return BadRequest($"Member limit reached ({memberCount} / {Limits.MaxMemberCount})."); + // Explicit bounds checks if (newMember.Name != null && newMember.Name.Length > Limits.MaxMemberNameLength) return BadRequest($"Member name too long ({newMember.Name.Length} > {Limits.MaxMemberNameLength}."); diff --git a/PluralKit.Bot/Commands/ImportExportCommands.cs b/PluralKit.Bot/Commands/ImportExportCommands.cs index 1d138601..e3c05264 100644 --- a/PluralKit.Bot/Commands/ImportExportCommands.cs +++ b/PluralKit.Bot/Commands/ImportExportCommands.cs @@ -101,8 +101,9 @@ namespace PluralKit.Bot.Commands // If passed system is null, it'll create a new one // (and that's okay!) var result = await _dataFiles.ImportSystem(data, ctx.System, ctx.Author.Id); - - if (ctx.System != null) + if (!result.Success) + await ctx.Reply($"{Emojis.Error} The provided system profile could not be imported. {result.Message}"); + else if (ctx.System != null) { await ctx.Reply($"{Emojis.Success} PluralKit has created a system for you based on the given file. Your system ID is `{result.System.Hid}`. Type `pk;system` for more information."); } diff --git a/PluralKit.Bot/Commands/MemberCommands.cs b/PluralKit.Bot/Commands/MemberCommands.cs index 8961b1f0..dbb80f59 100644 --- a/PluralKit.Bot/Commands/MemberCommands.cs +++ b/PluralKit.Bot/Commands/MemberCommands.cs @@ -45,13 +45,24 @@ namespace PluralKit.Bot.Commands if (!await ctx.PromptYesNo(msg)) throw new PKError("Member creation cancelled."); } + // Enforce per-system member limit + var memberCount = await _members.MemberCount(ctx.System); + if (memberCount >= Limits.MaxMemberCount) + throw Errors.MemberLimitReachedError; + // Create the member var member = await _members.Create(ctx.System, memberName); + memberCount++; // Send confirmation and space hint await ctx.Reply($"{Emojis.Success} Member \"{memberName.SanitizeMentions()}\" (`{member.Hid}`) registered! See the user guide for commands for editing this member: https://pluralkit.me/guide#member-management"); - if (memberName.Contains(" ")) await ctx.Reply($"{Emojis.Note} Note that this member's name contains spaces. You will need to surround it with \"double quotes\" when using commands referring to it, or just use the member's 5-character ID (which is `{member.Hid}`)."); - + if (memberName.Contains(" ")) + await ctx.Reply($"{Emojis.Note} Note that this member's name contains spaces. You will need to surround it with \"double quotes\" when using commands referring to it, or just use the member's 5-character ID (which is `{member.Hid}`)."); + if (memberCount >= Limits.MaxMemberCount) + await ctx.Reply($"{Emojis.Warn} You have reached the per-system member limit ({Limits.MaxMemberCount}). You will be unable to create additional members until existing members are deleted."); + else if (memberCount >= Limits.MaxMembersWarnThreshold) + await ctx.Reply($"{Emojis.Warn} You are approaching the per-system member limit ({memberCount} / {Limits.MaxMemberCount} members). Please review your member list for unused or duplicate members."); + await _proxyCache.InvalidateResultsForSystem(ctx.System); } diff --git a/PluralKit.Bot/Errors.cs b/PluralKit.Bot/Errors.cs index 2d5d79ed..74760569 100644 --- a/PluralKit.Bot/Errors.cs +++ b/PluralKit.Bot/Errors.cs @@ -22,7 +22,8 @@ namespace PluralKit.Bot { public static PKError DescriptionTooLongError(int length) => new PKError($"Description too long ({length}/{Limits.MaxDescriptionLength} characters)."); public static PKError MemberNameTooLongError(int length) => new PKError($"Member name too long ({length}/{Limits.MaxMemberNameLength} characters)."); public static PKError MemberPronounsTooLongError(int length) => new PKError($"Member pronouns too long ({length}/{Limits.MaxMemberNameLength} characters)."); - + public static PKError MemberLimitReachedError => new PKError($"System has reached the maximum number of members ({Limits.MaxMemberCount}). Please delete unused members first in order to create new ones."); + public static PKError InvalidColorError(string color) => new PKError($"\"{color.SanitizeMentions()}\" is not a valid color. Color must be in 6-digit RGB hex format (eg. #ff0000)."); public static PKError BirthdayParseError(string birthday) => new PKError($"\"{birthday.SanitizeMentions()}\" could not be parsed as a valid date. Try a format like \"2016-12-24\" or \"May 3 1996\"."); public static PKError ProxyMustHaveText => new PKSyntaxError("Example proxy message must contain the string 'text'."); diff --git a/PluralKit.Core/DataFiles.cs b/PluralKit.Core/DataFiles.cs index 8c839269..8e15c646 100644 --- a/PluralKit.Core/DataFiles.cs +++ b/PluralKit.Core/DataFiles.cs @@ -5,6 +5,7 @@ using System.Threading.Tasks; using Newtonsoft.Json; using NodaTime; using NodaTime.Text; +using PluralKit.Core; using Serilog; namespace PluralKit.Bot @@ -70,38 +71,22 @@ namespace PluralKit.Bot }; } - private async Task ExportMember(PKMember member) => new DataFileMember - { - Id = member.Hid, - Name = member.Name, - DisplayName = member.DisplayName, - Description = member.Description, - Birthday = member.Birthday != null ? Formats.DateExportFormat.Format(member.Birthday.Value) : null, - Pronouns = member.Pronouns, - Color = member.Color, - AvatarUrl = member.AvatarUrl, - Prefix = member.Prefix, - Suffix = member.Suffix, - Created = Formats.TimestampExportFormat.Format(member.Created), - MessageCount = await _members.MessageCount(member) - }; - - private async Task ExportSwitch(PKSwitch sw) => new DataFileSwitch - { - Members = (await _switches.GetSwitchMembers(sw)).Select(m => m.Hid).ToList(), - Timestamp = Formats.TimestampExportFormat.Format(sw.Timestamp) - }; - public async Task ImportSystem(DataFileSystem data, PKSystem system, ulong accountId) { // TODO: make atomic, somehow - we'd need to obtain one IDbConnection and reuse it // which probably means refactoring SystemStore.Save and friends etc - - var result = new ImportResult {AddedNames = new List(), ModifiedNames = new List()}; - var hidMapping = new Dictionary(); + var result = new ImportResult { + AddedNames = new List(), + ModifiedNames = new List(), + Success = true // Assume success unless indicated otherwise + }; + var dataFileToMemberMapping = new Dictionary(); + var unmappedMembers = new List(); // If we don't already have a system to save to, create one - if (system == null) system = await _systems.Create(data.Name); + if (system == null) + system = await _systems.Create(data.Name); + result.System = system; // Apply system info system.Name = data.Name; @@ -110,41 +95,56 @@ namespace PluralKit.Bot if (data.AvatarUrl != null) system.AvatarUrl = data.AvatarUrl; if (data.TimeZone != null) system.UiTz = data.TimeZone ?? "UTC"; await _systems.Save(system); - + // Make sure to link the sender account, too await _systems.Link(system, accountId); - // Apply members - // TODO: parallelize? - foreach (var dataMember in data.Members) + // Determine which members already exist and which ones need to be created + var existingMembers = await _members.GetBySystem(system); + foreach (var d in data.Members) { - // If member's given an ID, we try to look up the member with the given ID - PKMember member = null; - if (dataMember.Id != null) + // Try to look up the member with the given ID + var match = existingMembers.FirstOrDefault(m => m.Hid.Equals(d.Id)); + if (match == null) + match = existingMembers.FirstOrDefault(m => m.Name.Equals(d.Name)); // Try with the name instead + if (match != null) { - member = await _members.GetByHid(dataMember.Id); - - // ...but if it's a different system's member, we just make a new one anyway - if (member != null && member.System != system.Id) member = null; - } - - // Try to look up by name, too - if (member == null) member = await _members.GetByName(system, dataMember.Name); - - // And if all else fails (eg. fresh import from Tupperbox, etc) we just make a member lol - if (member == null) - { - member = await _members.Create(system, dataMember.Name); - result.AddedNames.Add(dataMember.Name); + dataFileToMemberMapping.Add(d.Id, match); // Relate the data file ID to the PKMember for importing switches + result.ModifiedNames.Add(d.Name); } else { - result.ModifiedNames.Add(dataMember.Name); + unmappedMembers.Add(d); // Track members that weren't found so we can create them all + result.AddedNames.Add(d.Name); } + } - // Keep track of what the data file's member ID maps to for switch import - if (!hidMapping.ContainsKey(dataMember.Id)) - hidMapping.Add(dataMember.Id, member); + // If creating the unmatched members would put us over the member limit, abort before creating any members + // new total: # in the system + (# in the file - # in the file that already exist) + if (data.Members.Count - dataFileToMemberMapping.Count + existingMembers.Count() >= Limits.MaxMemberCount) + { + result.Success = false; + result.Message = $"Import would exceed the maximum number of members ({Limits.MaxMemberCount})."; + result.AddedNames.Clear(); + result.ModifiedNames.Clear(); + return result; + } + + // Create all unmapped members in one transaction + // These consist of members from another PluralKit system or another framework (e.g. Tupperbox) + var membersToCreate = new Dictionary(); + unmappedMembers.ForEach(x => membersToCreate.Add(x.Id, x.Name)); + var newMembers = await _members.CreateMultiple(system, membersToCreate); + foreach (var member in newMembers) + dataFileToMemberMapping.Add(member.Key, member.Value); + + // Update members with data file properties + // TODO: parallelize? + foreach (var dataMember in data.Members) + { + dataFileToMemberMapping.TryGetValue(dataMember.Id, out PKMember member); + if (member == null) + continue; // Apply member info member.Name = dataMember.Name; @@ -161,7 +161,7 @@ namespace PluralKit.Bot if (dataMember.Birthday != null) { var birthdayParse = Formats.DateExportFormat.Parse(dataMember.Birthday); - member.Birthday = birthdayParse.Success ? (LocalDate?) birthdayParse.Value : null; + member.Birthday = birthdayParse.Success ? (LocalDate?)birthdayParse.Value : null; } await _members.Save(member); @@ -174,7 +174,7 @@ namespace PluralKit.Bot var timestamp = InstantPattern.ExtendedIso.Parse(sw.Timestamp).Value; var swMembers = new List(); swMembers.AddRange(sw.Members.Select(x => - hidMapping.FirstOrDefault(y => y.Key.Equals(x)).Value)); + dataFileToMemberMapping.FirstOrDefault(y => y.Key.Equals(x)).Value)); var mapped = new Tuple>(timestamp, swMembers); mappedSwitches.Add(mapped); } @@ -182,8 +182,6 @@ namespace PluralKit.Bot await _switches.RegisterSwitches(system, mappedSwitches); _logger.Information("Imported system {System}", system.Id); - - result.System = system; return result; } } @@ -193,6 +191,8 @@ namespace PluralKit.Bot public ICollection AddedNames; public ICollection ModifiedNames; public PKSystem System; + public bool Success; + public string Message; } public struct DataFileSystem diff --git a/PluralKit.Core/Limits.cs b/PluralKit.Core/Limits.cs index 8ced1346..5c1e5d8b 100644 --- a/PluralKit.Core/Limits.cs +++ b/PluralKit.Core/Limits.cs @@ -5,6 +5,8 @@ namespace PluralKit.Core { public static readonly int MaxSystemNameLength = 100; public static readonly int MaxSystemTagLength = MaxProxyNameLength - 1; + public static readonly int MaxMemberCount = 1000; + public static readonly int MaxMembersWarnThreshold = MaxMemberCount - 50; public static readonly int MaxDescriptionLength = 1000; public static readonly int MaxMemberNameLength = 100; // Fair bit larger than MaxProxyNameLength for bookkeeping public static readonly int MaxPronounsLength = 100; diff --git a/PluralKit.Core/Stores.cs b/PluralKit.Core/Stores.cs index 44aeeff0..22a978f5 100644 --- a/PluralKit.Core/Stores.cs +++ b/PluralKit.Core/Stores.cs @@ -128,6 +128,37 @@ namespace PluralKit { return member; } + public async Task> CreateMultiple(PKSystem system, Dictionary names) + { + using (var conn = await _conn.Obtain()) + using (var tx = conn.BeginTransaction()) + { + var results = new Dictionary(); + foreach (var name in names) + { + string hid; + do + { + hid = await conn.QuerySingleOrDefaultAsync("SELECT @Hid WHERE NOT EXISTS (SELECT id FROM members WHERE hid = @Hid LIMIT 1)", new + { + Hid = Utils.GenerateHid() + }); + } while (hid == null); + var member = await conn.QuerySingleAsync("INSERT INTO members (hid, system, name) VALUES (@Hid, @SystemId, @Name) RETURNING *", new + { + Hid = hid, + SystemID = system.Id, + Name = name.Value + }); + results.Add(name.Key, member); + } + + tx.Commit(); + _logger.Information("Created {MemberCount} members for system {SystemID}", names.Count(), system.Hid); + return results; + } + } + public async Task GetByHid(string hid) { using (var conn = await _conn.Obtain()) return await conn.QuerySingleOrDefaultAsync("select * from members where hid = @Hid", new { Hid = hid.ToLower() }); From 3d21adeec905617a1181214de0bf9ae3f3ecd6d7 Mon Sep 17 00:00:00 2001 From: Noko Date: Sun, 20 Oct 2019 02:22:22 -0500 Subject: [PATCH 2/4] Fix reversal of import response messages Some of the command rewrite changes resulted in the response messages for importing a system being swapped. When importing without an existing system (ctx.System == null), we want to display the "new system" message. Otherwise, show the count added/modified. --- PluralKit.Bot/Commands/ImportExportCommands.cs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/PluralKit.Bot/Commands/ImportExportCommands.cs b/PluralKit.Bot/Commands/ImportExportCommands.cs index e3c05264..d07bb964 100644 --- a/PluralKit.Bot/Commands/ImportExportCommands.cs +++ b/PluralKit.Bot/Commands/ImportExportCommands.cs @@ -103,12 +103,14 @@ namespace PluralKit.Bot.Commands var result = await _dataFiles.ImportSystem(data, ctx.System, ctx.Author.Id); if (!result.Success) await ctx.Reply($"{Emojis.Error} The provided system profile could not be imported. {result.Message}"); - else if (ctx.System != null) + else if (ctx.System == null) { + // We didn't have a system prior to importing, so give them the new system's ID await ctx.Reply($"{Emojis.Success} PluralKit has created a system for you based on the given file. Your system ID is `{result.System.Hid}`. Type `pk;system` for more information."); } else { + // We already had a system, so show them what changed await ctx.Reply($"{Emojis.Success} Updated {result.ModifiedNames.Count} members, created {result.AddedNames.Count} members. Type `pk;system list` to check!"); } } From 406f005b4f3f6cf087c10aafb0f3348383679180 Mon Sep 17 00:00:00 2001 From: Noko Date: Sun, 20 Oct 2019 14:38:43 -0500 Subject: [PATCH 3/4] Bulk import switches for pk;import We're now using binary import for switches and switch_members when importing a system profile, rather than importing them one switch at a time. This adds a pass-through method to the PerformanceTrackingConnection that can be used for other bulk import applications. --- PluralKit.Core/DataFiles.cs | 5 ++- PluralKit.Core/Stores.cs | 73 ++++++++++++++++++++++++++++++++++++- PluralKit.Core/Utils.cs | 5 +++ 3 files changed, 80 insertions(+), 3 deletions(-) diff --git a/PluralKit.Core/DataFiles.cs b/PluralKit.Core/DataFiles.cs index 8e15c646..b6e56945 100644 --- a/PluralKit.Core/DataFiles.cs +++ b/PluralKit.Core/DataFiles.cs @@ -179,9 +179,10 @@ namespace PluralKit.Bot mappedSwitches.Add(mapped); } // Import switches - await _switches.RegisterSwitches(system, mappedSwitches); + if (mappedSwitches.Any()) + await _switches.BulkImportSwitches(system, mappedSwitches); - _logger.Information("Imported system {System}", system.Id); + _logger.Information("Imported system {System}", system.Hid); return result; } } diff --git a/PluralKit.Core/Stores.cs b/PluralKit.Core/Stores.cs index 22a978f5..aaef353b 100644 --- a/PluralKit.Core/Stores.cs +++ b/PluralKit.Core/Stores.cs @@ -5,7 +5,7 @@ using System.Threading.Tasks; using App.Metrics.Logging; using Dapper; using NodaTime; - +using Npgsql; using PluralKit.Core; using Serilog; @@ -345,6 +345,77 @@ namespace PluralKit { } } + public async Task BulkImportSwitches(PKSystem system, ICollection>> switches) + { + // Read existing switches to enforce unique timestamps + var priorSwitches = await GetSwitches(system); + var lastSwitchId = priorSwitches.Any() + ? priorSwitches.Max(x => x.Id) + : 0; + + using (var conn = (PerformanceTrackingConnection) await _conn.Obtain()) + { + using (var tx = conn.BeginTransaction()) + { + // Import switches in bulk + using (var importer = conn.BeginBinaryImport("COPY switches (system, timestamp) FROM STDIN (FORMAT BINARY)")) + { + foreach (var sw in switches) + { + // If there's already a switch at this time, move on + if (priorSwitches.Any(x => x.Timestamp.Equals(sw.Item1))) + continue; + + // Otherwise, add it to the importer + importer.StartRow(); + importer.Write(system.Id, NpgsqlTypes.NpgsqlDbType.Integer); + importer.Write(sw.Item1, NpgsqlTypes.NpgsqlDbType.Timestamp); + } + importer.Complete(); // Commits the copy operation so dispose won't roll it back + } + + // Get all switches that were created above and don't have members for ID lookup + var switchesWithoutMembers = + await conn.QueryAsync(@" + SELECT switches.* + FROM switches + LEFT JOIN switch_members + ON switch_members.switch = switches.id + WHERE switches.id > @LastSwitchId + AND switches.system = @System + AND switch_members.id IS NULL", new { LastSwitchId = lastSwitchId, System = system.Id }); + + // Import switch_members in bulk + using (var importer = conn.BeginBinaryImport("COPY switch_members (switch, member) FROM STDIN (FORMAT BINARY)")) + { + // Iterate over the switches we created above and set their members + foreach (var pkSwitch in switchesWithoutMembers) + { + // If this isn't in our import set, move on + var sw = switches.FirstOrDefault(x => x.Item1.Equals(pkSwitch.Timestamp)); + if (sw == null) + continue; + + // Loop through associated members to add each to the switch + foreach (var m in sw.Item2) + { + // Skip switch-outs - these don't have switch_members + if (m == null) + continue; + importer.StartRow(); + importer.Write(pkSwitch.Id, NpgsqlTypes.NpgsqlDbType.Integer); + importer.Write(m.Id, NpgsqlTypes.NpgsqlDbType.Integer); + } + } + importer.Complete(); // Commits the copy operation so dispose won't roll it back + } + tx.Commit(); + } + } + + _logger.Information("Completed bulk import of switches for system {0}", system.Hid); + } + public async Task RegisterSwitches(PKSystem system, ICollection>> switches) { // Use a transaction here since we're doing multiple executed commands in one diff --git a/PluralKit.Core/Utils.cs b/PluralKit.Core/Utils.cs index b71c5b6a..31cd3f1b 100644 --- a/PluralKit.Core/Utils.cs +++ b/PluralKit.Core/Utils.cs @@ -488,6 +488,11 @@ namespace PluralKit _impl.Open(); } + public NpgsqlBinaryImporter BeginBinaryImport(string copyFromCommand) + { + return _impl.BeginBinaryImport(copyFromCommand); + } + public string ConnectionString { get => _impl.ConnectionString; From 593e9d9aef98df1a03093cd54aa2c8dbcd56edfb Mon Sep 17 00:00:00 2001 From: Noko Date: Sun, 20 Oct 2019 18:24:00 -0500 Subject: [PATCH 4/4] Fix edge case of importing exactly enough members to hit the cap It's OK if import will bring us to the cap (1000), just don't let the count go beyond that. --- PluralKit.Core/DataFiles.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/PluralKit.Core/DataFiles.cs b/PluralKit.Core/DataFiles.cs index b6e56945..8a9de1b3 100644 --- a/PluralKit.Core/DataFiles.cs +++ b/PluralKit.Core/DataFiles.cs @@ -121,7 +121,7 @@ namespace PluralKit.Bot // If creating the unmatched members would put us over the member limit, abort before creating any members // new total: # in the system + (# in the file - # in the file that already exist) - if (data.Members.Count - dataFileToMemberMapping.Count + existingMembers.Count() >= Limits.MaxMemberCount) + if (data.Members.Count - dataFileToMemberMapping.Count + existingMembers.Count() > Limits.MaxMemberCount) { result.Success = false; result.Message = $"Import would exceed the maximum number of members ({Limits.MaxMemberCount}).";