Support custom timeout on tasks (#2439)

src/main/java/net/dv8tion/jda/api/entities/Guild.java

@@ -2508,7 +2508,8 @@ default Task<List<Member>> findMembers(@Nonnull Predicate<? super Member> filter
             if (filter.test(member))
                 list.add(member);
         });
-        GatewayTask<List<Member>> task = new GatewayTask<>(future, reference::cancel);
+        GatewayTask<List<Member>> task = new GatewayTask<>(future, reference::cancel)
+                .onSetTimeout(timeout -> reference.setTimeout(Duration.ofMillis(timeout)));
         reference.onSuccess(it -> future.complete(list))
                  .onError(future::completeExceptionally);
         return task;

src/main/java/net/dv8tion/jda/api/utils/concurrent/Task.java

@@ -16,7 +16,11 @@
 
 package net.dv8tion.jda.api.utils.concurrent;
 
+import net.dv8tion.jda.internal.utils.Checks;
+
 import javax.annotation.Nonnull;
+import java.time.Duration;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
 /**
@@ -68,6 +72,48 @@ public interface Task<T>
     @Nonnull
     Task<T> onSuccess(@Nonnull Consumer<? super T> callback);
 
+    /**
+     * Change the timeout duration for this task.
+     * <br>This may be ignored for certain operations.
+     *
+     * <p>The provided timeout is relative to the start time of the task.
+     * If the time has already passed, this will immediately cancel the task.
+     *
+     * @param  timeout
+     *         The new timeout duration
+     *
+     * @throws IllegalArgumentException
+     *         If null is provided or the timeout is not positive
+     *
+     * @return The current Task instance for chaining
+     */
+    @Nonnull
+    Task<T> setTimeout(@Nonnull Duration timeout);
+
+    /**
+     * Change the timeout duration for this task.
+     * <br>This may be ignored for certain operations.
+     *
+     * <p>The provided timeout is relative to the start time of the task.
+     * If the time has already passed, this will immediately cancel the task.
+     *
+     * @param  timeout
+     *         The new timeout duration
+     * @param  unit
+     *         The time unit of the timeout
+     *
+     * @throws IllegalArgumentException
+     *         If null is provided or the timeout is not positive
+     *
+     * @return The current Task instance for chaining
+     */
+    @Nonnull
+    default Task<T> setTimeout(long timeout, TimeUnit unit)
+    {
+        Checks.notNull(unit, "TimeUnit");
+        return setTimeout(Duration.ofMillis(unit.toMillis(timeout)));
+    }
+
     /**
      * Blocks the current thread until the result is ready.
      * <br>This will not work on the default JDA event thread because it might depend on other events to be processed,

src/main/java/net/dv8tion/jda/internal/entities/GuildImpl.java

@@ -1118,12 +1118,12 @@ public Task<Void> loadMembers(@Nonnull Consumer<Member> callback)
 
         MemberChunkManager chunkManager = getJDA().getClient().getChunkManager();
         boolean includePresences = getJDA().isIntent(GatewayIntent.GUILD_PRESENCES);
-        CompletableFuture<Void> handler = chunkManager.chunkGuild(this, includePresences, (last, list) -> list.forEach(callback));
+        MemberChunkManager.ChunkRequest handler = chunkManager.chunkGuild(this, includePresences, (last, list) -> list.forEach(callback));
         handler.exceptionally(ex -> {
             WebSocketClient.LOG.error("Encountered exception trying to handle member chunk response", ex);
             return null;
         });
-        return new GatewayTask<>(handler, () -> handler.cancel(false));
+        return new GatewayTask<>(handler, () -> handler.cancel(false)).onSetTimeout(handler::setTimeout);
     }
 
     @Nonnull
@@ -1159,7 +1159,7 @@ public Task<List<Member>> retrieveMembersByIds(boolean includePresence, @Nonnull
         MemberChunkManager chunkManager = api.getClient().getChunkManager();
         List<Member> collect = new ArrayList<>(ids.length);
         CompletableFuture<List<Member>> result = new CompletableFuture<>();
-        CompletableFuture<Void> handle = chunkManager.chunkGuild(this, includePresence, ids, (last, list) -> {
+        MemberChunkManager.ChunkRequest handle = chunkManager.chunkGuild(this, includePresence, ids, (last, list) -> {
             collect.addAll(list);
             if (last)
                 result.complete(collect);
@@ -1171,7 +1171,7 @@ public Task<List<Member>> retrieveMembersByIds(boolean includePresence, @Nonnull
             return null;
         });
 
-        return new GatewayTask<>(result, () -> handle.cancel(false));
+        return new GatewayTask<>(result, () -> handle.cancel(false)).onSetTimeout(handle::setTimeout);
     }
 
     @Nonnull
@@ -1186,7 +1186,7 @@ public Task<List<Member>> retrieveMembersByPrefix(@Nonnull String prefix, int li
 
         List<Member> collect = new ArrayList<>(limit);
         CompletableFuture<List<Member>> result = new CompletableFuture<>();
-        CompletableFuture<Void> handle = chunkManager.chunkGuild(this, prefix, limit, (last, list) -> {
+        MemberChunkManager.ChunkRequest handle = chunkManager.chunkGuild(this, prefix, limit, (last, list) -> {
             collect.addAll(list);
             if (last)
                 result.complete(collect);
@@ -1198,7 +1198,7 @@ public Task<List<Member>> retrieveMembersByPrefix(@Nonnull String prefix, int li
             return null;
         });
 
-        return new GatewayTask<>(result, () -> handle.cancel(false));
+        return new GatewayTask<>(result, () -> handle.cancel(false)).onSetTimeout(handle::setTimeout);
     }
 
     @Nonnull

src/main/java/net/dv8tion/jda/internal/requests/MemberChunkManager.java

@@ -70,7 +70,7 @@ public void shutdown()
             timeoutHandle.cancel(false);
     }
 
-    public CompletableFuture<Void> chunkGuild(GuildImpl guild, boolean presence, BiConsumer<Boolean, List<Member>> handler)
+    public ChunkRequest chunkGuild(GuildImpl guild, boolean presence, BiConsumer<Boolean, List<Member>> handler)
     {
         init();
         DataObject request = DataObject.empty()
@@ -84,7 +84,7 @@ public CompletableFuture<Void> chunkGuild(GuildImpl guild, boolean presence, BiC
         return chunkRequest;
     }
 
-    public CompletableFuture<Void> chunkGuild(GuildImpl guild, String query, int limit, BiConsumer<Boolean, List<Member>> handler)
+    public ChunkRequest chunkGuild(GuildImpl guild, String query, int limit, BiConsumer<Boolean, List<Member>> handler)
     {
         init();
         DataObject request = DataObject.empty()
@@ -97,7 +97,7 @@ public CompletableFuture<Void> chunkGuild(GuildImpl guild, String query, int lim
         return chunkRequest;
     }
 
-    public CompletableFuture<Void> chunkGuild(GuildImpl guild, boolean presence, long[] userIds, BiConsumer<Boolean, List<Member>> handler)
+    public ChunkRequest chunkGuild(GuildImpl guild, boolean presence, long[] userIds, BiConsumer<Boolean, List<Member>> handler)
     {
         init();
         DataObject request = DataObject.empty()
@@ -152,13 +152,14 @@ private void sendChunkRequest(DataObject request)
         client.sendChunkRequest(request);
     }
 
-    private class ChunkRequest extends CompletableFuture<Void>
+    public class ChunkRequest extends CompletableFuture<Void>
     {
         private final BiConsumer<Boolean, List<Member>> handler;
         private final GuildImpl guild;
         private final DataObject request;
         private final long nonce;
         private long startTime;
+        private long timeout = MAX_CHUNK_AGE;
 
         public ChunkRequest(BiConsumer<Boolean, List<Member>> handler, GuildImpl guild, DataObject request)
         {
@@ -168,6 +169,12 @@ public ChunkRequest(BiConsumer<Boolean, List<Member>> handler, GuildImpl guild,
             this.request = request.put("nonce", getNonce());
         }
 
+        public ChunkRequest setTimeout(long timeout)
+        {
+            this.timeout = timeout;
+            return this;
+        }
+
         public boolean isNonce(String nonce)
         {
             return this.nonce == Long.parseLong(nonce);
@@ -183,6 +190,11 @@ public long getAge()
             return startTime <= 0 ? 0 : System.currentTimeMillis() - startTime;
         }
 
+        public boolean isExpired()
+        {
+            return getAge() > timeout;
+        }
+
         public DataObject getRequest()
         {
             startTime = System.currentTimeMillis();
@@ -241,7 +253,7 @@ public void run()
             MiscUtil.locked(lock, () ->
             {
                 requests.forEachValue(request -> {
-                    if (request.getAge() > MAX_CHUNK_AGE)
+                    if (request.isExpired())
                         request.completeExceptionally(new TimeoutException());
                     return true;
                 });

src/main/java/net/dv8tion/jda/internal/utils/concurrent/task/GatewayTask.java

@@ -24,21 +24,30 @@
 import org.slf4j.Logger;
 
 import javax.annotation.Nonnull;
+import java.time.Duration;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
+import java.util.function.LongConsumer;
 
 public class GatewayTask<T> implements Task<T>
 {
     private static final Logger log = JDALogger.getLog(Task.class);
     private final Runnable onCancel;
     private final CompletableFuture<T> future;
+    private LongConsumer setTimeout;
 
     public GatewayTask(CompletableFuture<T> future, Runnable onCancel)
     {
         this.future = future;
         this.onCancel = onCancel;
     }
 
+    public GatewayTask<T> onSetTimeout(LongConsumer setTimeout)
+    {
+        this.setTimeout = setTimeout;
+        return this;
+    }
+
     @Override
     public boolean isStarted()
     {
@@ -88,6 +97,18 @@ public Task<T> onSuccess(@Nonnull Consumer<? super T> callback)
         return this;
     }
 
+    @Nonnull
+    @Override
+    public Task<T> setTimeout(@Nonnull Duration timeout)
+    {
+        Checks.notNull(timeout, "Timeout");
+        long millis = timeout.toMillis();
+        Checks.positive(millis, "Timeout");
+        if (this.setTimeout != null)
+            this.setTimeout.accept(millis);
+        return this;
+    }
+
     @Nonnull
     @Override
     public T get()