Experimental - Cancel-able Menus

commons/src/main/java/com/jagrosh/jdautilities/commons/waiter/EventWaiter.java

@@ -24,9 +24,7 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Set;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -150,16 +148,19 @@ public boolean isShutdown()
      * @param  action
      *         The Consumer to perform an action when the condition Predicate returns {@code true}. Never null.
      *
+     * @return Future for canceling the waiting.
+     *         Using {@code Future#cancel(boolean)} with {@code mayInterruptIfRunning} set to {@code true} is not supported.
+     *
      * @throws IllegalArgumentException
      *         One of two reasons:
      *         <ul>
      *             <li>1) Either the {@code classType}, {@code condition}, or {@code action} was {@code null}.</li>
      *             <li>2) The internal threadpool is shut down, meaning that no more tasks can be submitted.</li>
      *         </ul>
      */
-    public <T extends Event> void waitForEvent(Class<T> classType, Predicate<T> condition, Consumer<T> action)
+    public <T extends Event> Future<Void> waitForEvent(Class<T> classType, Predicate<T> condition, Consumer<T> action)
     {
-        waitForEvent(classType, condition, action, -1, null, null);
+        return waitForEvent(classType, condition, action, -1, null, null);
     }
 
     /**
@@ -191,33 +192,36 @@ public <T extends Event> void waitForEvent(Class<T> classType, Predicate<T> cond
      *         The Runnable to run if the time runs out before a correct Event is thrown, or
      *         {@code null} if there is no action on timeout.
      *
+     * @return Future for canceling the waiting. This will call the timeoutAction if provided.
+     *         Using {@code Future#cancel(boolean)} with {@code mayInterruptIfRunning} set to {@code true} is not supported.
+     *
      * @throws IllegalArgumentException
      *         One of two reasons:
      *         <ul>
      *             <li>1) Either the {@code classType}, {@code condition}, or {@code action} was {@code null}.</li>
      *             <li>2) The internal threadpool is shut down, meaning that no more tasks can be submitted.</li>
      *         </ul>
      */
-    public <T extends Event> void waitForEvent(Class<T> classType, Predicate<T> condition, Consumer<T> action,
+    public <T extends Event> Future<Void> waitForEvent(Class<T> classType, Predicate<T> condition, Consumer<T> action,
                                                long timeout, TimeUnit unit, Runnable timeoutAction)
     {
         Checks.check(!isShutdown(), "Attempted to register a WaitingEvent while the EventWaiter's threadpool was already shut down!");
         Checks.notNull(classType, "The provided class type");
         Checks.notNull(condition, "The provided condition predicate");
         Checks.notNull(action, "The provided action consumer");
 
-        WaitingEvent we = new WaitingEvent<>(condition, action);
-        Set<WaitingEvent> set = waitingEvents.computeIfAbsent(classType, c -> new HashSet<>());
-        set.add(we);
+        WaitingEvent<T> we = new WaitingEvent<>(classType, condition, action, timeoutAction);
+        Future<Void> future = we.getFuture();
 
         if(timeout > 0 && unit != null)
         {
             threadpool.schedule(() ->
             {
-                if(set.remove(we) && timeoutAction != null)
-                    timeoutAction.run();
+                future.cancel(false);
             }, timeout, unit);
         }
+
+        return future;
     }
 
     @Override
@@ -268,26 +272,56 @@ public void shutdown()
 
         threadpool.shutdown();
     }
-    
+
     private class WaitingEvent<T extends Event>
     {
+        final Class<T> classType;
         final Predicate<T> condition;
         final Consumer<T> action;
+        final Runnable cancelAction;
+        final CompletableFuture<Void> future = new WaitingFuture();
 
-        WaitingEvent(Predicate<T> condition, Consumer<T> action)
+        WaitingEvent(Class<T> classType, Predicate<T> condition, Consumer<T> action, Runnable cancelAction)
         {
+            this.classType = classType;
             this.condition = condition;
             this.action = action;
+            this.cancelAction = cancelAction;
+
+
+            Set<WaitingEvent> set = waitingEvents.computeIfAbsent(classType, c -> new HashSet<>());
+            set.add(this);
         }
 
         boolean attempt(T event)
         {
             if(condition.test(event))
             {
                 action.accept(event);
+                future.complete(null);
                 return true;
             }
             return false;
         }
+
+        Future<Void> getFuture()
+        {
+            return future;
+        }
+
+        private class WaitingFuture extends CompletableFuture<Void>
+        {
+            @Override
+            public boolean cancel(boolean mayInterruptIfRunning)
+            {
+                if(mayInterruptIfRunning)
+                    throw new UnsupportedOperationException("EventWaiter#waitForEvent can not be cancelled with mayInterruptIfRunning set to true");
+                if(isDone() || isCancelled())
+                    return isCancelled();
+                if(waitingEvents.get(classType).remove(WaitingEvent.this) && cancelAction != null)
+                    cancelAction.run();
+                return super.cancel(false);
+            }
+        }
     }
 }

menu/src/main/java/com/jagrosh/jdautilities/menu/ButtonMenu.java

@@ -49,18 +49,16 @@ public class ButtonMenu extends Menu
     private final String description;
     private final List<String> choices;
     private final Consumer<ReactionEmote> action;
