feat: фоновый пинг апстримов в backoff-откате (probe_interval) — мгновенное возвращение в пул

2026-09-15 01:00:10 +03:00
parent 249b83eb51
commit c6a24a94b0
8 changed files with 452 additions and 0 deletions
@@ -0,0 +1,148 @@
+package pw.binom.llmproxy
+
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.jsonArray
+import kotlinx.serialization.json.jsonObject
+import kotlinx.serialization.json.jsonPrimitive
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNull
+import kotlin.time.Duration.Companion.milliseconds
+import kotlin.time.Duration.Companion.minutes
+import kotlin.time.Duration.Companion.seconds
+
+class HealthProberTest {
+
+    private val json = Json { ignoreUnknownKeys = true }
+
+    @Test
+    fun effectiveProbeIntervalPrefersUpstreamThenProviderThenDefault() {
+        val provider = ProviderConf("p", "https://x", "", null, null, probe_interval = 2.seconds)
+        val upstream = UpstreamConf("u", "p", "m", null, null, null, null, null, probe_interval = 5.milliseconds)
+        val upstreamNoOwn = UpstreamConf("u2", "p", "m")
+        assertEquals(5.milliseconds, effectiveProbeInterval(upstream, provider))
+        assertEquals(2.seconds, effectiveProbeInterval(upstreamNoOwn, provider))
+        assertEquals(DEFAULT_PROBE_INTERVAL, effectiveProbeInterval(upstreamNoOwn, null))
+    }
+
+    @Test
+    fun effectiveProbeIntervalZeroDisablesProber() {
+        val up = UpstreamConf("u", "p", "m", null, null, null, null, null, probe_interval = 0.milliseconds)
+        assertNull(effectiveProbeInterval(up, null))
+
+        val providerZero = ProviderConf("p", "https://x", "", null, null, probe_interval = 0.milliseconds)
+        assertNull(effectiveProbeInterval(UpstreamConf("u2", "p", "m"), providerZero))
+    }
+
+    @Test
+    fun buildProbeBodyIsMinimalChatWithModelReplacedAndPatchesApplied() {
+        val provider = ProviderConf(
+            "p", "https://x", "", null,
+            Json.parseToJsonElement("""{"provider":{"allow_fallbacks":false}}""").jsonObject,
+            probe_interval = 1.seconds,
+        )
+        val up = UpstreamConf(
+            "u", "p", "real-1", null,
+            Json.parseToJsonElement("""{"temperature":0.5}""").jsonObject,
+            null, null, null, 1.seconds,
+        )
+
+        val body = buildProbeBody(up, provider)
+
+        assertEquals("real-1", body["model"]?.jsonPrimitive?.content)
+        assertEquals("2+2=?", body["messages"]?.jsonArray?.first()?.jsonObject?.get("content")?.jsonPrimitive?.content)
+        assertEquals("true", body["stream"]?.jsonPrimitive?.content)
+        assertEquals("4", body["max_tokens"]?.jsonPrimitive?.content)
+        assertEquals("false", body["provider"]?.jsonObject?.get("allow_fallbacks")?.jsonPrimitive?.content)
+        assertEquals("0.5", body["temperature"]?.jsonPrimitive?.content)
+    }
+
+    @Test
+    fun buildProbeBodyIgnoresModelPatch() {
+        val up = UpstreamConf("u", "p", "real-1")
+        val provider = ProviderConf("p", "https://x")
+        val body = buildProbeBody(up, provider)
+        assertEquals(
+            """{"model":"real-1","messages":[{"role":"user","content":"2+2=?"}],"max_tokens":4,"stream":true}""",
+            body.toString(),
+        )
+    }
+
+    @Test
+    fun probeLoopRecoversUpstreamWhenPingSucceeds() = runTest {
+        val provider = ProviderConf("p", "https://x", "", null, null, backoff = 5.minutes)
+        val up = UpstreamConf("u", "p", "m", null, null, null, backoff = 5.minutes, probe_interval = 5.milliseconds)
+        val backoff = BackoffRegistry(mapOf("p" to provider), mapOf("u" to up))
+        backoff.recordFailure(up)
+        assertEquals(true, backoff.isCooling(up))
+
+        val metrics = MetricsRegistry()
+        val prober = HealthProber(
+            providersById = mapOf("p" to provider),
+            upstreams = listOf(up),
+            http = createHttpClient(),
+            backoff = backoff,
+            metrics = metrics,
+            scope = this,
+        )
+        var pings = 0
+        val job = launch {
+            prober.probeLoop(up, 5.milliseconds) { pinged ->
+                pings++
+                assertEquals("u", pinged.id)
+                true
+            }
+        }
+        runUntil { pings >= 1 }
+        // успех уже зафиксирован (recordSuccess) — даём циклу несколько
+        // итераций и убеждаемся, что повторных пингов не было
+        delay(10.milliseconds)
+        job.cancel()
+
+        assertEquals(1, pings)
+        assertEquals(false, backoff.isCooling(up))
+    }
+
+    @Test
+    fun probeLoopKeepsCoolingWhenPingFails() = runTest {
+        val provider = ProviderConf("p", "https://x", "", null, null, backoff = 5.minutes)
+        val up = UpstreamConf("u", "p", "m", null, null, null, backoff = 5.minutes, probe_interval = 5.milliseconds)
+        val backoff = BackoffRegistry(mapOf("p" to provider), mapOf("u" to up))
+        backoff.recordFailure(up)
+
+        val metrics = MetricsRegistry()
+        val prober = HealthProber(
+            providersById = mapOf("p" to provider),
+            upstreams = listOf(up),
+            http = createHttpClient(),
+            backoff = backoff,
+            metrics = metrics,
+            scope = this,
+        )
+        var pings = 0
+        val job = launch {
+            prober.probeLoop(up, 5.milliseconds) { _ ->
+                pings++
+                false
+            }
+        }
+        runUntil { pings >= 2 }
+        job.cancel()
+
+        assertEquals(true, backoff.isCooling(up))
+    }
+
+    private suspend fun TestScope.runUntil(condition: () -> Boolean) {
+        var i = 0
+        while (!condition() && i < 10_000) {
+            delay(1)
+            i++
+        }
+        if (!condition()) throw AssertionError("condition not met in time")
+    }
+
+}