-
Notifications
You must be signed in to change notification settings - Fork 11
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Showing
21 changed files
with
464 additions
and
42 deletions.
There are no files selected for viewing
This file was deleted.
Oops, something went wrong.
15 changes: 0 additions & 15 deletions
15
core/health/src/main/kotlin/dev/schlaubi/mikbot/core/health/check/KordHealthCheck.kt
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,6 @@ | ||
FROM --platform=$TARGETOS/$TARGETARCH eclipse-temurin:22-jre-alpine | ||
|
||
WORKDIR /usr/app | ||
COPY build/install/bot-kubernetes . | ||
|
||
ENTRYPOINT ["/usr/app/bin/mikmusic"] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,82 @@ | ||
# kubernetes | ||
|
||
Plugin providing tools for running Mikbot within kubernetes | ||
|
||
Also shouteout to @lucsoft, who helped me to set this up | ||
|
||
## health endpoint | ||
|
||
The plugin adds a HTTP endpoint, which you can use as a K8s probe, it will return 200 unless a shard or the DB connection | ||
is down | ||
|
||
```yaml | ||
livenessProbe: | ||
httpGet: | ||
path: /healthz | ||
port: mikbot | ||
startupProbe: | ||
httpGet: | ||
path: /healthz | ||
port: mikbot # remember to specify that port name somewhere | ||
# let's give our shards time to connect | ||
initialDelaySeconds: 15 | ||
``` | ||
### Own health checks | ||
You can add more checks by extending the `HealthCheck` [extension point](../../PLUGINS.md#what-are-extension-points) | ||
|
||
## Scaling | ||
|
||
The bot supports very basic distribution of shards through a stateful set, here is an example configuration | ||
|
||
The bot will automatically re-balance if enough new guilds are created that a new shard is required | ||
|
||
```yaml | ||
apiVersion: apps/v1 | ||
kind: StatefulSet | ||
metadata: | ||
name: test-bot | ||
namespace: default | ||
spec: | ||
serviceName: test-bot | ||
replicas: 4 | ||
ordinals: | ||
start: 0 | ||
selector: | ||
matchLabels: | ||
app: test-bot | ||
template: | ||
metadata: | ||
labels: | ||
app: test-bot | ||
spec: | ||
nodeSelector: | ||
fast: "true" | ||
containers: | ||
- name: test-bot | ||
image: ghcr.io/drschlaubi/k8s-test-bot:latest | ||
imagePullPolicy: Always | ||
livenessProbe: | ||
httpGet: | ||
path: /healthz | ||
port: mikbot | ||
startupProbe: | ||
httpGet: | ||
path: /healthz | ||
port: mikbot | ||
initialDelaySeconds: 15 | ||
ports: | ||
- containerPort: 8080 | ||
name: mikbot | ||
env: | ||
- name: TOTAL_SHARDS | ||
value: "8" | ||
- name: ENABLE_SCALING | ||
value: "true" | ||
- name: SHARDS_PER_POD # default value is true | ||
value: "2" | ||
- name: POD_ID | ||
valueFrom: | ||
fieldRef: | ||
fieldPath: metadata.labels['apps.kubernetes.io/pod-index'] | ||
``` |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,46 @@ | ||
plugins { | ||
`mikbot-module` | ||
`mikbot-publishing` | ||
alias(libs.plugins.kotlinx.serialization) | ||
com.google.devtools.ksp | ||
dev.schlaubi.mikbot.`gradle-plugin` | ||
`jvm-test-suite` | ||
} | ||
|
||
group = "dev.schlaubi.mikbot" | ||
version = mikbotVersion | ||
|
||
repositories { | ||
maven("https://jitpack.io") | ||
} | ||
|
||
dependencies { | ||
plugin(projects.core.ktor) | ||
|
||
implementation(libs.kubernetes.client) | ||
implementation(libs.kotlin.jsonpatch) | ||
|
||
testImplementation(kotlin("test-junit5")) | ||
testImplementation(libs.system.rules) | ||
testImplementation(libs.kord.core) | ||
} | ||
|
||
testing { | ||
suites { | ||
@Suppress("UnstableApiUsage") | ||
named<JvmTestSuite>("test") { | ||
useJUnitJupiter() | ||
} | ||
} | ||
} | ||
|
||
tasks { | ||
test { | ||
jvmArgs = | ||
listOf("--add-opens", "java.base/java.lang=ALL-UNNAMED", "--add-opens", "java.base/java.lang=ALL-UNNAMED") | ||
} | ||
} | ||
|
||
mikbotPlugin { | ||
description = "Plugin providing an /healthz endpoint used for health checking." | ||
} |
16 changes: 16 additions & 0 deletions
16
core/kubernetes/src/main/kotlin/dev/schlaubi/mikbot/core/health/Config.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,16 @@ | ||
package dev.schlaubi.mikbot.core.health | ||
|
||
import dev.schlaubi.mikbot.plugin.api.EnvironmentConfig | ||
|
||
@Suppress("PropertyName") | ||
open class Config : EnvironmentConfig() { | ||
val ENABLE_SCALING by getEnv(false, String::toBooleanStrict) | ||
val POD_ID by getEnv(transform = String::toInt) | ||
val SHARDS_PER_POD by getEnv(2, String::toInt) | ||
val TOTAL_SHARDS by getEnv(transform = String::toInt) | ||
val STATEFUL_SET_NAME by this | ||
val NAMESPACE by getEnv("default") | ||
val CONTAINER_NAME by this | ||
|
||
companion object : Config() | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
148 changes: 148 additions & 0 deletions
148
core/kubernetes/src/main/kotlin/dev/schlaubi/mikbot/core/health/Rebalancer.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,148 @@ | ||
package dev.schlaubi.mikbot.core.health | ||
|
||
import com.kotlindiscord.kord.extensions.commands.Arguments | ||
import com.kotlindiscord.kord.extensions.commands.converters.impl.optionalInt | ||
import com.kotlindiscord.kord.extensions.extensions.ephemeralSlashCommand | ||
import com.kotlindiscord.kord.extensions.extensions.event | ||
import com.reidsync.kxjsonpatch.generatePatch | ||
import dev.kord.common.KordConstants | ||
import dev.kord.core.event.guild.GuildCreateEvent | ||
import dev.kord.rest.json.response.BotGatewayResponse | ||
import dev.kord.rest.route.Route | ||
import dev.schlaubi.mikbot.core.health.check.ready | ||
import dev.schlaubi.mikbot.core.health.util.blocking | ||
import dev.schlaubi.mikbot.plugin.api.PluginContext | ||
import dev.schlaubi.mikbot.plugin.api.module.MikBotModule | ||
import dev.schlaubi.mikbot.plugin.api.owner.ownerOnly | ||
import dev.schlaubi.mikbot.plugin.api.util.discordError | ||
import io.ktor.client.request.get | ||
import io.ktor.client.request.header | ||
import io.ktor.client.statement.bodyAsText | ||
import io.ktor.http.HttpHeaders.Authorization | ||
import io.ktor.http.HttpHeaders.UserAgent | ||
import io.kubernetes.client.custom.V1Patch | ||
import io.kubernetes.client.openapi.apis.AppsV1Api | ||
import io.kubernetes.client.openapi.models.V1StatefulSet | ||
import io.kubernetes.client.util.PatchUtils | ||
import kotlinx.serialization.encodeToString | ||
import kotlinx.serialization.json.Json | ||
import kotlinx.serialization.json.addJsonObject | ||
import kotlinx.serialization.json.buildJsonObject | ||
import kotlinx.serialization.json.put | ||
import kotlinx.serialization.json.putJsonArray | ||
import kotlinx.serialization.json.putJsonObject | ||
import kotlin.time.Duration.Companion.minutes | ||
import kotlin.time.TimeSource | ||
import io.kubernetes.client.util.Config as KubeConfig | ||
|
||
private val startup = TimeSource.Monotonic.markNow() | ||
|
||
private class RebalanceArguments() : Arguments() { | ||
val forceTo by optionalInt { | ||
name = "force_to" | ||
description = "commands.rebalance.arguments.force_to.description" | ||
} | ||
} | ||
|
||
class RebalancerExtension(context: PluginContext) : MikBotModule(context) { | ||
override val name: String = "rebalancer" | ||
override val bundle: String = "kubernetes" | ||
|
||
private val kubeClient = KubeConfig.defaultClient() | ||
private val kubeApi = AppsV1Api(kubeClient) | ||
|
||
override suspend fun setup() { | ||
command() | ||
|
||
// Only register self-update on the first pod | ||
if (Config.POD_ID == 0) { | ||
eventListener() | ||
} | ||
} | ||
|
||
private suspend fun eventListener() = event<GuildCreateEvent> { | ||
check { | ||
failIfNot(ready) | ||
// Do not call listener when pod just started | ||
// this is to prevent initial guild_creates to cause this | ||
failIf(startup.elapsedNow() > 5.minutes) | ||
} | ||
action { | ||
val (_, newTotalShards) = getGatewayInfo() | ||
if (newTotalShards > Config.TOTAL_SHARDS) { | ||
reBalance(newTotalShards) | ||
} | ||
} | ||
} | ||
|
||
private suspend fun command() = ephemeralSlashCommand(::RebalanceArguments) { | ||
name = "rebalance" | ||
description = "commands.rebalance.description" | ||
ownerOnly() | ||
|
||
action { | ||
val (_, newTotalShards) = getGatewayInfo() | ||
if (arguments.forceTo != null) { | ||
reBalance(arguments.forceTo!!) | ||
} else if (newTotalShards > Config.TOTAL_SHARDS) { | ||
reBalance(newTotalShards) | ||
} else if (arguments.forceTo != null) { | ||
reBalance(arguments.forceTo!!) | ||
} else { | ||
discordError(translate("commands.rebalance.already_balanced")) | ||
} | ||
|
||
respond { | ||
content = translate("commands.rebalance.done") | ||
} | ||
} | ||
|
||
} | ||
|
||
private suspend fun getGatewayInfo(): BotGatewayResponse { | ||
val response = kord.resources.httpClient.get("${Route.baseUrl}${Route.GatewayBotGet.path}") { | ||
header(UserAgent, KordConstants.USER_AGENT) | ||
header(Authorization, "Bot ${kord.resources.token}") | ||
} | ||
|
||
return Json.decodeFromString(response.bodyAsText()) | ||
} | ||
|
||
private suspend fun reBalance(newTotalShards: Int) { | ||
val json = | ||
generateStatefulSetSpec(replicas = newTotalShards / Config.SHARDS_PER_POD, totalShards = newTotalShards) | ||
|
||
val jsonPatch = generateStatefulSetSpec().generatePatch(json) | ||
|
||
val patch = V1Patch(Json.encodeToString(jsonPatch)) | ||
|
||
blocking { | ||
PatchUtils.patch(V1StatefulSet::class.java, { | ||
kubeApi.patchNamespacedStatefulSet(Config.STATEFUL_SET_NAME, Config.NAMESPACE, patch) | ||
.buildCall(null) | ||
}, "application/json-patch+json", kubeClient) | ||
} | ||
} | ||
} | ||
|
||
private fun generateStatefulSetSpec(replicas: Int = 0, totalShards: Int = 0) = buildJsonObject { | ||
putJsonObject("spec") { | ||
put("replicas", replicas) | ||
|
||
putJsonObject("template") { | ||
putJsonObject("spec") { | ||
putJsonArray("containers") { | ||
addJsonObject { | ||
put("name", Config.CONTAINER_NAME) | ||
putJsonArray("env") { | ||
addJsonObject { | ||
put("name", "TOTAL_SHARDS") | ||
put("value", totalShards.toString()) | ||
} | ||
} | ||
} | ||
} | ||
} | ||
} | ||
} | ||
} |
Oops, something went wrong.