feat(relay): graceful drain on LocalRelayServer.stop

Closes the "100ms grace can drop in-flight EVENTs" gap from the audit. Behavioural change: - stop() default grace is now 5 s (was 100 ms) and total timeout 10 s (was 1 s). A SQLite write + OK reply round-trip easily fits. - stop() first sends a NOTICE("closing: relay is shutting down — please reconnect later") to every connected client, so well- behaved clients can reconnect rather than hammering a dead socket. - Active client sessions are tracked in a ConcurrentHashMap-backed set populated by the WebSocket handler's connect/finally pair. Exposed read-only via [activeSessionCount] so operators (and tests) can observe lifecycle. - stop() is now idempotent. Tests (4 new in GracefulShutdownTest): - activeSessionCountTracksConnectAndDisconnect — counter goes 0 → 1 on subscribe, → 0 on disconnect. - stopSendsShutdownNoticeToActiveClients — connected client receives a NOTICE whose message starts with "closing:". - stopIsIdempotent — second stop() is a safe no-op. - rawWsClientObservesNoticeBeforeServerCloses — bare OkHttp ws client sees the NOTICE frame before the server closes the socket (proves the order: NOTICE first, then engine.stop). Total :quartz-relay tests: 66, 0 failures.
2026-05-07 02:22:39 +00:00
parent 633d0c3c74
commit 65311590f2
2 changed files with 276 additions and 4 deletions
@@ -20,6 +20,8 @@
 */
 package com.vitorpamplona.quartz.relay

+import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.NoticeMessage
+import com.vitorpamplona.quartz.nip01Core.relay.server.RelaySession
 import io.ktor.http.ContentType
 import io.ktor.http.HttpHeaders
 import io.ktor.http.HttpStatusCode
@@ -37,6 +39,7 @@ import io.ktor.websocket.Frame
 import io.ktor.websocket.readText
 import kotlinx.coroutines.channels.consumeEach
 import kotlinx.coroutines.runBlocking
+import java.util.concurrent.ConcurrentHashMap

 /**
 * Hosts a [Relay] over a real `ws://` endpoint backed by Ktor + CIO.
@@ -73,6 +76,17 @@ class LocalRelayServer(
    private var engine: CIOApplicationEngine? = null
    private var resolvedPort: Int = -1

+    /**
+     * Active client sessions, registered when their WebSocket handler
+     * runs and removed on disconnect. Exposed (read-only) so [stop] can
+     * NOTICE every connected client during graceful drain, and so tests
+     * can assert lifecycle bookkeeping.
+     */
+    private val activeSessions: MutableSet<RelaySession> = ConcurrentHashMap.newKeySet()
+
+    /** Number of WebSocket sessions currently connected to the server. */
+    val activeSessionCount: Int get() = activeSessions.size
+
    /** `ws://host:port/path` — only valid after [start]. */
    val url: String
        get() {
@@ -119,6 +133,7 @@ class LocalRelayServer(
                                // dispatcher; trySend never blocks the relay thread.
                                outgoing.trySend(Frame.Text(json))
                            }
+                        activeSessions.add(session)
                        try {
                            incoming.consumeEach { frame ->
                                if (frame is Frame.Text) {
@@ -126,6 +141,7 @@ class LocalRelayServer(
                                }
                            }
                        } finally {
+                            activeSessions.remove(session)
                            session.close()
                        }
                    }
@@ -145,13 +161,45 @@ class LocalRelayServer(
        return this
    }

-    /** Stops the engine. Safe to call multiple times. */
+    /**
+     * Graceful shutdown. Safe to call multiple times.
+     *
+     * 1. Sends a NOTICE("closing: …") to every currently-connected
+     *    client so well-behaved clients know to reconnect later.
+     * 2. Stops the Ktor engine: rejects new connections immediately,
+     *    then waits up to [gracePeriodMillis] for active WebSocket
+     *    handlers to finish whatever they're processing (so an in-flight
+     *    `EVENT` lands its `OK` reply before the socket dies). After
+     *    the grace window, in-progress handlers are cancelled and the
+     *    engine waits up to [timeoutMillis] - [gracePeriodMillis] for
+     *    that cancellation to complete.
+     *
+     * Defaults to 5 s grace / 10 s total — generous enough that a
+     * SQLite write + reply round-trip can land for typical event
+     * sizes. Override either with a tighter budget if your operator
+     * knows their workload.
+     */
    fun stop(
-        gracePeriodMillis: Long = 100,
-        timeoutMillis: Long = 1_000,
+        gracePeriodMillis: Long = 5_000,
+        timeoutMillis: Long = 10_000,
    ) {
-        engine?.stop(gracePeriodMillis, timeoutMillis)
+        val e = engine ?: return
+        notifyShutdown()
+        e.stop(gracePeriodMillis, timeoutMillis)
        engine = null
        resolvedPort = -1
    }
+
+    /**
+     * Best-effort NOTICE to every active client. Failures are
+     * swallowed — a flaky socket on its way out is exactly the case
+     * where a NOTICE will fail anyway, and the client's read of the
+     * close frame is the authoritative shutdown signal.
+     */
+    private fun notifyShutdown() {
+        val notice = NoticeMessage("closing: relay is shutting down — please reconnect later")
+        activeSessions.forEach { session ->
+            runCatching { session.send(notice) }
+        }
+    }
 }
