Fix FlatMapRestAction predicate with complete or submit (#2636)

discord-jda · Mar 30, 2024 · e7bafe9 · e7bafe9
1 parent ce6e2e2
commit e7bafe9
Showing 1 changed file with 19 additions and 2 deletions.
diff --git a/src/main/java/net/dv8tion/jda/internal/requests/restaction/operator/FlatMapRestAction.java b/src/main/java/net/dv8tion/jda/internal/requests/restaction/operator/FlatMapRestAction.java
@@ -21,6 +21,7 @@
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -61,14 +62,30 @@ public void queue(@Nullable Consumer<? super O> success, @Nullable Consumer<? su
     @Override
     public O complete(boolean shouldQueue) throws RateLimitedException
     {
-        return supply(action.complete(shouldQueue)).complete(shouldQueue);
+        I complete = action.complete(shouldQueue);
+
+        if (this.condition != null && !this.condition.test(complete))
+        {
+            throw new CancellationException("FlatMap condition failed");
+        }
+        return supply(complete).complete(shouldQueue);
     }
 
     @Nonnull
     @Override
     public CompletableFuture<O> submit(boolean shouldQueue)
     {
         return action.submit(shouldQueue)
-                .thenCompose((result) -> supply(result).submit(shouldQueue));
+                .thenCompose((result) ->
+                {
+                    if (condition != null && !condition.test(result))
+                    {
+                        CompletableFuture<O> future = new CompletableFuture<>();
+                        future.cancel(true);
+
+                        return future;
+                    }
+                    return supply(result).submit(shouldQueue);
+                });
     }
 }