-    private final Consumer<Message> finalAction;
-
+
     ButtonMenu(EventWaiter waiter, Set<User> users, Set<Role> roles, long timeout, TimeUnit unit,
                Color color, String text, String description, List<String> choices, Consumer<ReactionEmote> action, Consumer<Message> finalAction)
     {
-        super(waiter, users, roles, timeout, unit);
+        super(waiter, users, roles, timeout, unit, finalAction);
         this.color = color;
         this.text = text;
         this.description = description;
         this.choices = choices;
         this.action = action;
-        this.finalAction = finalAction;
     }
 
     /**
@@ -94,6 +92,7 @@ public void display(Message message)
     private void initialize(RestAction<Message> ra)
     {
         ra.queue(m -> {
+            setAttachedMessage(m);
             for(int i=0; i<choices.size(); i++)
             {
                 // Get the emote to display.
@@ -113,7 +112,7 @@ private void initialize(RestAction<Message> ra)
                 {
                     // This is the last reaction added.
                     r.queue(v -> {
-                        waiter.waitForEvent(MessageReactionAddEvent.class, event -> {
+                        setCancelFuture(waiter.waitForEvent(MessageReactionAddEvent.class, event -> {
                             // If the message is not the same as the ButtonMenu
                             // currently being displayed.
                             if(!event.getMessageId().equals(m.getId()))
@@ -139,8 +138,8 @@ private void initialize(RestAction<Message> ra)
 
                             // Preform the specified action with the ReactionEmote
                             action.accept(event.getReaction().getReactionEmote());
-                            finalAction.accept(m);
-                        }, timeout, unit, () -> finalAction.accept(m));
+                            finalizeMenu();
+                        }, timeout, unit, this::finalizeMenu));
                     });
                 }
             }

menu/src/main/java/com/jagrosh/jdautilities/menu/Menu.java

@@ -18,7 +18,9 @@
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 import com.jagrosh.jdautilities.commons.waiter.EventWaiter;
 import net.dv8tion.jda.core.entities.*;
@@ -60,14 +62,19 @@ public abstract class Menu
     protected Set<Role> roles;
     protected final long timeout;
     protected final TimeUnit unit;
-
-    protected Menu(EventWaiter waiter, Set<User> users, Set<Role> roles, long timeout, TimeUnit unit)
+
+    private final Consumer<Message> finalAction;
+    private Future<Void> cancelFuture;
+    private Message attachedMessage;
+
+    protected Menu(EventWaiter waiter, Set<User> users, Set<Role> roles, long timeout, TimeUnit unit, Consumer<Message> finalAction)
     {
         this.waiter = waiter;
         this.users = users;
         this.roles = roles;
         this.timeout = timeout;
         this.unit = unit;
+        this.finalAction = finalAction;
     }
 
     /**
@@ -89,6 +96,95 @@ protected Menu(EventWaiter waiter, Set<User> users, Set<Role> roles, long timeou
      */
     public abstract void display(Message message);
 
+    /**
+     * Returns the message, this menu was attached to.
+     * This is set <b>asynchronously</b> after a call to a {@code display()} or equivalent method.
+     * <br>The attached message will be reset, when the menu times out or is cancelled.
+     *
+     * @return The message, this menu was attached to or {@code null} if not yet attached or already timed out/cancelled.
+     */
+    public Message getAttachedMessage()
+    {
+        return attachedMessage;
+    }
+
+    /**
+     * Cancels the menu. This stops the EventWaiter from waiting for events and behaves exactly the same way as if the menu timed out.
+     * That includes calling the final/cancel action if provided.
+     */
+    public void cancel()
+    {
+        if(cancelFuture == null)
+            throw new IllegalStateException("Menu was not yet displayed or is already cancelled");
+        cancelFuture.cancel(false);
+        cancelFuture = null;
+    }
+
+    /**
+     * Internally used to set the attached message.
+     * Should be used by the corresponding Menu implementation as soon as the message was created/edited.
+     *
+     * @param  attachedMessage
+     *         The message, where this menu was attached to
+     */
+    protected final void setAttachedMessage(Message attachedMessage)
+    {
+        this.attachedMessage = attachedMessage;
+    }
+
+    /**
+     * Internally used to set the cancelFuture used to cancel the menu.
+     * Should be used by the corresponding Menu implementation as soon as {@code EventWaiter#waitForEvent} was used.
+     *
+     * @param  cancelFuture
+     *         The cancel Future returned from {@code EventWaiter#waitForEvent}
+     */
+    protected final void setCancelFuture(Future<Void> cancelFuture)
+    {
+        this.cancelFuture = cancelFuture;
+    }
+
+    /**
+     * Returns the finalAction given via constructor.
+     *
+     * @return The finalAction given via constructor.
+     */
+    protected final Consumer<Message> getFinalAction()
+    {
+        return finalAction;
+    }
+
+    /**
+     * Calls the final action and cleans up the attached message and cancelFuture (sets them to {@code null}).
+     * <br>The actual Menu implementation should should call this method whenever the waiting loop is about to exit
+     * and the final action should be called (e.g. as timeout action for the EventWaiter).
+     *
+     * <p>This method is a shortcut for using {@link #finalizeMenu(boolean) finalizeMenu(true)}
+     *
+     * @see #finalizeMenu(boolean)
+     */
+    protected void finalizeMenu()
+    {
+        finalizeMenu(true);
+    }
+
+    /**
+     * Cleans up the attached message and cancelFuture (sets them to {@code null}) and optionally calls the final action.
+     * <br>The actual Menu implementation should should call this method whenever the waiting loop is about to exit
+     * and {@link #finalizeMenu() finalizeMenu()} is not applicable.
+     *
+     * @see #finalizeMenu()
+     */
+    protected void finalizeMenu(boolean callFinalAction)
+    {
+        Message tmp = this.attachedMessage;
+        this.attachedMessage = null;
+        //also clean up cancelFuture if not already
+        this.cancelFuture = null;
+        if(callFinalAction && finalAction != null)
+            finalAction.accept(tmp);
+    }
+
     /**
      * Checks to see if the provided {@link net.dv8tion.jda.core.entities.User User}
      * is valid to interact with this Menu.<p>