From a6d5890a4baccebf9197c1162a36e5bcf420e761 Mon Sep 17 00:00:00 2001
From: Gareth Coles <gdude2002@gmail.com>
Date: Fri, 31 May 2024 18:05:51 +0100
Subject: [PATCH] Unload extensions on close, add shutdown hook

---
 .idea/compiler.xml                            |  3 ++
 .../kord/extensions/ExtensibleBot.kt          | 44 ++++++++++++++++---
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/.idea/compiler.xml b/.idea/compiler.xml
index 1f2de3224d..beaca081ca 100644
--- a/.idea/compiler.xml
+++ b/.idea/compiler.xml
@@ -11,6 +11,9 @@
       <module name="kord-extensions.extra-modules.mappings" target="14" />
       <module name="kord-extensions.extra-modules.mappings.main" target="14" />
       <module name="kord-extensions.extra-modules.mappings.test" target="14" />
+      <module name="kord-extensions.web.backend" target="17" />
+      <module name="kord-extensions.web.backend.main" target="13" />
+      <module name="kord-extensions.web.backend.test" target="13" />
     </bytecodeTargetLevel>
   </component>
 </project>
\ No newline at end of file
diff --git a/kord-extensions/src/main/kotlin/com/kotlindiscord/kord/extensions/ExtensibleBot.kt b/kord-extensions/src/main/kotlin/com/kotlindiscord/kord/extensions/ExtensibleBot.kt
index b3cec9b7ea..ab68d3b436 100644
--- a/kord-extensions/src/main/kotlin/com/kotlindiscord/kord/extensions/ExtensibleBot.kt
+++ b/kord-extensions/src/main/kotlin/com/kotlindiscord/kord/extensions/ExtensibleBot.kt
@@ -50,6 +50,9 @@ import kotlinx.serialization.json.Json
 import kotlinx.serialization.json.decodeFromJsonElement
 import org.koin.core.component.inject
 import org.koin.dsl.bind
+import java.lang.IllegalArgumentException
+import java.lang.IllegalStateException
+import kotlin.concurrent.thread
 
 /**
  * An extensible bot, wrapping a Kord instance.
@@ -96,6 +99,13 @@ public open class ExtensibleBot(
 	/** @suppress **/
 	public open val logger: KLogger = KotlinLogging.logger {}
 
+	/** @suppress **/
+	public open val shutdownHook: Thread = thread(false) {
+		runBlocking {
+			close()
+		}
+	}
+
 	/** @suppress Function that sets up the bot early on, called by the builder. **/
 	public open suspend fun setup() {
 		val kord = settings.kordBuilder(token) {
@@ -139,6 +149,15 @@ public open class ExtensibleBot(
 
 		if (!initialized) registerListeners()
 
+		@Suppress("TooGenericExceptionCaught")
+		try {
+			Runtime.getRuntime().addShutdownHook(shutdownHook)
+		} catch (e: IllegalArgumentException) {
+			logger.debug(e) { "Shutdown hook already added or thread is running." }
+		} catch (e: Exception) {
+			logger.warn(e) { "Unable to add shutdown hook." }
+		}
+
 		getKoin().get<Kord>().login {
 			this.presence(settings.presenceBuilder)
 			this.intents = Intents(settings.intentsBuilder!!)
@@ -148,7 +167,7 @@ public open class ExtensibleBot(
 	/**
 	 * Stop the bot by logging out [Kord].
 	 *
-	 * This will leave the Koin context intact, so subsequent restarting of the bot is possible.
+	 * This will leave extensions loaded and the Koin context intact, so later restarting of the bot is possible.
 	 *
 	 * @see close
 	 **/
@@ -157,19 +176,32 @@ public open class ExtensibleBot(
 	}
 
 	/**
-	 * Stop the bot by shutting down [Kord] and removing its Koin context.
+	 * Stop the bot by unloading extensions, shutting down [Kord], and removing its Koin context.
 	 *
-	 * Restarting the bot after closing will result in undefined behavior
+	 * Restarting the bot after closing will result in undefined behaviour
 	 * because the Koin context needed to start will no longer exist.
 	 *
-	 * If a bot has been closed, then it must be fully rebuilt to start again.
-	 *
-	 * If a new bot is going to be built, then the previous bot must be closed first.
+	 * You must fully rebuild closed bots to start again.
+	 * Likewise, you must close previous bots before building a new one.
 	 *
 	 * @see stop
 	 **/
 	public open suspend fun close() {
+		@Suppress("TooGenericExceptionCaught")
+		try {
+			Runtime.getRuntime().removeShutdownHook(shutdownHook)
+		} catch (e: IllegalStateException) {
+			logger.debug { "Shutdown in progress, unable to remove shutdown hook." }
+		} catch (e: Exception) {
+			logger.warn(e) { "Failed to remove shutdown hook." }
+		}
+
+		extensions.keys.forEach {
+			unloadExtension(it)
+		}
+
 		getKoin().get<Kord>().shutdown()
+
 		KordExContext.stopKoin()
 	}