Merge pull request #3872 from vitorpamplona/perf/outbox-no-quadratic-publish

Stop the outbox getting slower with every publish

quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/nip01Core/relay/client/pool/PoolEventOutbox.kt

@@ -27,32 +27,61 @@ import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.OkMessage
 import com.vitorpamplona.quartz.nip01Core.relay.commands.toRelay.Command
 import com.vitorpamplona.quartz.nip01Core.relay.commands.toRelay.EventCmd
 import com.vitorpamplona.quartz.nip01Core.relay.normalizer.NormalizedRelayUrl
+import com.vitorpamplona.quartz.utils.cache.LargeCache
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.update
 import kotlin.concurrent.Volatile
 
 class PoolEventOutbox {
-    // @Volatile so the polling path (INostrClient.pendingPublishRelaysFor)
-    // sees current state from threads that didn't write the map. Mutations
-    // still happen on NostrClient's IO scope; this only closes the
-    // visibility gap for cross-thread readers.
-    @Volatile
-    private var eventOutbox = mapOf<HexKey, PoolEventOutboxState>()
+    /**
+     * Pending publishes, keyed by event id.
+     *
+     * A concurrent map, NOT a copy-on-write immutable one. It used to be
+     * `@Volatile var eventOutbox = mapOf(...)` reassigned with
+     * `eventOutbox + Pair(...)`, which copies EVERY entry on EVERY publish —
+     * so publishing N events cost O(N^2). Measured on a bulk push with ~970k
+     * entries resident: 22.7ms per event, of which ~20.5ms was this map, and
+     * the rate decayed as the outbox grew (45.6 -> 44.6 -> 43.2 ev/s across
+     * three windows). The store fetch behind the same loop cost 1.2ms.
+     *
+     * [LargeCache] is ConcurrentHashMap on JVM/Android, so put/get/remove are
+     * O(1) and cross-thread visibility no longer needs the volatile republish.
+     */
+    private val eventOutbox = LargeCache<HexKey, PoolEventOutboxState>()
     val relays = MutableStateFlow(setOf<NormalizedRelayUrl>())
 
+    /**
+     * Removals since the relay set was last rebuilt.
+     *
+     * Deciding whether a relay may leave [relays] means asking whether ANY
+     * remaining entry still wants it — O(outbox), and doing that per publish
+     * is the second half of the quadratic. Additions stay exact and cheap (a
+     * union of the event's own relays); removals are swept in batches, because
+     * keeping a relay in the set slightly too long only means holding a
+     * connection a little longer, while scanning a million entries to retire
+     * it promptly costs the whole push.
+     */
+    @Volatile
+    private var pendingSweep = 0
+
+    companion object {
+        /** Removals between full relay-set rebuilds — see [pendingSweep]. */
+        private const val SWEEP_EVERY = 256
+    }
+
     fun needsToUpdateRelays(): Boolean {
         val currentRelays = relays.value
 
         var relaysToRemoveCounter = 0
 
         currentRelays.forEach { currentRelay ->
-            if (eventOutbox.values.none { currentRelay in it.relaysRemaining }) {
+            if (eventOutbox.values().none { currentRelay in it.relaysRemaining }) {
                 relaysToRemoveCounter++
             }
         }
 
         var relaysToAddCounter = 0
-        eventOutbox.values.forEach { outboxState ->
+        eventOutbox.values().forEach { outboxState ->
             if (outboxState.relaysRemaining.any { it !in currentRelays }) {
                 relaysToAddCounter++
             }
@@ -67,13 +96,13 @@ class PoolEventOutbox {
                 val relaysToRemove = mutableSetOf<NormalizedRelayUrl>()
 
                 currentRelays.forEach { currentRelay ->
-                    if (eventOutbox.values.none { currentRelay in it.relaysRemaining }) {
+                    if (eventOutbox.values().none { currentRelay in it.relaysRemaining }) {
                         relaysToRemove.add(currentRelay)
                     }
                 }
 
                 val relaysToAdd = mutableSetOf<NormalizedRelayUrl>()
-                eventOutbox.values.forEach { outboxState ->
+                eventOutbox.values().forEach { outboxState ->
                     outboxState.relaysRemaining.forEach { relay ->
                         if (relay !in relaysToAdd && relay !in currentRelays) {
                             relaysToAdd.add(relay)
@@ -88,7 +117,7 @@ class PoolEventOutbox {
 
     fun activeOutboxCacheFor(url: NormalizedRelayUrl): Set<HexKey> {
         val myEvents = mutableSetOf<HexKey>()
-        eventOutbox.forEach { (eventId, outboxCache) ->
+        eventOutbox.forEach { eventId, outboxCache ->
             if (url in outboxCache.relaysRemaining) {
                 myEvents.add(eventId)
             }
@@ -103,7 +132,7 @@ class PoolEventOutbox {
      */
     fun activeOutboxEventsFor(url: NormalizedRelayUrl): List<Event> {
         val myEvents = mutableListOf<Event>()
-        eventOutbox.forEach { (_, outboxCache) ->
+        eventOutbox.forEach { _, outboxCache ->
             if (url in outboxCache.relaysRemaining) {
                 myEvents.add(outboxCache.event)
             }
@@ -117,20 +146,41 @@ class PoolEventOutbox {
      * Callers can poll this after publish to detect when relays ack: the set shrinks
      * as OKs arrive, then the entry is removed from the outbox (returns null).
      */
-    fun pendingRelaysFor(eventId: HexKey): Set<NormalizedRelayUrl>? = eventOutbox[eventId]?.relaysLeft()
+    fun pendingRelaysFor(eventId: HexKey): Set<NormalizedRelayUrl>? = eventOutbox.get(eventId)?.relaysLeft()
 
     fun markAsSending(
         event: Event,
         relays: Set<NormalizedRelayUrl>,
     ): Set<NormalizedRelayUrl> {
-        val currentOutbox = eventOutbox[event.id]
+        val currentOutbox = eventOutbox.get(event.id)
         if (currentOutbox == null) {
-            eventOutbox = eventOutbox + Pair(event.id, PoolEventOutboxState(event, relays))
+            eventOutbox.put(event.id, PoolEventOutboxState(event, relays))
         } else {
             currentOutbox.updateRelays(relays)
         }
-        updateRelays()
-        return eventOutbox[event.id]?.remainingRelays() ?: emptySet()
+        // Additions only, and only what is genuinely new: the union is over
+        // this event's relays, never over the whole outbox.
+        addRelays(relays)
+        return eventOutbox.get(event.id)?.remainingRelays() ?: emptySet()
+    }
+
+    /** Union [wanted] into [relays], touching the flow only when it actually changes. */
+    private fun addRelays(wanted: Set<NormalizedRelayUrl>) {
+        val missing = wanted - relays.value
+        if (missing.isNotEmpty()) relays.update { it + missing }
+    }
+
+    /**
+     * An entry left the outbox. Retiring its relays needs a full scan, so that
+     * is amortised across [SWEEP_EVERY] removals — and always run once the
+     * outbox empties, which is the case that must not linger.
+     */
+    private fun onRemoved() {
+        pendingSweep++
+        if (pendingSweep >= SWEEP_EVERY || eventOutbox.isEmpty()) {
+            pendingSweep = 0
+            updateRelays()
+        }
     }
 
     /** Records a send attempt. Returns the event if this attempt exhausted its retry budget for
@@ -139,11 +189,11 @@ class PoolEventOutbox {
         id: HexKey,
         url: NormalizedRelayUrl,
     ): Event? {
-        val waiting = eventOutbox[id] ?: return null
+        val waiting = eventOutbox.get(id) ?: return null
         val gaveUp = waiting.newTry(url)
         if (waiting.isDone()) {
-            eventOutbox = eventOutbox - waiting.event.id
-            updateRelays()
+            eventOutbox.remove(waiting.event.id)
+            onRemoved()
         }
         return if (gaveUp) waiting.event else null
     }
@@ -154,12 +204,12 @@ class PoolEventOutbox {
         success: Boolean,
         message: String,
     ) {
-        val waiting = eventOutbox[id]
+        val waiting = eventOutbox.get(id)
         if (waiting != null) {
             waiting.newResponse(url, success, message)
             if (waiting.isDone()) {
-                eventOutbox = eventOutbox - waiting.event.id
-                updateRelays()
+                eventOutbox.remove(waiting.event.id)
+                onRemoved()
             }
         }
     }
@@ -171,8 +221,8 @@ class PoolEventOutbox {
         relay: NormalizedRelayUrl,
         sync: (Command) -> Unit,
     ) {
-        eventOutbox.forEach {
-            it.value.forEachUnsentEvent(relay) {
+        eventOutbox.forEach { _, outboxCache ->
+            outboxCache.forEachUnsentEvent(relay) {
                 sync(EventCmd(it))
             }
         }
@@ -210,15 +260,16 @@ class PoolEventOutbox {
         relay: NormalizedRelayUrl,
         errorMessage: String,
     ) {
-        eventOutbox.forEach {
-            if (relay in it.value.relaysRemaining) {
-                newResponse(it.key, relay, false, errorMessage)
+        eventOutbox.forEach { id, outboxCache ->
+            if (relay in outboxCache.relaysRemaining) {
+                newResponse(id, relay, false, errorMessage)
             }
         }
     }
 
     fun destroy() {
-        eventOutbox = emptyMap()
+        eventOutbox.clear()
+        pendingSweep = 0
         relays.tryEmit(emptySet())
     }
 }

quartz/src/commonTest/kotlin/com/vitorpamplona/quartz/nip01Core/relay/client/pool/PoolEventOutboxScaleTest.kt

@@ -0,0 +1,121 @@
+/*
+ * 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.nip01Core.relay.client.pool
+
+import com.vitorpamplona.quartz.nip01Core.core.Event
+import com.vitorpamplona.quartz.nip01Core.relay.normalizer.NormalizedRelayUrl
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+import kotlin.time.TimeSource
+
+/**
+ * The outbox must not get slower as it fills.
+ *
+ * It used to: the map was immutable and every `markAsSending` rebuilt it with
+ * `eventOutbox + Pair(...)`, so publishing N events copied 1 + 2 + … + N
+ * entries. Measured on a real bulk push at ~970k entries resident, that was
+ * ~20.5ms of the 22.7ms each event cost, and the rate visibly decayed as the
+ * backlog grew (45.6 -> 44.6 -> 43.2 ev/s over three windows). The relay-set
+ * bookkeeping was the other half — two full scans of every entry, per publish.
+ *
+ * This asserts the SHAPE of the cost rather than a wall-clock budget: a
+ * quadratic makes the second half of a run dramatically slower than the first,
+ * whatever the machine. A constant factor cannot be pinned in a unit test, but
+ * a growth curve can.
+ */
+class PoolEventOutboxScaleTest {
+    private val relay = NormalizedRelayUrl("wss://scale.relay.test")
+
+    private fun event(i: Int) =
+        Event(
+            id = i.toString(16).padStart(64, '0'),
+            pubKey = "00".repeat(32),
+            createdAt = 1_700_000_000L,
+            kind = 1,
+            tags = emptyArray(),
+            content = "hello",
+            sig = "00".repeat(64),
+        )
+
+    @Test
+    fun `publishing stays flat as the outbox fills`() {
+        val outbox = PoolEventOutbox()
+        val relays = setOf(relay)
+        val clock = TimeSource.Monotonic
+        val sample = 2_000
+        val total = 60_000
+
+        fun publishRange(
+            from: Int,
+            until: Int,
+        ) {
+            for (i in from until until) outbox.markAsSending(event(i), relays)
+        }
+
+        // Equal-sized windows at the START and the END of a long run. Halves
+        // would not do: over 20k publishes the average backlog only grows from
+        // ~7k to ~17k, a 2.4x expected ratio that hides inside JIT noise. Here
+        // the late window carries ~29x the backlog of the early one, so a
+        // per-entry cost shows up as a per-entry cost.
+        repeat(sample) { outbox.markAsSending(event(it), relays) } // warm up
+        val early =
+            clock.markNow().let { start ->
+                publishRange(sample, sample * 2)
+                start.elapsedNow()
+            }
+        publishRange(sample * 2, total - sample)
+        val late =
+            clock.markNow().let { start ->
+                publishRange(total - sample, total)
+                start.elapsedNow()
+            }
+
+        assertEquals(total, outbox.activeOutboxCacheFor(relay).size, "every publish is tracked")
+
+        val ratio = late.inWholeMicroseconds.toDouble() / early.inWholeMicroseconds.coerceAtLeast(1)
+        assertTrue(
+            ratio < 5.0,
+            "cost per publish must not grow with the backlog: first $sample took ${early.inWholeMilliseconds}ms at " +
+                "~$sample entries, last $sample took ${late.inWholeMilliseconds}ms at ~$total entries (ratio $ratio)",
+        )
+    }
+
+    @Test
+    fun `the relay set still reflects what is pending`() {
+        val outbox = PoolEventOutbox()
+        val a = NormalizedRelayUrl("wss://a.relay.test")
+        val b = NormalizedRelayUrl("wss://b.relay.test")
+
+        outbox.markAsSending(event(1), setOf(a))
+        assertEquals(setOf(a), outbox.relays.value, "a publish adds its relay immediately")
+
+        outbox.markAsSending(event(2), setOf(b))
+        assertEquals(setOf(a, b), outbox.relays.value, "a second relay joins without a rebuild")
+
+        // Draining every entry must clear the set — the sweep is batched, but
+        // emptying the outbox forces it, so a finished push does not strand a
+        // connection open forever.
+        outbox.newResponse(event(1).id, a, true, "")
+        outbox.newResponse(event(2).id, b, true, "")
+        assertEquals(emptySet(), outbox.relays.value, "an empty outbox wants no relays")
+    }
+}