Merge pull request #493 from danthe1st/fix-purge

fix purges deleting too much, purge cancellation

Author: danthe1st

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<Long, Set<RunningPurge>> 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<RunningPurge> 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<Void> 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<Message> 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<Message> 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<Message> 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;
+		}
+	}
 }