diff --git a/src/main/java/net/discordjug/javabot/systems/moderation/PurgeCommand.java b/src/main/java/net/discordjug/javabot/systems/moderation/PurgeCommand.java index f259ca635..5fab74220 100644 --- a/src/main/java/net/discordjug/javabot/systems/moderation/PurgeCommand.java +++ b/src/main/java/net/discordjug/javabot/systems/moderation/PurgeCommand.java @@ -29,10 +29,16 @@ import java.time.LocalDateTime; import java.time.OffsetDateTime; import java.time.format.DateTimeFormatter; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.Set; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Collectors; /** @@ -42,6 +48,8 @@ public class PurgeCommand extends ModerateCommand { private static final Path ARCHIVE_DIR = Path.of("purgeArchives"); private final ExecutorService asyncPool; + + private final Map> currentPurges = Collections.synchronizedMap(new HashMap<>()); /** * The constructor of this class, which sets the corresponding {@link net.dv8tion.jda.api.interactions.commands.build.SlashCommandData}. @@ -51,8 +59,8 @@ public class PurgeCommand extends ModerateCommand { public PurgeCommand(BotConfig botConfig, ExecutorService asyncPool) { super(botConfig); this.asyncPool = asyncPool; - setModerationSlashCommandData(Commands.slash("purge", "Deletes messages from a channel.") - .addOption(OptionType.INTEGER, "amount", "Number of messages to remove.", true) + setModerationSlashCommandData(Commands.slash("purge", "Bulk-deletes messages from a channel. Use /purge 0 to stop all purges.") + .addOption(OptionType.INTEGER, "amount", "Number of messages to remove. Set this to 0 in order to stop all running purges.", true) .addOption(OptionType.USER, "user", "The user whose messages to remove. If left blank, messages from any user are removed.", false) .addOption(OptionType.BOOLEAN, "archive", "Whether the removed messages should be saved in an archive. This defaults to true, if left blank.", false) ); @@ -69,10 +77,34 @@ protected ReplyCallbackAction handleModerationCommand(@NotNull SlashCommandInter Long amount = (amountOption == null) ? 1 : amountOption.getAsLong(); User user = (userOption == null) ? null : userOption.getAsUser(); int maxAmount = config.getPurgeMaxMessageCount(); - if (amount == null || amount < 1 || amount > maxAmount) { + if (amount == null || amount > maxAmount) { return Responses.warning(event, "Invalid amount. Should be between 1 and " + maxAmount + ", inclusive."); } - asyncPool.submit(() -> this.purge(amount, user, event.getUser(), archive, event.getChannel(), config.getLogChannel())); + if (amount == 0) { + Set purges = currentPurges.get(event.getGuild().getIdLong()); + if (purges == null) { + return Responses.warning(event, "Cannot stop purge as no purge is currently running."); + } else { + int count = 0; + for (RunningPurge purge : purges) { + if (purge.cancelled().compareAndSet(false, true)) { + count++; + } + } + return Responses.success(event, "Purge stopped", count + " purge(s) have been stopped."); + } + } + RunningPurge runningPurge = new RunningPurge(event.getIdLong(), new AtomicBoolean()); + CompletableFuture future = CompletableFuture.runAsync( + () -> this.purge(amount, user, event.getUser(), archive, event.getChannel(), config.getLogChannel(), runningPurge.cancelled()), + asyncPool); + currentPurges + .computeIfAbsent(event.getGuild().getIdLong(), l -> Collections.synchronizedSet(new HashSet<>())) + .add(runningPurge); + future.whenComplete((success, failure) -> + currentPurges.get(event.getGuild().getIdLong()) + .remove(runningPurge) + ); StringBuilder sb = new StringBuilder(); sb.append(amount > 1 ? "Up to " + amount + " messages " : "1 message "); if (user != null) { @@ -92,20 +124,42 @@ protected ReplyCallbackAction handleModerationCommand(@NotNull SlashCommandInter * @param archive Whether to create an archive file for the purge. * @param channel The channel to remove messages from. * @param logChannel The channel to write log messages to during the purge. + * @param cancelled {@code true} indicates the purge is cancelled, else {@code false} */ - private void purge(long amount, @Nullable User user, User initiatedBy, boolean archive, MessageChannel channel, TextChannel logChannel) { + private void purge(long amount, @Nullable User user, User initiatedBy, boolean archive, MessageChannel channel, TextChannel logChannel, AtomicBoolean cancelled) { MessageHistory history = channel.getHistory(); String timestamp = LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd_HH-mm-ss")); String file = String.format("purge_%s_%s.txt", channel.getName(), timestamp); PrintWriter archiveWriter = archive ? createArchiveWriter(channel, logChannel, file) : null; - List messages; OffsetDateTime startTime = OffsetDateTime.now(); long count = 0; logChannel.sendMessageFormat("Starting purge of channel %s, initiated by %s", channel.getAsMention(), initiatedBy.getAsMention()) .queue(); + count = performDeletion(amount, user, channel, logChannel, history, archiveWriter, count, cancelled); + if (archiveWriter != null) { + archiveWriter.close(); + } + MessageCreateAction action = logChannel.sendMessage(String.format( + "Purge of channel %s has completed. %d messages have been removed, and the purge took %s.", + channel.getAsMention(), + count, + new TimeUtils().formatDurationToNow(startTime) + )); + if (archive) { + action.addFiles(FileUpload.fromData(ARCHIVE_DIR.resolve(file).toFile())); + } + action.queue(); + } + + private long performDeletion(long amount, User user, MessageChannel channel, TextChannel logChannel, + MessageHistory history, PrintWriter archiveWriter, long count, AtomicBoolean cancelled) { int lastEmptyIterations = 0; + List messages; do { messages = history.retrievePast((int) Math.min(100, user==null ? amount : Math.max(amount, 10))).complete(); + if(cancelled.get()) { + return count; + } if (!messages.isEmpty()) { int messagesRemoved = removeMessages(messages, user, archiveWriter, amount - count); count += messagesRemoved; @@ -114,27 +168,15 @@ private void purge(long amount, @Nullable User user, User initiatedBy, boolean a messagesRemoved, channel.getAsMention(), count - )).queue(); + )).complete(); if (messagesRemoved == 0) { lastEmptyIterations++; }else { lastEmptyIterations = 0; } } - } while (!messages.isEmpty() && amount > count && lastEmptyIterations <= 20); - if (archiveWriter != null) { - archiveWriter.close(); - } - MessageCreateAction action = logChannel.sendMessage(String.format( - "Purge of channel %s has completed. %d messages have been removed, and the purge took %s.", - channel.getAsMention(), - count, - new TimeUtils().formatDurationToNow(startTime) - )); - if (archive) { - action.addFiles(FileUpload.fromData(ARCHIVE_DIR.resolve(file).toFile())); - } - action.queue(); + } while (!cancelled.get() && !messages.isEmpty() && amount > count && lastEmptyIterations <= 20); + return count; } /** @@ -162,7 +204,7 @@ private int removeMessages(List messages, @Nullable User user, @Nullabl for (Message msg : msgs) { archiveMessage(archiveWriter, msg); } - entry.getKey().purgeMessages(messages); + entry.getKey().purgeMessages(msgs); } } return count; @@ -208,4 +250,11 @@ private void archiveMessage(PrintWriter writer, Message message) { message.getContentRaw() ); } + + private record RunningPurge(long id, AtomicBoolean cancelled) { + @Override + public final int hashCode() { + return (int) id; + } + } } \ No newline at end of file