@@ -0,0 +1,224 @@
+/*
+ * Copyright (c) 2025 Vitor Pamplona
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of
+ * this software and associated documentation files (the "Software"), to deal in
+ * the Software without restriction, including without limitation the rights to use,
+ * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
+ * Software, and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+ * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+ * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
+ * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+ * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+package com.vitorpamplona.quartz.relay
+
+import com.vitorpamplona.quartz.nip01Core.relay.client.NostrClient
+import com.vitorpamplona.quartz.nip01Core.relay.client.listeners.RelayConnectionListener
+import com.vitorpamplona.quartz.nip01Core.relay.client.reqs.SubscriptionListener
+import com.vitorpamplona.quartz.nip01Core.relay.client.single.IRelayClient
+import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.Message
+import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.NoticeMessage
+import com.vitorpamplona.quartz.nip01Core.relay.filters.Filter
+import com.vitorpamplona.quartz.nip01Core.relay.normalizer.NormalizedRelayUrl
+import com.vitorpamplona.quartz.nip01Core.relay.normalizer.normalizeRelayUrl
+import com.vitorpamplona.quartz.nip01Core.relay.sockets.okhttp.BasicOkHttpWebSocket
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.Channel.Factory.UNLIMITED
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withTimeout
+import kotlinx.coroutines.withTimeoutOrNull
+import okhttp3.OkHttpClient
+import kotlin.test.AfterTest
+import kotlin.test.BeforeTest
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNotNull
+import kotlin.test.assertTrue
+
+/**
+ * Tests [LocalRelayServer.stop] honours the graceful-shutdown contract:
+ *  1. Active clients receive a `NOTICE` warning of imminent shutdown.
+ *  2. The active session counter accurately tracks open WS sessions.
+ *  3. After `stop()` returns, no sessions remain registered.
+ */
+class GracefulShutdownTest {
+    private lateinit var relay: Relay
+    private lateinit var server: LocalRelayServer
+    private lateinit var scope: CoroutineScope
+    private lateinit var client: NostrClient
+
+    private val httpClient = OkHttpClient.Builder().build()
+
+    @BeforeTest
+    fun setup() {
+        val placeholder = "ws://127.0.0.1:7771/".normalizeRelayUrl()
+        relay = Relay(url = placeholder)
+        server = LocalRelayServer(relay, host = "127.0.0.1", port = 0).start()
+        scope = CoroutineScope(Dispatchers.Default + SupervisorJob())
+        val builder = BasicOkHttpWebSocket.Builder { _ -> httpClient }
+        client = NostrClient(builder, scope)
+    }
+
+    @AfterTest
+    fun teardown() {
+        client.disconnect()
+        scope.cancel()
+        // server may already be stopped by the test; calling stop()
+        // again is a no-op.
+        server.stop(gracePeriodMillis = 200, timeoutMillis = 500)
+        relay.close()
+    }
+
+    @Test
+    fun activeSessionCountTracksConnectAndDisconnect() =
+        runBlocking {
+            assertEquals(0, server.activeSessionCount, "no clients yet")
+
+            // Open a connection by subscribing — wait for EOSE so we
+            // know the WebSocket handshake completed and the relay
+            // session has registered.
+            val gotEose = Channel<Unit>(UNLIMITED)
+            val relayUrl = server.url.normalizeRelayUrl()
+            client.subscribe(
+                "track-1",
+                mapOf(relayUrl to listOf(Filter(kinds = listOf(1)))),
+                object : SubscriptionListener {
+                    override fun onEose(
+                        relay: NormalizedRelayUrl,
+                        forFilters: List<Filter>?,
+                    ) {
+                        gotEose.trySend(Unit)
+                    }
+                },
+            )
+            withTimeout(5000) { gotEose.receive() }
+
+            assertEquals(1, server.activeSessionCount, "one connected session")
+
+            client.unsubscribe("track-1")
+            client.disconnect()
+
+            // Disconnect happens asynchronously on the relay side; allow
+            // a short window for the handler's `finally` block to run.
+            withTimeoutOrNull(2000) {
+                while (server.activeSessionCount > 0) kotlinx.coroutines.delay(10)
+            }
+            assertEquals(0, server.activeSessionCount, "session must be removed after disconnect")
+        }
+
+    @Test
+    fun stopSendsShutdownNoticeToActiveClients() =
+        runBlocking {
+            val noticeChannel = Channel<NoticeMessage>(UNLIMITED)
+            val gotEose = Channel<Unit>(UNLIMITED)
+            val listener =
+                object : RelayConnectionListener {
+                    override fun onIncomingMessage(
+                        relay: IRelayClient,
+                        msgStr: String,
+                        msg: Message,
+                    ) {
+                        if (msg is NoticeMessage) noticeChannel.trySend(msg)
+                    }
+                }
+            client.addConnectionListener(listener)
+
+            val relayUrl = server.url.normalizeRelayUrl()
+            client.subscribe(
+                "notice-watch",
+                mapOf(relayUrl to listOf(Filter(kinds = listOf(1)))),
+                object : SubscriptionListener {
+                    override fun onEose(
+                        relay: NormalizedRelayUrl,
+                        forFilters: List<Filter>?,
+                    ) {
+                        gotEose.trySend(Unit)
+                    }
+                },
+            )
+            withTimeout(5000) { gotEose.receive() }
+            assertEquals(1, server.activeSessionCount)
+
+            // Trigger graceful shutdown.
+            server.stop(gracePeriodMillis = 1_000, timeoutMillis = 2_000)
+
+            val notice = withTimeout(5000) { noticeChannel.receive() }
+            assertNotNull(notice)
+            assertTrue(
+                notice.message.startsWith("closing:"),
+                "expected NOTICE to start with 'closing:', got '${notice.message}'",
+            )
+        }
+
+    @Test
+    fun stopIsIdempotent() {
+        // First call shuts the engine down.
+        server.stop(gracePeriodMillis = 100, timeoutMillis = 500)
+        // Second call must be a safe no-op (no exception).
+        server.stop(gracePeriodMillis = 100, timeoutMillis = 500)
+    }
+
+    /**
+     * Sanity check on the grace window: a bare-bones ws client that
+     * connects and never sends anything should *receive* the shutdown
+     * NOTICE before the server fully closes the socket. Uses Ktor's
+     * client-agnostic OkHttp transport directly so we can observe the
+     * raw frames.
+     */
+    @Test
+    fun rawWsClientObservesNoticeBeforeServerCloses() =
+        runBlocking {
+            val httpUrl =
+                server.url
+                    .replace("ws://", "http://")
+            val request =
+                okhttp3.Request
+                    .Builder()
+                    .url(httpUrl)
+                    .build()
+
+            val frames = Channel<String>(UNLIMITED)
+            val socket =
+                httpClient.newWebSocket(
+                    request,
+                    object : okhttp3.WebSocketListener() {
+                        override fun onMessage(
+                            ws: okhttp3.WebSocket,
+                            text: String,
+                        ) {
+                            frames.trySend(text)
+                        }
+                    },
+                )
+
+            try {
+                // Wait until the relay sees the connection.
+                withTimeoutOrNull(2000) {
+                    while (server.activeSessionCount == 0) kotlinx.coroutines.delay(10)
+                }
+                assertEquals(1, server.activeSessionCount)
+
+                server.stop(gracePeriodMillis = 1_000, timeoutMillis = 2_000)
+
+                val text = withTimeout(3000) { frames.receive() }
+                assertTrue(
+                    text.contains("\"NOTICE\"") && text.contains("closing"),
+                    "expected a NOTICE frame, got: $text",
+                )
+            } finally {
+                socket.cancel()
+            }
+        }
+}