Rework resuming (#32)

* Add tracking of Session state on update/connect

* Add delay to moving orphaned Links

* Implement post-resume synchronisation

* Add ResumeSynchronizationEvent

NB: This is potentially breaking as it replaces a `sealed` attribute

* Get session via PATCH instead of GET

* Add LavalinkNode.getCachedPlayers()

* Abandon resuming after failed reconnection attempt

* Handle reconnect even on abnormal closure

* Change how resume reconnect failure is handled

* Defensively call onResumeReconnectFailed(node)

* Fix removing cached player from new node instead of old

Also manages the LinkState a bit more consistently

* Create new Link instances upon resume synchronization

Only if one doesn't already exist for the guild

Author: freyacodes

src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkClient.kt

@@ -1,9 +1,9 @@
 package dev.arbjerg.lavalink.client
 
+import dev.arbjerg.lavalink.client.event.ClientEvent
 import dev.arbjerg.lavalink.client.loadbalancing.ILoadBalancer
 import dev.arbjerg.lavalink.client.loadbalancing.VoiceRegion
 import dev.arbjerg.lavalink.client.loadbalancing.builtin.DefaultLoadBalancer
-import dev.arbjerg.lavalink.client.event.ClientEvent
 import dev.arbjerg.lavalink.client.player.LavalinkPlayer
 import dev.arbjerg.lavalink.internal.ReconnectTask
 import dev.arbjerg.lavalink.protocol.v4.VoiceState
@@ -110,6 +110,19 @@ class LavalinkClient(val userId: Long) : Closeable, Disposable {
             Link(guildId, loadBalancer.selectNode(region, guildId))
         }
     }
