From ff8b88fc96976d61f1c6873c37b80bb8f709bb9f Mon Sep 17 00:00:00 2001
From: kikkia <10608318+kikkia@users.noreply.github.com>
Date: Fri, 7 Mar 2025 18:48:54 +0900
Subject: [PATCH] Add functionality to nodes to transfer all players to other
 nodes.

---
 .../arbjerg/lavalink/client/LavalinkClient.kt |  2 +-
 .../arbjerg/lavalink/client/LavalinkNode.kt   | 26 ++++++++++++++++++-
 .../lavalink/internal/LavalinkSocket.kt       |  1 +
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkClient.kt b/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkClient.kt
index c59df05..c4ad724 100644
--- a/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkClient.kt
+++ b/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkClient.kt
@@ -196,7 +196,7 @@ class LavalinkClient(val userId: Long) : Closeable, Disposable {
         transferNodes(node)
     }
 
-    private fun transferNodes(node: LavalinkNode) {
+    internal fun transferNodes(node: LavalinkNode) {
         linkMap.forEach { (_, link) ->
             if (link.node == node) {
                 val voiceRegion = link.cachedPlayer?.voiceRegion
diff --git a/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkNode.kt b/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkNode.kt
index 49e0617..b9c5ef7 100644
--- a/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkNode.kt
+++ b/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkNode.kt
@@ -70,6 +70,10 @@ class LavalinkNode(
     var available: Boolean = false
         internal set
 
+    // TODO: Maybe it would make sense to add some node state enums like the links have.
+    var transferring: Boolean = false
+        internal set
+
     /**
      * A local player cache, allows us to not call the rest client every time we need a player.
      */
@@ -183,7 +187,10 @@ class LavalinkNode(
         return rest.destroyPlayer(guildId)
             .doOnSuccess {
                 removeCachedPlayer(guildId)
-                lavalink.removeDestroyedLink(guildId)
+                // When transferring to a new node preserve the link
+                if (!transferring) {
+                    lavalink.removeDestroyedLink(guildId)
+                }
             }
     }
 
@@ -473,6 +480,23 @@ class LavalinkNode(
         }
     }
 
+    /**
+     * Puts the node into a transferring state, then starts transferring all players to other nodes.
+     * Once the node is put into this state, a reconnect or calling transferComplete() removes it from that state.
+     * Transferring state does not remove links from cache when deleting off this node, unless the transfer also fails.
+     */
+    fun transferPlayersToOtherNodes() {
+        transferring = true
+        lavalink.transferNodes(this)
+    }
+
+    /**
+     * Manually removes this node from the transfer state.
+     */
+    fun transferComplete() {
+        transferring = false
+    }
+
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
         if (javaClass != other?.javaClass) return false
diff --git a/src/main/kotlin/dev/arbjerg/lavalink/internal/LavalinkSocket.kt b/src/main/kotlin/dev/arbjerg/lavalink/internal/LavalinkSocket.kt
index a0e1e49..5815adb 100644
--- a/src/main/kotlin/dev/arbjerg/lavalink/internal/LavalinkSocket.kt
+++ b/src/main/kotlin/dev/arbjerg/lavalink/internal/LavalinkSocket.kt
@@ -66,6 +66,7 @@ class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Clos
 
                 node.sessionId = sessionId
                 node.available = true
+                node.transferring = false
                 logger.info("${node.name} is ready with session id $sessionId")
 
                 node.playerCache.values.forEach { player ->