From 02c4e912efd9947e875684db5c27e95e082cad44 Mon Sep 17 00:00:00 2001 From: tterrag1098 Date: Thu, 29 Apr 2021 15:18:04 -0400 Subject: [PATCH] Fix startup procedure: Fixes potential ghost death and console handler --- src/main/java/com/tterrag/k9/K9.java | 19 +- .../k9/commands/api/CommandRegistrar.java | 191 +++++++++--------- 2 files changed, 110 insertions(+), 100 deletions(-) diff --git a/src/main/java/com/tterrag/k9/K9.java b/src/main/java/com/tterrag/k9/K9.java index 581283a2..eb6be861 100644 --- a/src/main/java/com/tterrag/k9/K9.java +++ b/src/main/java/com/tterrag/k9/K9.java @@ -130,8 +130,8 @@ public Mono start() { GatewayBootstrap gateway = client.gateway() .setEventDispatcher(ReplayingEventDispatcher.builder() .replayEventFilter(e -> e instanceof ReadyEvent) - .eventScheduler(Schedulers.boundedElastic()) - .build()) + .eventScheduler(Schedulers.boundedElastic()) + .build()) .setEnabledIntents(IntentSet.of( Intent.GUILDS, Intent.GUILD_MEMBERS, Intent.GUILD_PRESENCES, Intent.GUILD_MESSAGES, Intent.GUILD_MESSAGE_REACTIONS, @@ -193,8 +193,15 @@ public Mono start() { return Mono.fromRunnable(commands::slurpCommands) .then(gateway.login()) - .flatMap(c -> Mono.when(onInitialReady.apply(c.getEventDispatcher()), services.start(c)).thenReturn(c)) - .flatMap(this::teardown); + .flatMap(c -> + Mono.when(onInitialReady.apply(c.getEventDispatcher()), services.start(c)) + .doOnError(t -> log.error("Unexpected error received in main bot subscriber:", t)) + .doOnTerminate(() -> log.error("Unexpected completion of main bot subscriber!")) + .zipWith(teardown(c)) + .thenReturn(c) + .onErrorResume($ -> teardown(c).thenReturn(c))) + .flatMap(c -> Mono.fromRunnable(commands::onShutdown) + .then(c.logout())); } private boolean isUser(MessageCreateEvent evt) { @@ -210,7 +217,7 @@ private Mono teardown(GatewayDiscordClient gatewayClient) { while (scan.hasNextLine()) { if (scan.nextLine().equals("stop")) { scan.close(); - System.exit(0); + return null; // Empty completion will bubble up to zip below } } Threads.sleep(100); @@ -226,7 +233,7 @@ private Mono teardown(GatewayDiscordClient gatewayClient) { return Mono.zip(consoleHandler, gatewayClient.onDisconnect()) .then() - .doOnTerminate(() -> log.error("Unexpected completion of main bot subscriber!")); + .doOnError(t -> log.error("Disconnect listener error:", t)); } public static String getVersion() { diff --git a/src/main/java/com/tterrag/k9/commands/api/CommandRegistrar.java b/src/main/java/com/tterrag/k9/commands/api/CommandRegistrar.java index 5ea138ec..638480a6 100644 --- a/src/main/java/com/tterrag/k9/commands/api/CommandRegistrar.java +++ b/src/main/java/com/tterrag/k9/commands/api/CommandRegistrar.java @@ -8,6 +8,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; @@ -41,59 +42,60 @@ @Slf4j public class CommandRegistrar { - + @NonNull - static final File DATA_FOLDER = NullHelper.notnullJ(Paths.get("command_data").toFile(), "Path#toFile"); - static { - DATA_FOLDER.mkdirs(); - } - - private final K9 k9; - - private final Map commands = Maps.newTreeMap(); - private final CommandControl ctrl = new CommandControl(); - - private final @NonNull GsonBuilder builder = new GsonBuilder(); - private @NonNull Gson gson = new Gson(); - - private boolean finishedDefaultSlurp; - private boolean locked; - - private Disposable autoSaveSubscriber; - - public CommandRegistrar(K9 k9) { - this.k9 = k9; - } + static final File DATA_FOLDER = NullHelper.notnullJ(Paths.get("command_data").toFile(), "Path#toFile"); + static { + DATA_FOLDER.mkdirs(); + } + + private final K9 k9; + + private final Map commands = Maps.newTreeMap(); + private final CommandControl ctrl = new CommandControl(); + + private final @NonNull GsonBuilder builder = new GsonBuilder(); + private @NonNull Gson gson = new Gson(); + + private boolean finishedDefaultSlurp; + private boolean locked; + private AtomicBoolean shutdown = new AtomicBoolean(false); + + private Disposable autoSaveSubscriber; + + public CommandRegistrar(K9 k9) { + this.k9 = k9; + } - public Mono invokeCommand(MessageCreateEvent evt, String name, String argstr) { - Optional commandReq = findCommand(evt.getGuildId().orElse(null), name); - - ICommand command = commandReq.filter(c -> !c.admin() || evt.getMessage().getAuthor().map(this::isAdmin).orElse(false)).orElse(null); - if (command == null) { - return Mono.empty(); - } - - CommandContext ctx = new CommandContext(k9, evt); - - if (!command.requirements().matches(ctx).block()) { - return evt.getMessage().getChannel() - .flatMap(c -> c.createMessage("You do not have permission to use this command!")) - .delayElement(Duration.ofSeconds(5)) - .flatMap(m -> m.delete()) - .thenReturn(command); - } - -// evt.getMessage().getChannel().flatMap(c -> c.type()).subscribe(); - - argstr = Strings.nullToEmpty(argstr); - - Map flags = new HashMap<>(); - Map, String> args = new HashMap<>(); - - Map keyToFlag = command.getFlags().stream().collect(Collectors.toMap(Flag::name, f -> f)); - Map longKeyToFlag = command.getFlags().stream().collect(Collectors.toMap(Flag::longFormName, f -> f)); + public Mono invokeCommand(MessageCreateEvent evt, String name, String argstr) { + Optional commandReq = findCommand(evt.getGuildId().orElse(null), name); + + ICommand command = commandReq.filter(c -> !c.admin() || evt.getMessage().getAuthor().map(this::isAdmin).orElse(false)).orElse(null); + if (command == null) { + return Mono.empty(); + } + + CommandContext ctx = new CommandContext(k9, evt); + + if (!command.requirements().matches(ctx).block()) { + return evt.getMessage().getChannel() + .flatMap(c -> c.createMessage("You do not have permission to use this command!")) + .delayElement(Duration.ofSeconds(5)) + .flatMap(m -> m.delete()) + .thenReturn(command); + } + +// evt.getMessage().getChannel().flatMap(c -> c.type()).subscribe(); + + argstr = Strings.nullToEmpty(argstr); + + Map flags = new HashMap<>(); + Map, String> args = new HashMap<>(); + + Map keyToFlag = command.getFlags().stream().collect(Collectors.toMap(Flag::name, f -> f)); + Map longKeyToFlag = command.getFlags().stream().collect(Collectors.toMap(Flag::longFormName, f -> f)); - Matcher matcher = Patterns.FLAGS.matcher(argstr); + Matcher matcher = Patterns.FLAGS.matcher(argstr); while (matcher.find()) { String flagname = matcher.group(2); List foundFlags; @@ -165,14 +167,14 @@ public Mono invokeCommand(MessageCreateEvent evt, String name, String return ctx.reply("Unexpected error processing command: " + e).thenReturn(command); // TODO should this be different? } } - - public boolean isAdmin(User user) { - return k9.isAdmin(user.getId()); - } - - public Optional findCommand(CommandContext ctx, String name) { - return findCommand(ctx.getGuildId().orElse(null), name); - } + + public boolean isAdmin(User user) { + return k9.isAdmin(user.getId()); + } + + public Optional findCommand(CommandContext ctx, String name) { + return findCommand(ctx.getGuildId().orElse(null), name); + } public Optional findCommand(@Nullable Snowflake guild, String name) { if (guild != null && ctrl.getData(guild).getCommandBlacklist().contains(name)) { @@ -188,7 +190,7 @@ public void slurpCommands() { } } - @SneakyThrows + @SneakyThrows public void slurpCommands(@NonNull String packagename) { if (locked) { throw new IllegalStateException("Cannot slurp commands in locked registrar."); @@ -199,28 +201,28 @@ public void slurpCommands(@NonNull String packagename) { return; // ?? } ClassPath classpath = ClassPath.from(loader); - for (ClassInfo foo : classpath.getTopLevelClassesRecursive(packagename)) { - if (!foo.getName().equals(getClass().getName())) { - Class c = foo.load(); - if (c.isAnnotationPresent(Command.class)) { - log.info("Found annotation command: {}", c.getName()); - registerCommand((ICommand) c.newInstance()); - } - } - } - } - - public void registerCommand(ICommand command) { - if (locked) { - throw new IllegalStateException("Cannot register command to locked registrar."); - } - if (!command.isTransient()) { - commands.put(command.getName(), command); - command.gatherParsers(builder); - command.onRegister(k9); - } - command.getChildren().forEach(this::registerCommand); - } + for (ClassInfo foo : classpath.getTopLevelClassesRecursive(packagename)) { + if (!foo.getName().equals(getClass().getName())) { + Class c = foo.load(); + if (c.isAnnotationPresent(Command.class)) { + log.info("Found annotation command: {}", c.getName()); + registerCommand((ICommand) c.newInstance()); + } + } + } + } + + public void registerCommand(ICommand command) { + if (locked) { + throw new IllegalStateException("Cannot register command to locked registrar."); + } + if (!command.isTransient()) { + commands.put(command.getName(), command); + command.gatherParsers(builder); + command.onRegister(k9); + } + command.getChildren().forEach(this::registerCommand); + } public void unregisterCommand(ICommand command) { commands.remove(command.getName()); @@ -249,19 +251,20 @@ private void saveAll() { } } - public void onShutdown() { - saveAll(); - for (ICommand c : commands.values()) { - c.onShutdown(); - } - if (autoSaveSubscriber != null) { - autoSaveSubscriber.dispose(); - } - } - - public Iterable getCommands(Optional guild) { - return getCommands(guild.orElse(null)); - } + public void onShutdown() { + if (shutdown.getAndSet(true)) return; + saveAll(); + for (ICommand c : commands.values()) { + c.onShutdown(); + } + if (autoSaveSubscriber != null) { + autoSaveSubscriber.dispose(); + } + } + + public Iterable getCommands(Optional guild) { + return getCommands(guild.orElse(null)); + } public Iterable getCommands(@Nullable Snowflake guild) { if (guild == null) {