Refactor interactive event handlers

PluralKit.Bot/Bot.cs

@@ -72,16 +72,20 @@ namespace PluralKit.Bot
             {
                 var serviceScope = _services.BeginLifetimeScope();
                 
-                // Find an event handler that can handle the type of event (<T>) we're given
-                var handler = serviceScope.Resolve<IEventHandler<T>>();
-                
                 // Also, find a Sentry enricher for the event type (if one is present), and ask it to put some event data in the Sentry scope
                 var sentryEnricher = serviceScope.ResolveOptional<ISentryEnricher<T>>();
                 sentryEnricher?.Enrich(serviceScope.Resolve<Scope>(), evt);
                 
+                // Find an event handler that can handle the type of event (<T>) we're given
+                var handler = serviceScope.Resolve<IEventHandler<T>>();
+                var queue = serviceScope.ResolveOptional<HandlerQueue<T>>();
                 try
                 {
-                    await handler.Handle(evt);
+                    // Delegate to the queue to see if it wants to handle this event
+                    // the TryHandle call returns true if it's handled the event
+                    // Usually it won't, so just pass it on to the main handler
+                    if (queue == null || !await queue.TryHandle(evt))
+                        await handler.Handle(evt);
                 }
                 catch (Exception exc)
                 {

PluralKit.Bot/CommandSystem/Context.cs

@@ -304,5 +304,7 @@ namespace PluralKit.Bot
                 return null;
             }
         }
+
+        public IComponentContext Services => _provider;
     }
 }

PluralKit.Bot/Commands/SystemEdit.cs

