From 634541e2dc31f15b7ccaa7a6df149aab56a8d24e Mon Sep 17 00:00:00 2001
From: devoxin <15076404+devoxin@users.noreply.github.com>
Date: Fri, 6 Sep 2024 23:56:24 +0100
Subject: [PATCH] Add shutdown hook to close clients properly (#1102)

---
 .../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)