From 13d8220e4d9a680fe6cc36e9f11489e2464e9af7 Mon Sep 17 00:00:00 2001
From: Freya Arbjerg <freya.arbjerg@bluestacks.com>
Date: Thu, 14 Dec 2023 13:32:03 +0100
Subject: [PATCH] Make load balancer public

---
 .../audio/internal/AbstractLavakord.kt        |  7 ++--
 .../lavakord/audio/internal/LoadBalancer.kt   | 35 ++++++++++++++-----
 2 files changed, 32 insertions(+), 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
index 0f645cbb..06fbcd47 100644
--- 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
index 2bf25f76..92285391 100644
--- 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.
  */