From 305aa3e8d9c6702f89d45f33c05394611a5796dc Mon Sep 17 00:00:00 2001
From: TheMode <themode@outlook.fr>
Date: Wed, 2 Jun 2021 12:12:14 +0200
Subject: [PATCH] Add support for listener expiration count

---
 .../minestom/server/event/EventListener.java  | 35 ++++++++++++++++---
 .../net/minestom/server/event/EventNode.java  |  8 +++--
 src/test/java/demo/PlayerInit.java            |  1 +
 3 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/src/main/java/net/minestom/server/event/EventListener.java b/src/main/java/net/minestom/server/event/EventListener.java
index 03c6ab9e4..b960d788e 100644
--- a/src/main/java/net/minestom/server/event/EventListener.java
+++ b/src/main/java/net/minestom/server/event/EventListener.java
@@ -1,20 +1,23 @@
 package net.minestom.server.event;
 
+import net.minestom.server.utils.time.UpdateOption;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 public class EventListener<T extends Event> {
 
     protected final Class<T> type;
-    protected final Consumer<T> combined;
+    protected final Function<T, Result> executor;
 
-    private EventListener(@NotNull Class<T> type, @NotNull Consumer<T> combined) {
+    private EventListener(@NotNull Class<T> type, @NotNull Function<T, Result> executor) {
         this.type = type;
-        this.combined = combined;
+        this.executor = executor;
     }
 
     public static <T extends Event> EventListener.Builder<T> of(@NotNull Class<T> eventType) {
@@ -26,6 +29,8 @@ public class EventListener<T extends Event> {
         private final Class<T> eventType;
 
         private List<Predicate<T>> filters = new ArrayList<>();
+        private int expirationCount;
+        private UpdateOption expirationTime;
         private Consumer<T> handler;
 
         protected Builder(Class<T> eventType) {
@@ -37,25 +42,47 @@ public class EventListener<T extends Event> {
             return this;
         }
 
+        public EventListener.Builder<T> expirationCount(int expirationCount) {
+            this.expirationCount = expirationCount;
+            return this;
+        }
+
         public EventListener.Builder<T> handler(Consumer<T> handler) {
             this.handler = handler;
             return this;
         }
 
         public EventListener<T> build() {
+            AtomicInteger expirationCount = new AtomicInteger(this.expirationCount);
+            final boolean hasExpirationCount = expirationCount.get() > 0;
+
+            final var filters = new ArrayList<>(this.filters);
+            final var handler = this.handler;
             return new EventListener<>(eventType, event -> {
                 // Filtering
                 if (!filters.isEmpty()) {
                     if (filters.stream().anyMatch(filter -> !filter.test(event))) {
                         // Cancelled
-                        return;
+                        return Result.INVALID;
                     }
                 }
                 // Handler
                 if (handler != null) {
                     handler.accept(event);
                 }
+                // Expiration check
+                if (hasExpirationCount && expirationCount.decrementAndGet() == 0) {
+                    return Result.EXPIRED;
+                }
+                return Result.SUCCESS;
             });
         }
     }
+
+
+    enum Result {
+        SUCCESS,
+        INVALID,
+        EXPIRED
+    }
 }
diff --git a/src/main/java/net/minestom/server/event/EventNode.java b/src/main/java/net/minestom/server/event/EventNode.java
index dbba4c1b1..6726f4fc9 100644
--- a/src/main/java/net/minestom/server/event/EventNode.java
+++ b/src/main/java/net/minestom/server/event/EventNode.java
@@ -51,8 +51,12 @@ public class EventNode<T extends Event> {
         }
         final var listeners = listenerMap.get(event.getClass());
         if (listeners != null && !listeners.isEmpty()) {
-            listeners.forEach(eventListener ->
-                    eventListener.combined.accept(event));
+            listeners.forEach(listener -> {
+                final EventListener.Result result = listener.executor.apply(event);
+                if (result == EventListener.Result.EXPIRED) {
+                    listeners.remove(listener);
+                }
+            });
         }
         this.children.forEach(eventNode -> eventNode.call(event));
     }
diff --git a/src/test/java/demo/PlayerInit.java b/src/test/java/demo/PlayerInit.java
index af355b3d2..0a763acbf 100644
--- a/src/test/java/demo/PlayerInit.java
+++ b/src/test/java/demo/PlayerInit.java
@@ -140,6 +140,7 @@ public class PlayerInit {
         var node = EventNode.create(PlayerEvent.class);
         node.addListener(EventListener.of(PlayerTickEvent.class)
                 .handler(playerTickEvent -> System.out.println("Player tick!"))
+                .expirationCount(2)
                 .build());
 
         var empty = EventNode.create();