Make load balancer public

botlabs-gg · Dec 14, 2023 · 13d8220 · 13d8220
1 parent b3e5c4f
commit 13d8220
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 10 deletions.
diff --git a/core/src/commonMain/kotlin/dev/schlaubi/lavakord/audio/internal/AbstractLavakord.kt b/core/src/commonMain/kotlin/dev/schlaubi/lavakord/audio/internal/AbstractLavakord.kt
@@ -132,8 +132,11 @@ public abstract class AbstractLavakord internal constructor(
     }
 
 
-    @Suppress("LeakingThis")
-    private val loadBalancer: LoadBalancer = LoadBalancer(options.loadBalancer.penaltyProviders, this)
+    /**
+     * The load balancer used to select nodes
+     */
+    @Suppress("LeakingThis", "MemberVisibilityCanBePrivate")
+    public val loadBalancer: LoadBalancer = LoadBalancer(options.loadBalancer.penaltyProviders, this)
 
     override val nodes: List<Node>
         get() = nodesMap.values.toList()

diff --git a/core/src/commonMain/kotlin/dev/schlaubi/lavakord/audio/internal/LoadBalancer.kt b/core/src/commonMain/kotlin/dev/schlaubi/lavakord/audio/internal/LoadBalancer.kt
@@ -4,22 +4,24 @@ import dev.schlaubi.lavakord.LavaKord
 import dev.schlaubi.lavakord.audio.Node
 import kotlin.math.pow
 
-internal class LoadBalancer(
+public class LoadBalancer(
     private val penaltyProviders: List<PenaltyProvider>,
     private val lavakord: LavaKord
 ) {
 
-    fun determineBestNode(guildId: ULong): Node? = lavakord.nodes
+    internal fun determineBestNode(guildId: ULong): Node? = lavakord.nodes
         .asSequence()
         .filter(Node::available)
-        .minByOrNull { calculatePenalties(it, penaltyProviders, guildId) }
+        .minByOrNull { calculatePenalties(it, guildId).sum }
 
-    // Inspired by: https://github.com/freyacodes/Lavalink-Client/blob/master/src/main/java/lavalink/client/io/LavalinkLoadBalancer.java#L111
-    private fun calculatePenalties(
+    /**
+     * Calculate the penalties for a given guild
+     * Adapted from https://github.com/freyacodes/Lavalink-Client/blob/master/src/main/java/lavalink/client/io/LavalinkLoadBalancer.java#L111
+     */
+    public fun calculatePenalties(
         node: Node,
-        penaltyProviders: List<PenaltyProvider>,
         guildId: ULong
-    ): Int {
+    ): Penalties {
         val playerPenalty: Int
         val cpuPenalty: Int
         val deficitFramePenalty: Int
@@ -46,10 +48,27 @@ internal class LoadBalancer(
                 nullFramePenalty = 0
             }
         }
-        return playerPenalty + cpuPenalty + deficitFramePenalty + nullFramePenalty + customPenalties
+        return Penalties(playerPenalty, cpuPenalty, deficitFramePenalty, nullFramePenalty, customPenalties)
     }
 }
 
+/** Result of penalties used for load balancing */
+public data class Penalties(
+    /** Penalty due to number of players */
+    val playerPenalty: Int,
+    /** Penalty due to high CPU */
+    val cpuPenalty: Int,
+    /** Penalty due to Lavalink struggling to send frames */
+    val deficitFramePenalty: Int,
+    /** Penalty due to Lavaplayer struggling to provide frames */
+    val nullFramePenalty: Int,
+    /** Penalties from [PenaltyProvider]*/
+    val customPenalties: Int
+) {
+    /** The sum of all penalties */
+    val sum: Int = playerPenalty + cpuPenalty + deficitFramePenalty + nullFramePenalty + customPenalties
+}
+
 /**
  * Interface to accept custom penalties for [Node]s.
  */