+    /**
+     * Get or crate a [Link] between a guild and a node.
+     *
+     * The requested [LavalinkNode] is only assigned if a new [Link] is created
+     *
+     * @param guildId The id of the guild
+     * @param node the node to initially assign the [Link] to if a new one is created
+     */
+    internal fun getOrCreateLink(guildId: Long, node: LavalinkNode): Link {
+        return linkMap.getOrPut(guildId) {
+            Link(guildId, node)
+        }
+    }
 
     /**
      * Returns a [Link] if it exists in the cache.
@@ -169,6 +182,21 @@ class LavalinkClient(val userId: Long) : Closeable, Disposable {
             return
         }
 
+        val session = node.cachedSession
+        val canResume = session != null && session.resuming && session.timeoutSeconds > 0
+        if (canResume) {
+            // This causes onResumeReconnectFailed(node) to be called if the next reconnect fails
+            node.ws.onResumableConnectionDisconnected()
+        } else {
+            transferNodes(node)
+        }
+    }
+
+    internal fun onResumeReconnectFailed(node: LavalinkNode) {
+        transferNodes(node)
+    }
+
+    private fun transferNodes(node: LavalinkNode) {
         linkMap.forEach { (_, link) ->
             if (link.node == node) {
                 val voiceRegion = link.cachedPlayer?.voiceRegion

src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkNode.kt

@@ -1,6 +1,7 @@
 package dev.arbjerg.lavalink.client
 
 import dev.arbjerg.lavalink.client.event.ClientEvent
+import dev.arbjerg.lavalink.client.event.ResumeSynchronizationEvent
 import dev.arbjerg.lavalink.client.http.HttpBuilder
 import dev.arbjerg.lavalink.client.player.*
 import dev.arbjerg.lavalink.client.player.Track
@@ -17,12 +18,14 @@ import kotlinx.serialization.serializer
 import okhttp3.Call
 import okhttp3.OkHttpClient
 import okhttp3.Response
+import org.slf4j.LoggerFactory
 import reactor.core.Disposable
 import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
 import reactor.core.publisher.Sinks
 import reactor.core.publisher.Sinks.Many
 import reactor.kotlin.core.publisher.toMono
+import reactor.util.retry.Retry
 import java.io.Closeable
 import java.io.IOException
 import java.time.Duration
@@ -39,12 +42,14 @@ class LavalinkNode(
     private val nodeOptions: NodeOptions,
     val lavalink: LavalinkClient
 ) : Disposable, Closeable {
+    private val logger = LoggerFactory.getLogger(LavalinkNode::class.java)
     // "safe" uri with all paths removed
     val baseUri = "${nodeOptions.serverUri.scheme}://${nodeOptions.serverUri.host}:${nodeOptions.serverUri.port}"
 
     val name = nodeOptions.name
     val regionFilter = nodeOptions.regionFilter
     val password = nodeOptions.password
+    internal var cachedSession: Session? = null
 
     var sessionId: String? = nodeOptions.sessionId
         internal set
@@ -237,10 +242,12 @@ class LavalinkNode(
     }
 
     /**
-     * Enables resuming. This causes Lavalink to continue playing for [duration], during which
+     * Enables resuming. This causes Lavalink to continue playing for [timeout] amount of time, during which
      *  we can reconnect without losing our session data. */
     fun enableResuming(timeout: Duration): Mono<Session> {
-        return rest.patchSession(Session(resuming = true, timeout.seconds))
+        return rest.patchSession(Session(resuming = true, timeout.seconds)).doOnSuccess {
+            cachedSession = it
+        }
     }
 
     /**
@@ -249,7 +256,9 @@ class LavalinkNode(
      * This is the default behavior, reversing calls to [enableResuming].
      */
     fun disableResuming(): Mono<Session> {
-        return rest.patchSession(Session(resuming = false, timeoutSeconds = 0))
+        return rest.patchSession(Session(resuming = false, timeoutSeconds = 0)).doOnSuccess {
+            cachedSession = it
+        }
     }
 
     /**
@@ -429,6 +438,41 @@ class LavalinkNode(
         lavalink.transferOrphansTo(this)
     }
 
+    internal fun synchronizeAfterResume() {
+        getPlayers()
+            .retryWhen(Retry.fixedDelay(3, Duration.ofSeconds(1)))
+            .map { players ->
+            val remoteGuildIds = players.map { it.guildId }
+
+            players.forEach { player ->
+                playerCache[player.guildId] = player
+
+                val link = lavalink.getOrCreateLink(player.guildId, node = this)
+                if (link.node != this) return@forEach
+
+                link.state = if (player.state.connected) {
+                    LinkState.CONNECTED
+                } else {
+                    LinkState.DISCONNECTED
+                }
+            }
+
+            val missingIds = playerCache.keys().toList() - remoteGuildIds.toSet()
+            missingIds.forEach { guildId ->
+                playerCache.remove(guildId)
+                val link = lavalink.getLinkIfCached(guildId) ?: return@forEach
+                if (link.node == this) link.state = LinkState.DISCONNECTED
+            }
+
+            ResumeSynchronizationEvent(this, failureReason = null)
+        }.doOnError {
+            logger.error("Failure while attempting synchronization with $this", it)
+            sink.tryEmitNext(ResumeSynchronizationEvent(this, failureReason = it))
+        }.subscribe {
+            sink.tryEmitNext(it)
+        }
+    }
+
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
         if (javaClass != other?.javaClass) return false

src/main/kotlin/dev/arbjerg/lavalink/client/event/events.kt

@@ -7,7 +7,7 @@ import dev.arbjerg.lavalink.client.player.toCustom
 import dev.arbjerg.lavalink.protocol.v4.*
 import dev.arbjerg.lavalink.protocol.v4.Message.EmittedEvent.TrackEndEvent.AudioTrackEndReason
 
-internal fun Message.toClientEvent(node: LavalinkNode) = when (this) {
+internal fun Message.toClientEvent(node: LavalinkNode): ClientEvent = when (this) {
     is Message.ReadyEvent -> ReadyEvent(node, resumed, sessionId)
     is Message.EmittedEvent.TrackEndEvent -> TrackEndEvent(node, guildId.toLong(), track.toCustom(), reason)
     is Message.EmittedEvent.TrackExceptionEvent -> TrackExceptionEvent(node, guildId.toLong(), track.toCustom(), exception.toCustom())
@@ -18,12 +18,19 @@ internal fun Message.toClientEvent(node: LavalinkNode) = when (this) {
     is Message.StatsEvent -> StatsEvent(node, frameStats, players, playingPlayers, uptime, memory, cpu)
 }
 
-sealed class ClientEvent(open val node: LavalinkNode)
+abstract class ClientEvent(open val node: LavalinkNode)
 
 // Normal events
 data class ReadyEvent(override val node: LavalinkNode, val resumed: Boolean, val sessionId: String)
     : ClientEvent(node)
 
+/**
+ * Represents a successful or failed synchronization after a [ReadyEvent] with [ReadyEvent.resumed] set to true.
+ *
+ * Whether it is successful depends on whether [failureReason] is null.
+ */
+data class ResumeSynchronizationEvent(override val node: LavalinkNode, val failureReason: Throwable?) : ClientEvent(node)
+
 data class PlayerUpdateEvent(override val node: LavalinkNode, val guildId: Long, val state: PlayerState)
     : ClientEvent(node)
 