@@ -169,9 +169,9 @@ namespace PluralKit.Bot
         public async Task Delete(Context ctx) {
             ctx.CheckSystem();
 
-            var msg = await ctx.Reply($"{Emojis.Warn} Are you sure you want to delete your system? If so, reply to this message with your system's ID (`{ctx.System.Hid}`).\n**Note: this action is permanent.**");
+            await ctx.Reply($"{Emojis.Warn} Are you sure you want to delete your system? If so, reply to this message with your system's ID (`{ctx.System.Hid}`).\n**Note: this action is permanent.**");
-            var reply = await ctx.AwaitMessage(ctx.Channel, ctx.Author, timeout: TimeSpan.FromMinutes(1));
+            if (!await ctx.ConfirmWithReply(ctx.System.Hid))
-            if (reply.Content != ctx.System.Hid) throw new PKError($"System deletion cancelled. Note that you must reply with your system ID (`{ctx.System.Hid}`) *verbatim*.");
+                throw new PKError($"System deletion cancelled. Note that you must reply with your system ID (`{ctx.System.Hid}`) *verbatim*.");
 
             await _data.DeleteSystem(ctx.System);
             await ctx.Reply($"{Emojis.Success} System deleted.");

PluralKit.Bot/Modules.cs

@@ -54,6 +54,10 @@ namespace PluralKit.Bot
             builder.RegisterType<MessageEdited>().As<IEventHandler<MessageUpdateEventArgs>>();
             builder.RegisterType<ReactionAdded>().As<IEventHandler<MessageReactionAddEventArgs>>();
             
+            // Event handler queue
+            builder.RegisterType<HandlerQueue<MessageCreateEventArgs>>().AsSelf().SingleInstance();
+            builder.RegisterType<HandlerQueue<MessageReactionAddEventArgs>>().AsSelf().SingleInstance();
+            
             // Bot services
             builder.RegisterType<EmbedService>().AsSelf().SingleInstance();
             builder.RegisterType<ProxyService>().AsSelf().SingleInstance();

PluralKit.Bot/Utils/ContextUtils.cs

@@ -1,22 +1,71 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 
+using Autofac;
+
 using DSharpPlus;
 using DSharpPlus.Entities;
 using DSharpPlus.EventArgs;
 using DSharpPlus.Exceptions;
 
+using NodaTime;
+
 using PluralKit.Core;
 
 namespace PluralKit.Bot {
     public static class ContextUtils {
-        public static async Task<bool> PromptYesNo(this Context ctx, DiscordMessage message, DiscordUser user = null, TimeSpan? timeout = null) {
+        public static async Task<bool> PromptYesNo(this Context ctx, DiscordMessage message, DiscordUser user = null, Duration? timeout = null)
+        {
+            var cts = new CancellationTokenSource();
+            if (user == null) user = ctx.Author;
+            if (timeout == null) timeout = Duration.FromMinutes(5);
+            
             // "Fork" the task adding the reactions off so we don't have to wait for them to be finished to start listening for presses
             var _ = message.CreateReactionsBulk(new[] {Emojis.Success, Emojis.Error});
-            var reaction = await ctx.AwaitReaction(message, user ?? ctx.Author, r => r.Emoji.Name == Emojis.Success || r.Emoji.Name == Emojis.Error, timeout ?? TimeSpan.FromMinutes(1));
+            
-            return reaction.Emoji.Name == Emojis.Success;
+            bool ReactionPredicate(MessageReactionAddEventArgs e)
+            {
+                if (e.Channel.Id != message.ChannelId || e.Message.Id != message.Id) return false;
+                if (e.User.Id != user.Id) return false;
+                return true;
+            }
+
+            bool MessagePredicate(MessageCreateEventArgs e)
+            {
+                if (e.Channel.Id != message.ChannelId) return false;
+                if (e.Author.Id != user.Id) return false;
+
+                var strings = new [] {"y", "yes", "n", "no"};
+                foreach (var str in strings)
+                    if (e.Message.Content.Equals(str, StringComparison.InvariantCultureIgnoreCase))
+                        return true;
+                
+                return false;
+            }
+
+            var messageTask = ctx.Services.Resolve<HandlerQueue<MessageCreateEventArgs>>().WaitFor(MessagePredicate, timeout, cts.Token);
+            var reactionTask = ctx.Services.Resolve<HandlerQueue<MessageReactionAddEventArgs>>().WaitFor(ReactionPredicate, timeout, cts.Token);
+            
+            var theTask = await Task.WhenAny(messageTask, reactionTask);
+            cts.Cancel();
+
+            if (theTask == messageTask)
+            {
+                var responseMsg = (await messageTask).Message;
+                var positives = new[] {"y", "yes"};
+                foreach (var p in positives)
+                    if (responseMsg.Content.Equals(p, StringComparison.InvariantCultureIgnoreCase))
+                        return true;
+                return false;
+            }
+
+            if (theTask == reactionTask) 
+                return (await reactionTask).Emoji.Name == Emojis.Success;
+
+            return false;
         }
 
         public static async Task<MessageReactionAddEventArgs> AwaitReaction(this Context ctx, DiscordMessage message, DiscordUser user = null, Func<MessageReactionAddEventArgs, bool> predicate = null, TimeSpan? timeout = null) {
@@ -37,33 +86,15 @@ namespace PluralKit.Bot {
             }
         }
 
-        public static async Task<DiscordMessage> AwaitMessage(this Context ctx, DiscordChannel channel, DiscordUser user = null, Func<DiscordMessage, bool> predicate = null, TimeSpan? timeout = null) {
-            var tcs = new TaskCompletionSource<DiscordMessage>();
-            Task Inner(MessageCreateEventArgs args)
-            {
-                var msg = args.Message;
-                if (channel != msg.Channel) return Task.CompletedTask; // Ignore messages in a different channel
-                if (user != null && user != msg.Author) return Task.CompletedTask; // Ignore messages from other users
-                if (predicate != null && !predicate.Invoke(msg)) return Task.CompletedTask; // Check predicate
-                tcs.SetResult(msg);
-                return Task.CompletedTask;
-            }
-
-            ctx.Shard.MessageCreated += Inner;
-            try
-            {
-                return await (tcs.Task.TimeoutAfter(timeout));
-            }
-            finally
-            {
-                ctx.Shard.MessageCreated -= Inner;
-            }
-        }
-        
         public static async Task<bool> ConfirmWithReply(this Context ctx, string expectedReply)
         {
-            var msg = await ctx.AwaitMessage(ctx.Channel, ctx.Author, timeout: TimeSpan.FromMinutes(1));
+            bool Predicate(MessageCreateEventArgs e) =>
-            return string.Equals(msg.Content, expectedReply, StringComparison.InvariantCultureIgnoreCase);
+                e.Author == ctx.Author && e.Channel.Id == ctx.Channel.Id;
+            
+            var msg = await ctx.Services.Resolve<HandlerQueue<MessageCreateEventArgs>>()
+                .WaitFor(Predicate, Duration.FromMinutes(1));
+            
+            return string.Equals(msg.Message.Content, expectedReply, StringComparison.InvariantCultureIgnoreCase);
         }
 
         public static async Task Paginate<T>(this Context ctx, IAsyncEnumerable<T> items, int totalCount, int itemsPerPage, string title, Func<DiscordEmbedBuilder, IEnumerable<T>, Task> renderer) {

PluralKit.Bot/Utils/HandlerQueue.cs

@@ -0,0 +1,79 @@
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+using NodaTime;
+
+namespace PluralKit.Bot
+{
+    public class HandlerQueue<T>
+    {
+        private readonly List<HandlerEntry> _handlers = new List<HandlerEntry>();
+
+        public HandlerEntry Add(Func<T, Task<bool>> handler)
+        {
+            var entry = new HandlerEntry {Handler = handler};
+            _handlers.Add(entry);
+            return entry;
+        }
+
+        public async Task<T> WaitFor(Func<T, bool> predicate, Duration? timeout = null, CancellationToken ct = default)
+        {
+            var timeoutTask = Task.Delay(timeout?.ToTimeSpan() ?? TimeSpan.FromMilliseconds(-1), ct);
+            var tcs = new TaskCompletionSource<T>();
+
+            Task<bool> Handler(T e)
+            {
+                var matches = predicate(e);
+                if (matches) tcs.SetResult(e);
+                return Task.FromResult(matches);
+            }
+
+            var entry = new HandlerEntry {Handler = Handler};
+            _handlers.Add(entry);
+
+            // Wait for either the event task or the timeout task
+            // If the timeout task finishes first, raise, otherwise pass event through
+            try
+            {
+                var theTask = await Task.WhenAny(timeoutTask, tcs.Task);
+                if (theTask == timeoutTask)
+                    throw new TimeoutException();
+            }
+            finally
+            {
+                entry.Remove();
+            }
+
+            return await tcs.Task;
+        }
+
+        public async Task<bool> TryHandle(T evt)
+        {
+            _handlers.RemoveAll(he => !he.Alive);
+
+            var now = SystemClock.Instance.GetCurrentInstant();
+            foreach (var entry in _handlers)
+            {
+                if (entry.Expiry < now) entry.Alive = false;
+                else if (entry.Alive && await entry.Handler(evt))
+                {
+                    entry.Alive = false;
+                    return true;
+                }
+            }
+
+            return false;
+        }
+
+        public class HandlerEntry
+        {
+            internal Func<T, Task<bool>> Handler;
+            internal bool Alive = true;
+            internal Instant Expiry = SystemClock.Instance.GetCurrentInstant() + Duration.FromMinutes(30);
+
+            public void Remove() => Alive = false;
+        }
+    }
+}