From 29d7b5a16471a6a57baa83cbd96589698aeb565a Mon Sep 17 00:00:00 2001
From: tterrag1098 <tterrag1098@gmail.com>
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     |  3 +++
 2 files changed, 16 insertions(+), 6 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<Void> start() {
         GatewayBootstrap<GatewayOptions> 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<Void> 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<Void> 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<Void> 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..664192b1 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;
@@ -58,6 +59,7 @@ public class CommandRegistrar {
 	
 	private boolean finishedDefaultSlurp;
 	private boolean locked;
+	private AtomicBoolean shutdown = new AtomicBoolean(false);
 	
 	private Disposable autoSaveSubscriber;
 	
@@ -250,6 +252,7 @@ private void saveAll() {
     }
 
 	public void onShutdown() {
+		if (shutdown.getAndSet(true)) return;
 	    saveAll();
 		for (ICommand c : commands.values()) {
 		    c.onShutdown();