From 76103de110ee7b5f85e8c8a280b654f1341d60fd Mon Sep 17 00:00:00 2001 From: kikkia Date: Tue, 26 Mar 2024 22:51:49 -0700 Subject: [PATCH] update node penalty naming to be more inline with function, break it out into an interface to allow for overriding. --- .../arbjerg/lavalink/client/LavalinkNode.kt | 5 ++-- .../arbjerg/lavalink/client/NodeOptions.kt | 9 +++++++ .../builtin/DefaultLoadBalancer.kt | 2 +- .../builtin/DefaultNodeHealthProvider.kt} | 19 ++++++++++--- .../builtin/INodeHealthProvider.kt | 27 +++++++++++++++++++ .../lavalink/internal/LavalinkSocket.kt | 2 +- 6 files changed, 56 insertions(+), 8 deletions(-) rename src/main/kotlin/dev/arbjerg/lavalink/{internal/loadbalancing/Penalties.kt => client/loadbalancing/builtin/DefaultNodeHealthProvider.kt} (81%) create mode 100644 src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/INodeHealthProvider.kt diff --git a/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkNode.kt b/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkNode.kt index a30c57d..8b9fd90 100644 --- a/src/main/kotlin/dev/arbjerg/lavalink/client/LavalinkNode.kt +++ b/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 diff --git a/src/main/kotlin/dev/arbjerg/lavalink/client/NodeOptions.kt b/src/main/kotlin/dev/arbjerg/lavalink/client/NodeOptions.kt index ae62745..f462da0 100644 --- a/src/main/kotlin/dev/arbjerg/lavalink/client/NodeOptions.kt +++ b/src/main/kotlin/dev/arbjerg/lavalink/client/NodeOptions.kt @@ -1,6 +1,7 @@ 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 @@ -8,12 +9,14 @@ 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) } } diff --git a/src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/DefaultLoadBalancer.kt b/src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/DefaultLoadBalancer.kt index 981040c..6050654 100644 --- a/src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/DefaultLoadBalancer.kt +++ b/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!") } } diff --git a/src/main/kotlin/dev/arbjerg/lavalink/internal/loadbalancing/Penalties.kt b/src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/DefaultNodeHealthProvider.kt similarity index 81% rename from src/main/kotlin/dev/arbjerg/lavalink/internal/loadbalancing/Penalties.kt rename to src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/DefaultNodeHealthProvider.kt index 3f1085d..08e5f29 100644 --- a/src/main/kotlin/dev/arbjerg/lavalink/internal/loadbalancing/Penalties.kt +++ b/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 + } } diff --git a/src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/INodeHealthProvider.kt b/src/main/kotlin/dev/arbjerg/lavalink/client/loadbalancing/builtin/INodeHealthProvider.kt new file mode 100644 index 0000000..c4ecd6b --- /dev/null +++ b/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 +} diff --git a/src/main/kotlin/dev/arbjerg/lavalink/internal/LavalinkSocket.kt b/src/main/kotlin/dev/arbjerg/lavalink/internal/LavalinkSocket.kt index 3fb6f60..7681cb1 100644 --- a/src/main/kotlin/dev/arbjerg/lavalink/internal/LavalinkSocket.kt +++ b/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 -> {