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)