Refactor Node Penalties into an interface that can be overriden by consumers

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

@@ -7,7 +7,8 @@ import dev.arbjerg.lavalink.client.player.Track
 import dev.arbjerg.lavalink.client.player.toCustom
 import dev.arbjerg.lavalink.internal.*
 import dev.arbjerg.lavalink.internal.error.RestException
-import dev.arbjerg.lavalink.internal.loadbalancing.Penalties
+import dev.arbjerg.lavalink.client.loadbalancing.builtin.DefaultNodeHealthProvider
+import dev.arbjerg.lavalink.client.loadbalancing.builtin.INodeHealthProvider
 import dev.arbjerg.lavalink.internal.toLavalinkPlayer
 import dev.arbjerg.lavalink.protocol.v4.*
 import kotlinx.serialization.DeserializationStrategy
@@ -57,7 +58,7 @@ class LavalinkNode(
     val ws = LavalinkSocket(this)
 
     // Stuff for load balancing
-    val penalties = Penalties(this)
+    var nodeHealth: INodeHealthProvider = nodeOptions.nodeHealthProvider ?: DefaultNodeHealthProvider(this)
     var stats: Stats? = null
         internal set
     var available: Boolean = false

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

@@ -1,19 +1,22 @@
 package dev.arbjerg.lavalink.client
 
 import dev.arbjerg.lavalink.client.loadbalancing.IRegionFilter
+import dev.arbjerg.lavalink.client.loadbalancing.builtin.INodeHealthProvider
 import dev.arbjerg.lavalink.internal.TIMEOUT_MS
 import java.net.URI
 
 data class NodeOptions private constructor(val name: String,
                        val serverUri: URI,
                        val password: String,
                        val regionFilter: IRegionFilter?,
+                       val nodeHealthProvider: INodeHealthProvider?,
                        val httpTimeout: Long) {
     data class Builder(
         private var name: String? = null,
         private var serverUri: URI? = null,
         private var password: String? = null,
         private var regionFilter: IRegionFilter? = null,
+        private var nodeHealthProvider: INodeHealthProvider? = null,
         private var httpTimeout: Long = TIMEOUT_MS,
     ) {
         fun setName(name: String) = apply { this.name = name }
@@ -39,6 +42,11 @@ data class NodeOptions private constructor(val name: String,
          */
         fun setRegionFilter(regionFilter: IRegionFilter?) = apply { this.regionFilter = regionFilter }
 
+        /**
+         * Sets a custom node health provider for the node. Used in loadbalancing/traffic routing (Default: none)
+         */
+        fun setNodeHealthProvider(nodeHealthProvider: INodeHealthProvider?) = apply { this.nodeHealthProvider = nodeHealthProvider }
+
         /**
          * Sets the http total call timeout. (Default: 10000ms)
          * @param httpTimeout - timeout in ms
@@ -55,6 +63,7 @@ data class NodeOptions private constructor(val name: String,
                 serverUri!!,
                 password!!,
                 regionFilter,
+                nodeHealthProvider,
                 httpTimeout)
         }
     }

src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/DefaultLoadBalancer.kt

@@ -33,7 +33,7 @@ class DefaultLoadBalancer(private val client: LavalinkClient) : ILoadBalancer {
 
         // TODO: Probably should enforce that no nodes go above the max
         return nodes.filter { it.available }.minByOrNull { node ->
-            node.penalties.calculateTotal() + penaltyProviders.sumOf { it.getPenalty(node, region) }
+            node.nodeHealth.calculateTotalHealthPenalty() + penaltyProviders.sumOf { it.getPenalty(node, region) }
         } ?: throw IllegalStateException("No available nodes!")
     }
 }

src/main/kotlin/dev/arbjerg/lavalink/internal/loadbalancing/Penalties.kt → src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/DefaultNodeHealthProvider.kt

@@ -1,7 +1,9 @@
-package dev.arbjerg.lavalink.internal.loadbalancing
+package dev.arbjerg.lavalink.client.loadbalancing.builtin
 
 import dev.arbjerg.lavalink.client.LavalinkNode
 import dev.arbjerg.lavalink.client.loadbalancing.MAX_ERROR
+import dev.arbjerg.lavalink.internal.loadbalancing.MetricService
+import dev.arbjerg.lavalink.internal.loadbalancing.MetricType
 import dev.arbjerg.lavalink.protocol.v4.Message
 import kotlin.math.pow
 
@@ -13,10 +15,10 @@ import kotlin.math.pow
 // Tracks stuck per minute
 // Track exceptions per minute
 // loads failed per minute
-data class Penalties(val node: LavalinkNode) {
+data class DefaultNodeHealthProvider(val node: LavalinkNode): INodeHealthProvider {
     private val metricService = MetricService()
 
-    fun handleTrackEvent(event: Message.EmittedEvent) {
+    override fun handleTrackEvent(event: Message.EmittedEvent) {
         when (event) {
             is Message.EmittedEvent.TrackStartEvent -> {
                 metricService.trackMetric(MetricType.LOAD_ATTEMPT)
@@ -42,7 +44,7 @@ data class Penalties(val node: LavalinkNode) {
         }
     }
 
-    fun calculateTotal(): Int {
+    override fun calculateTotalHealthPenalty(): Int {
         val stats = node.stats
 
         if (!node.available || stats == null) {
@@ -87,4 +89,13 @@ data class Penalties(val node: LavalinkNode) {
 
         return playerPenalty + cpuPenalty + deficitFramePenalty + nullFramePenalty + trackStuckPenalty + trackExceptionPenalty + loadFailedPenalty
     }
+
+    override fun isHealthy(): Boolean {
+        val metrics = metricService.getCurrentMetrics()
+        val loadsAttempted = metrics[MetricType.LOAD_ATTEMPT] ?: 0
+        val loadsFailed = metrics[MetricType.LOAD_FAILED] ?: 0
+
+        // When the node fails to load anything, we consider it to be unhealthy
+        return !(loadsAttempted > 0 && loadsAttempted == loadsFailed) && node.available
+    }
 }

src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/INodeHealthProvider.kt

@@ -0,0 +1,27 @@
+package dev.arbjerg.lavalink.client.loadbalancing.builtin
+
+import dev.arbjerg.lavalink.client.loadbalancing.MAX_ERROR
+import dev.arbjerg.lavalink.protocol.v4.Message
+
+interface INodeHealthProvider {
+    /**
+     * Called for each event on the node.
+     */
+    fun handleTrackEvent(event: Message.EmittedEvent)
+
+    /**
+     * Calculate the penalty for the node based off of its health.
+     *
+     * Return value should never exceed [MAX_ERROR]. Lower means to take preference.
+     *
+     * @return A number between 0 and [MAX_ERROR] (inclusive), using numbers outside of this range may cause errors.
+     */
+    fun calculateTotalHealthPenalty(): Int
+
+    /**
+     * Gives a simple answer if the node is considered healthy.
+     *
+     * @return true if the node is in a healthy state
+     */
+    fun isHealthy() : Boolean
+}

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

@@ -104,7 +104,7 @@ class LavalinkSocket(private val node: LavalinkNode) : WebSocketListener(), Clos
                     else -> {}
                 }
 
-                node.penalties.handleTrackEvent(event)
+                node.nodeHealth.handleTrackEvent(event)
             }
 
             else -> {