src/main/kotlin/dev/arbjerg/lavalink/internal/LavalinkRestClient.kt

@@ -76,6 +76,15 @@ class LavalinkRestClient(val node: LavalinkNode) {
         }.toMono()
     }
 
+    fun getSession(): Mono<Session> {
+        return newRequest {
+            path("/v4/sessions/${node.sessionId}")
+            // Using patch with an empty object is a dirty hack because GET is not supported for this resource
+            // 7 years younger me should have known better ~Freya
+            patch("{}".toRequestBody("application/json".toMediaType()))
+        }.toMono()
+    }
+
     /**
      * Make a request to the lavalink node. This is internal to keep it looking nice in kotlin. Java compatibility is in the node class.
      */

src/main/kotlin/dev/arbjerg/lavalink/internal/LavalinkSocket.kt

@@ -18,6 +18,7 @@ import java.io.EOFException
 import java.net.ConnectException
 import java.net.SocketException
 import java.net.SocketTimeoutException
+import java.util.concurrent.atomic.AtomicBoolean
 
 class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Closeable {
     private val logger = LoggerFactory.getLogger(LavalinkSocket::class.java)
@@ -26,11 +27,15 @@ class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Clos
 
     var mayReconnect = true
     var lastReconnectAttempt = 0L
+    @Volatile
     private var reconnectsAttempted = 0
     val reconnectInterval: Int
         get() = reconnectsAttempted * 2000 - 2000
     var open: Boolean = false
         private set
+    @Volatile
+    private var hasEverConnected = false
+    private val isAttemptingResume = AtomicBoolean(node.sessionId != null)
 
     init {
         connect()
@@ -40,6 +45,8 @@ class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Clos
         logger.info("${node.name} has been connected!")
         open = true
         reconnectsAttempted = 0
+        hasEverConnected = true
+        isAttemptingResume.set(false)
     }
 
     override fun onMessage(webSocket: WebSocket, text: String) {
@@ -75,6 +82,17 @@ class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Clos
                         .subscribe()
                 }
 
+                if (!resumed) {
+                    node.cachedSession = null
+                }
+                if (node.cachedSession == null) {
+                    node.rest.getSession().subscribe { node.cachedSession = it }
+                }
+
+                if (resumed) {
+                    node.synchronizeAfterResume()
+                }
+
                 // Move players from older, unavailable nodes to ourselves.
                 node.transferOrphansToSelf()
             }
@@ -138,21 +156,18 @@ class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Clos
         if (mayReconnect) {
             logger.info("${node.name} disconnected, reconnecting in ${reconnectInterval / 1000} seconds")
         }
-
-        node.available = false
-        open = false
     }
 
     override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
-        handleFailureTrhowable(t)
+        handleFailureThrowable(t)
 
         node.available = false
         open = false
 
         node.lavalink.onNodeDisconnected(node)
     }
 
-    private fun handleFailureTrhowable(t: Throwable) {
+    private fun handleFailureThrowable(t: Throwable) {
         when(t) {
             is EOFException -> {
                 logger.debug("Got disconnected from ${node.name}, trying to reconnect", t)
@@ -180,10 +195,19 @@ class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Clos
                 logger.error("Unknown error on ${node.name}", t)
             }
         }
+
+        if (hasEverConnected && isAttemptingResume.getAndSet(false)) {
+            try {
+                node.lavalink.onResumeReconnectFailed(node)
+            } catch (e: Exception) {
+                logger.error("Exception after giving up on resuming", e)
+            }
+        }
     }
 
     override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
         node.available = false
+        open = false
         node.lavalink.onNodeDisconnected(node)
 
         if (code == 1000) {
@@ -203,7 +227,6 @@ class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Clos
                 reason
             )
         }
-
     }
 
     fun attemptReconnect() {
@@ -224,7 +247,7 @@ class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Clos
             .addHeader("Client-Name", "Lavalink-Client/${CLIENT_VERSION}")
             .addHeader("User-Id", node.lavalink.userId.toString())
             .apply {
-                if (node.sessionId != null) {
+                if (node.sessionId != null && isAttemptingResume.get()) {
                     addHeader("Session-Id", node.sessionId!!)
                 }
             }
@@ -241,4 +264,8 @@ class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Clos
         socket?.close(1000, "Client shutdown")
         socket?.cancel()
     }
+
+    internal fun onResumableConnectionDisconnected() {
+        isAttemptingResume.set(true)
+    }
 }