From 6c81821432fb4c7dbeff8624b0a50bc080aa80e2 Mon Sep 17 00:00:00 2001
From: Devoxin <lukel9914@gmail.com>
Date: Thu, 5 Sep 2024 18:46:20 +0100
Subject: [PATCH] Add shutdown hook to close clients properly

---
 .../java/lavalink/server/io/ShutdownHandler.kt   | 16 ++++++++++++++++
 .../main/java/lavalink/server/io/SocketServer.kt |  4 ++++
 2 files changed, 20 insertions(+)
 create mode 100644 LavalinkServer/src/main/java/lavalink/server/io/ShutdownHandler.kt

diff --git a/LavalinkServer/src/main/java/lavalink/server/io/ShutdownHandler.kt b/LavalinkServer/src/main/java/lavalink/server/io/ShutdownHandler.kt
new file mode 100644
index 000000000..a6ebaefb0
--- /dev/null
+++ b/LavalinkServer/src/main/java/lavalink/server/io/ShutdownHandler.kt
@@ -0,0 +1,16 @@
+package lavalink.server.io
+
+import org.springframework.web.socket.CloseStatus
+
+class ShutdownHandler(private val socketServer: SocketServer) : Thread("lavalink-shutdown-handler") {
+    init {
+        isDaemon = false // we want this thread to block shutdown until it has finished running
+    }
+
+    override fun run() {
+        socketServer.contexts.forEach {
+            // don't care about exceptions here, the JVM's shutting down anyway.
+            it.runCatching { closeWebSocket(CloseStatus.GOING_AWAY.code) }
+        }
+    }
+}
diff --git a/LavalinkServer/src/main/java/lavalink/server/io/SocketServer.kt b/LavalinkServer/src/main/java/lavalink/server/io/SocketServer.kt
index 6ffbdfc61..4f2a72743 100644
--- a/LavalinkServer/src/main/java/lavalink/server/io/SocketServer.kt
+++ b/LavalinkServer/src/main/java/lavalink/server/io/SocketServer.kt
@@ -56,6 +56,10 @@ final class SocketServer(
     private val statsCollector = StatsCollector(this)
     private val charPool = ('a'..'z') + ('0'..'9')
 
+    init {
+        Runtime.getRuntime().addShutdownHook(ShutdownHandler(this))
+    }
+
     companion object {
         private val log = LoggerFactory.getLogger(SocketServer::class.java)