diff --git a/src/main/java/net/dv8tion/jda/api/entities/Guild.java b/src/main/java/net/dv8tion/jda/api/entities/Guild.java index 3caa8aa4ae..e7cfa53ca7 100644 --- a/src/main/java/net/dv8tion/jda/api/entities/Guild.java +++ b/src/main/java/net/dv8tion/jda/api/entities/Guild.java @@ -2508,7 +2508,8 @@ default Task> findMembers(@Nonnull Predicate filter if (filter.test(member)) list.add(member); }); - GatewayTask> task = new GatewayTask<>(future, reference::cancel); + GatewayTask> task = new GatewayTask<>(future, reference::cancel) + .onSetTimeout(timeout -> reference.setTimeout(Duration.ofMillis(timeout))); reference.onSuccess(it -> future.complete(list)) .onError(future::completeExceptionally); return task; diff --git a/src/main/java/net/dv8tion/jda/api/utils/concurrent/Task.java b/src/main/java/net/dv8tion/jda/api/utils/concurrent/Task.java index 87ca1454b1..c5317d71dc 100644 --- a/src/main/java/net/dv8tion/jda/api/utils/concurrent/Task.java +++ b/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 @Nonnull Task onSuccess(@Nonnull Consumer callback); + /** + * Change the timeout duration for this task. + *
This may be ignored for certain operations. + * + *

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 setTimeout(@Nonnull Duration timeout); + + /** + * Change the timeout duration for this task. + *
This may be ignored for certain operations. + * + *

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 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. *
This will not work on the default JDA event thread because it might depend on other events to be processed, diff --git a/src/main/java/net/dv8tion/jda/internal/entities/GuildImpl.java b/src/main/java/net/dv8tion/jda/internal/entities/GuildImpl.java index ef18eea01d..77a17301a1 100644 --- a/src/main/java/net/dv8tion/jda/internal/entities/GuildImpl.java +++ b/src/main/java/net/dv8tion/jda/internal/entities/GuildImpl.java @@ -1118,12 +1118,12 @@ public Task loadMembers(@Nonnull Consumer callback) MemberChunkManager chunkManager = getJDA().getClient().getChunkManager(); boolean includePresences = getJDA().isIntent(GatewayIntent.GUILD_PRESENCES); - CompletableFuture 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> retrieveMembersByIds(boolean includePresence, @Nonnull MemberChunkManager chunkManager = api.getClient().getChunkManager(); List collect = new ArrayList<>(ids.length); CompletableFuture> result = new CompletableFuture<>(); - CompletableFuture 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> 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> retrieveMembersByPrefix(@Nonnull String prefix, int li List collect = new ArrayList<>(limit); CompletableFuture> result = new CompletableFuture<>(); - CompletableFuture 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> 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 diff --git a/src/main/java/net/dv8tion/jda/internal/requests/MemberChunkManager.java b/src/main/java/net/dv8tion/jda/internal/requests/MemberChunkManager.java index 1f7e623fd2..e4fc441e65 100644 --- a/src/main/java/net/dv8tion/jda/internal/requests/MemberChunkManager.java +++ b/src/main/java/net/dv8tion/jda/internal/requests/MemberChunkManager.java @@ -70,7 +70,7 @@ public void shutdown() timeoutHandle.cancel(false); } - public CompletableFuture chunkGuild(GuildImpl guild, boolean presence, BiConsumer> handler) + public ChunkRequest chunkGuild(GuildImpl guild, boolean presence, BiConsumer> handler) { init(); DataObject request = DataObject.empty() @@ -84,7 +84,7 @@ public CompletableFuture chunkGuild(GuildImpl guild, boolean presence, BiC return chunkRequest; } - public CompletableFuture chunkGuild(GuildImpl guild, String query, int limit, BiConsumer> handler) + public ChunkRequest chunkGuild(GuildImpl guild, String query, int limit, BiConsumer> handler) { init(); DataObject request = DataObject.empty() @@ -97,7 +97,7 @@ public CompletableFuture chunkGuild(GuildImpl guild, String query, int lim return chunkRequest; } - public CompletableFuture chunkGuild(GuildImpl guild, boolean presence, long[] userIds, BiConsumer> handler) + public ChunkRequest chunkGuild(GuildImpl guild, boolean presence, long[] userIds, BiConsumer> handler) { init(); DataObject request = DataObject.empty() @@ -152,13 +152,14 @@ private void sendChunkRequest(DataObject request) client.sendChunkRequest(request); } - private class ChunkRequest extends CompletableFuture + public class ChunkRequest extends CompletableFuture { private final BiConsumer> 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> handler, GuildImpl guild, DataObject request) { @@ -168,6 +169,12 @@ public ChunkRequest(BiConsumer> 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; }); diff --git a/src/main/java/net/dv8tion/jda/internal/utils/concurrent/task/GatewayTask.java b/src/main/java/net/dv8tion/jda/internal/utils/concurrent/task/GatewayTask.java index 3256b9ebb5..54b4acb362 100644 --- a/src/main/java/net/dv8tion/jda/internal/utils/concurrent/task/GatewayTask.java +++ b/src/main/java/net/dv8tion/jda/internal/utils/concurrent/task/GatewayTask.java @@ -24,14 +24,17 @@ 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 implements Task { private static final Logger log = JDALogger.getLog(Task.class); private final Runnable onCancel; private final CompletableFuture future; + private LongConsumer setTimeout; public GatewayTask(CompletableFuture future, Runnable onCancel) { @@ -39,6 +42,12 @@ public GatewayTask(CompletableFuture future, Runnable onCancel) this.onCancel = onCancel; } + public GatewayTask onSetTimeout(LongConsumer setTimeout) + { + this.setTimeout = setTimeout; + return this; + } + @Override public boolean isStarted() { @@ -88,6 +97,18 @@ public Task onSuccess(@Nonnull Consumer callback) return this; } + @Nonnull + @Override + public Task 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()