Optimize QWP symbol dictionary recovery

Author: glasstiger

core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpWebSocketSender.java

@@ -3942,7 +3942,7 @@ private void resetSymbolDictStateForNewConnection() {
      * <ol>
      *   <li>the slot's persisted {@code .symbol-dict} -- its intact prefix; then</li>
      *   <li>the surviving frames' OWN delta sections, for every id above that prefix
-     *       ({@link CursorSendEngine#collectReplaySymbolsAbove}).</li>
+     *       ({@link CursorSendEngine#addRecoveredSymbolsTo}).</li>
      * </ol>
      * Those are exactly the two sources, in exactly the order, that the send loop's mirror
      * is built from: its constructor seeds {@code sentDictCount} from the same dictionary,
@@ -3976,7 +3976,7 @@ private void resetSymbolDictStateForNewConnection() {
      * <b>What still fails clean.</b> A genuine GAP: the ids below a surviving frame's delta
      * start were introduced by frames that were acked and TRIMMED away, so they lived only in
      * the lost dictionary and nothing can rebuild them.
-     * {@code collectReplaySymbolsAbove} returns -1 for that and we throw. It is the same
+     * {@code addRecoveredSymbolsTo} returns -1 for that and we throw. It is the same
      * condition the send loop's replay guard ({@code deltaStart > sentDictCount}) trips on, so
      * producer and drainer now agree on exactly which slots are recoverable, instead of the
      * producer rejecting slots the drainer drains.
@@ -3990,17 +3990,13 @@ private void seedGlobalDictionaryFromPersisted(PersistedSymbolDict pd) {
         //    mirror also seeds sentDictCount from.
         int baseline = 0;
         if (pd != null && pd.size() > 0) {
-            ObjList<String> persisted = pd.readLoadedSymbols();
-            for (int i = 0, n = persisted.size(); i < n; i++) {
-                globalSymbolDictionary.addRecoveredSymbol(persisted.getQuick(i));
-            }
+            pd.addLoadedSymbolsTo(globalSymbolDictionary);
             baseline = globalSymbolDictionary.size();
         }
         // 2. Everything the surviving frames define above that prefix, straight out of their
         //    own delta sections -- the same bytes, in the same order, accumulateSentDict will
         //    feed the mirror as those frames go back on the wire.
-        ObjList<String> fromFrames = new ObjList<>();
-        long coverage = cursorEngine.collectReplaySymbolsAbove(baseline, fromFrames);
+        long coverage = cursorEngine.addRecoveredSymbolsTo(baseline, globalSymbolDictionary);
         if (coverage < 0) {
             // A gap: the surviving frames reference ids below their own delta start,
             // introduced by frames since acked and trimmed away, and the persisted
@@ -4013,8 +4009,8 @@ private void seedGlobalDictionaryFromPersisted(PersistedSymbolDict pd) {
             // "resend required" for delivered data AND -- because such a slot is fully
             // drained -- let build()'s connect-path close unlink the (already-delivered)
             // bytes the quarantine claims to preserve. So DON'T throw: seed the intact
-            // prefix only, leaving fromFrames unused exactly as the send loop's mirror
-            // does on a -1 (it adds nothing), so the producer baseline and the mirror's
+            // prefix only; addRecoveredSymbolsTo adds nothing on a -1 exactly as the
+            // send loop's mirror does, so the producer baseline and the mirror's
             // sentDictCount still agree by construction. The producer resumes above the
             // prefix and the fully-drained slot is cleaned up on close.
             long ackedFsn = cursorEngine.ackedFsn();
@@ -4040,9 +4036,6 @@ private void seedGlobalDictionaryFromPersisted(PersistedSymbolDict pd) {
                             + "holds them; the recovered dictionary holds only "
                             + (pd == null ? 0 : pd.size()) + " id(s) -- resend the affected data");
         }
-        for (int i = 0, n = fromFrames.size(); i < n; i++) {
-            globalSymbolDictionary.addRecoveredSymbol(fromFrames.getQuick(i));
-        }
         // Producer baseline == the coverage the replay will establish == the mirror's
         // sentDictCount once those frames have gone out. The first new frame therefore
         // starts its delta exactly at the tip, and the replay guard passes.

core/src/main/java/io/questdb/client/cutlass/qwp/client/sf/cursor/CursorSendEngine.java

@@ -24,6 +24,7 @@
 
 package io.questdb.client.cutlass.qwp.client.sf.cursor;
 
+import io.questdb.client.cutlass.qwp.client.GlobalSymbolDictionary;
 import io.questdb.client.cutlass.qwp.protocol.QwpConstants;
 import io.questdb.client.std.Compat;
 import io.questdb.client.std.Files;
@@ -760,6 +761,24 @@ public long collectReplaySymbolsAbove(int baseline, ObjList<String> out) {
         );
     }
 
+    /**
+     * Decodes the cached recovery suffix directly into the producer's global
+     * dictionary. Recovery always builds the analysis with the persisted
+     * prefix size as its baseline, so no intermediate cardinality-sized list is
+     * needed on the production path.
+     */
+    public long addRecoveredSymbolsTo(int baseline, GlobalSymbolDictionary target) {
+        if (recoveredFrameAnalysis == null) {
+            return baseline;
+        }
+        RecoveredFrameAnalysis analysis = checkedRecoveryAnalysis(baseline);
+        long coverage = analysis.coverage();
+        if (coverage >= 0L) {
+            analysis.addDecodedSymbolsTo(target);
+        }
+        return coverage;
+    }
+
     long recoveredSymbolCoverage(int baseline) {
         return checkedRecoveryAnalysis(baseline).coverage();
     }
@@ -785,6 +804,11 @@ public long recoveryFramesVisited() {
         return recoveredFrameAnalysis == null ? 0L : recoveredFrameAnalysis.framesVisited();
     }
 
+    @TestOnly
+    public long recoverySymbolEntriesVisited() {
+        return recoveredFrameAnalysis == null ? 0L : recoveredFrameAnalysis.symbolEntriesVisited();
+    }
+
     private RecoveredFrameAnalysis checkedRecoveryAnalysis(int baseline) {
         if (recoveredFrameAnalysis == null || recoveredFrameAnalysis.baseline() != baseline) {
             throw new IllegalStateException("recovery symbol baseline mismatch [expected="

core/src/main/java/io/questdb/client/cutlass/qwp/client/sf/cursor/CursorWebSocketSendLoop.java

@@ -299,6 +299,9 @@ public final class CursorWebSocketSendLoop implements QuietCloseable {
     private long sentDictBytesAddr;
     private int sentDictBytesCapacity;
     private int sentDictBytesLen;
+    // False only while an orphan-drainer loop borrows the persisted prefix.
+    // Any growth first performs a copy-on-write; only owned buffers are freed.
+    private boolean sentDictBytesOwned;
     private int sentDictCount;
     // Reusable staging frame for dictionary catch-up chunks. A reconnect may split a
     // large dictionary into many chunks, and every later reconnect repeats that split;
@@ -721,24 +724,17 @@ public CursorWebSocketSendLoop(WebSocketClient client, CursorSendEngine engine,
         // tracks. When the dictionary could not be opened, pd is null, the mirror
         // simply starts empty and grows from the frames themselves.
         PersistedSymbolDict pd = engine.getPersistedSymbolDict();
+        int persistedPrefixLen = 0;
         if (pd != null && pd.size() > 0) {
             int len = pd.loadedEntriesLen();
             if (len > 0) {
-                // COPY the persisted dictionary's loaded-entries buffer into this
-                // loop's own mirror rather than taking ownership of it. The engine
-                // (and its PersistedSymbolDict) OUTLIVES this loop on the orphan
-                // drainer path: BackgroundDrainer builds a fresh send loop per wire
-                // session against the same engine on a durable-ack capability-gap
-                // recycle. A one-shot ownership transfer would leave every loop
-                // after the first with an EMPTY mirror -- it would then send no
-                // reconnect catch-up, and the first replayed delta frame
-                // (deltaStart > 0) would trip the torn-dict guard, falsely
-                // quarantining a healthy slot. Copying keeps the dictionary's
-                // loaded entries intact for the engine's lifetime so every
-                // recycled loop re-seeds; pd.close() (at engine close) frees the
-                // dictionary's copy, this loop frees its own copy on exit.
-                sentDictBytesAddr = Unsafe.malloc(len, MemoryTag.NATIVE_DEFAULT);
-                Unsafe.getUnsafe().copyMemory(pd.loadedEntriesAddr(), sentDictBytesAddr, len);
+                // Seed by reference. The foreground loop takes ownership after
+                // construction succeeds; orphan-drainer loops keep borrowing because
+                // one engine can create several sessions during capability-gap
+                // recycling. A borrowed mirror performs copy-on-write if a recovered
+                // frame suffix must extend it.
+                persistedPrefixLen = len;
+                sentDictBytesAddr = pd.loadedEntriesAddr();
                 sentDictBytesCapacity = len;
                 sentDictBytesLen = len;
                 // Set the count only alongside the bytes so sentDictCount can
@@ -786,16 +782,25 @@ public CursorWebSocketSendLoop(WebSocketClient client, CursorSendEngine engine,
                     }
                 }
             } catch (Throwable t) {
-                if (sentDictBytesAddr != 0) {
-                    Unsafe.free(sentDictBytesAddr, sentDictBytesCapacity, MemoryTag.NATIVE_DEFAULT);
-                    sentDictBytesAddr = 0;
-                    sentDictBytesCapacity = 0;
-                    sentDictBytesLen = 0;
-                    sentDictCount = 0;
-                }
+                releaseSentDictBytes();
                 throw t;
             }
         }
+        // QwpWebSocketSender seeds the producer dictionary before constructing
+        // its one foreground loop, so that loop can own the recovered prefix and
+        // eliminate the engine-side duplicate. BackgroundDrainer must retain the
+        // prefix for later recycled loops and therefore keeps borrowing it.
+        if (persistedPrefixLen > 0 && catchUpCapGapPolicy == CatchUpCapGapPolicy.RETRY_FOREVER) {
+            long persistedPrefixAddr = pd.takeLoadedEntries();
+            if (sentDictBytesOwned) {
+                // Copy-on-grow already produced the combined mirror. Retire the
+                // no-longer-needed persisted prefix after successful construction.
+                Unsafe.free(persistedPrefixAddr, persistedPrefixLen, MemoryTag.NATIVE_DEFAULT);
+            } else {
+                assert persistedPrefixAddr == sentDictBytesAddr;
+                sentDictBytesOwned = true;
+            }
+        }
         this.fsnAtZero = fsnAtZero;
         this.parkNanos = parkNanos;
         this.reconnectFactory = reconnectFactory;
@@ -1131,15 +1136,7 @@ public synchronized void close() {
         // thread may still be mid-send, so touching the mirror here would race.
         // A duplicate close observes sentDictBytesAddr == 0 and skips.
         if (loopNeverRan && sentDictBytesAddr != 0) {
-            Unsafe.free(sentDictBytesAddr, sentDictBytesCapacity, MemoryTag.NATIVE_DEFAULT);
-            sentDictBytesAddr = 0;
-            sentDictBytesCapacity = 0;
-            sentDictBytesLen = 0;
-            // Reset the count alongside the buffer so the mirror stays all-or-
-            // nothing: a hypothetical close()-then-start() (start() has no closed
-            // guard) must not observe a non-zero sentDictCount against a freed
-            // buffer and drive setWireBaselineWithCatchUp into a null-mirror catch-up.
-            sentDictCount = 0;
+            releaseSentDictBytes();
         }
         if (loopNeverRan) {
             freeCatchUpFrameBuffer();
@@ -2048,11 +2045,7 @@ private void ioLoop() {
             // The symbol-dict mirror is I/O-thread-owned; free it here, on the
             // owning thread's exit path, after the last send that could touch it.
             if (sentDictBytesAddr != 0) {
-                Unsafe.free(sentDictBytesAddr, sentDictBytesCapacity, MemoryTag.NATIVE_DEFAULT);
-                sentDictBytesAddr = 0;
-                sentDictBytesCapacity = 0;
-                sentDictBytesLen = 0;
-                sentDictCount = 0; // keep the mirror all-or-nothing (see close())
+                releaseSentDictBytes();
             }
             freeCatchUpFrameBuffer();
             shutdownLatch.countDown();
@@ -2438,10 +2431,34 @@ private void ensureSentDictCapacity(long required) {
         if (newCap > MAX_SENT_DICT_BYTES) {
             newCap = MAX_SENT_DICT_BYTES;
         }
-        sentDictBytesAddr = Unsafe.realloc(sentDictBytesAddr, sentDictBytesCapacity, (int) newCap, MemoryTag.NATIVE_DEFAULT);
+        if (sentDictBytesOwned) {
+            sentDictBytesAddr = Unsafe.realloc(
+                    sentDictBytesAddr,
+                    sentDictBytesCapacity,
+                    (int) newCap,
+                    MemoryTag.NATIVE_DEFAULT);
+        } else {
+            long newAddr = Unsafe.malloc((int) newCap, MemoryTag.NATIVE_DEFAULT);
+            if (sentDictBytesLen > 0) {
+                Unsafe.getUnsafe().copyMemory(sentDictBytesAddr, newAddr, sentDictBytesLen);
+            }
+            sentDictBytesAddr = newAddr;
+            sentDictBytesOwned = true;
+        }
         sentDictBytesCapacity = (int) newCap;
     }
 
+    private void releaseSentDictBytes() {
+        if (sentDictBytesAddr != 0 && sentDictBytesOwned) {
+            Unsafe.free(sentDictBytesAddr, sentDictBytesCapacity, MemoryTag.NATIVE_DEFAULT);
+        }
+        sentDictBytesAddr = 0;
+        sentDictBytesCapacity = 0;
+        sentDictBytesLen = 0;
+        sentDictBytesOwned = false;
+        sentDictCount = 0;
+    }
+
     private long readVarintAt(long p, long limit) {
         long value = 0;
         int shift = 0;

core/src/main/java/io/questdb/client/cutlass/qwp/client/sf/cursor/MmapSegment.java

@@ -364,8 +364,9 @@ public static MmapSegment openExisting(FilesFacade ff, String path) {
                 throw new MmapSegmentException(
                         "bad baseSeq in " + path + ": " + baseSeq);
             }
-            long lastGood = scanFrames(addr, fileSize);
-            long count = countFrames(addr, lastGood);
+            FrameScan scan = scanFrames(addr, fileSize);
+            long lastGood = scan.lastGood;
+            long count = scan.frameCount;
             long tornTail = detectTornTail(addr, lastGood, fileSize);
             if (tornTail > 0) {
                 LOG.warn("SF segment {}: torn tail of {} bytes at offset {} "
@@ -907,22 +908,23 @@ private static boolean isMmapAccessFault(Throwable t) {
     }
 
     /**
-     * Forward scan that returns the offset just past the last frame whose
-     * CRC verifies. A torn-tail frame (declared length runs past EOF, or
-     * CRC mismatch) leaves both cursors at the start of that frame; the
-     * next {@link #tryAppend} will overwrite it. The scan only reads from
-     * the mapping — no syscalls.
+     * Forward scan that returns the offset just past the last frame whose CRC
+     * verifies and the number of verified frames. A torn-tail frame (declared
+     * length runs past EOF, or CRC mismatch) leaves both cursors at the start of
+     * that frame; the next {@link #tryAppend} will overwrite it. The scan only
+     * reads from the mapping — no syscalls.
      */
-    private static long scanFrames(long addr, long fileSize) {
+    private static FrameScan scanFrames(long addr, long fileSize) {
         long pos = HEADER_SIZE;
+        long frameCount = 0L;
         try {
             while (pos + FRAME_HEADER_SIZE <= fileSize) {
                 int crcRead = Unsafe.getUnsafe().getInt(addr + pos);
                 int payloadLen = Unsafe.getUnsafe().getInt(addr + pos + 4);
                 // Defensive: a corrupt length field could be enormous or negative,
                 // both of which would otherwise overrun the mapping.
                 if (payloadLen < 0 || pos + FRAME_HEADER_SIZE + payloadLen > fileSize) {
-                    return pos;
+                    return new FrameScan(pos, frameCount);
                 }
                 // CRC over the contiguous (payloadLen, payload) pair, folded
                 // via Unsafe reads rather than the native Crc32c.update.
@@ -945,9 +947,10 @@ private static long scanFrames(long addr, long fileSize) {
                 // table CRC here is immaterial.
                 int crcCalc = crc32cRecovery(addr + pos + 4, 4L + payloadLen);
                 if (crcCalc != crcRead) {
-                    return pos;
+                    return new FrameScan(pos, frameCount);
                 }
                 pos += FRAME_HEADER_SIZE + payloadLen;
+                frameCount++;
             }
         } catch (InternalError e) {
             // The read at `pos` hit a mapped page that is not backed by real
@@ -973,7 +976,7 @@ private static long scanFrames(long addr, long fileSize) {
                             + "if this recurs.",
                     pos, fileSize);
         }
-        return pos;
+        return new FrameScan(pos, frameCount);
     }
 
     /**
@@ -1049,19 +1052,13 @@ private static long detectTornTail(long addr, long lastGood, long fileSize) {
         return 0L;
     }
 
-    /**
-     * Counts frames in {@code [HEADER_SIZE, lastGood)}. Walks the framing in
-     * lockstep with {@link #scanFrames} (which already validated CRCs); so
-     * this is just length-driven traversal, no CRC re-check.
-     */
-    private static long countFrames(long addr, long lastGood) {
-        long pos = HEADER_SIZE;
-        long count = 0;
-        while (pos < lastGood) {
-            int payloadLen = Unsafe.getUnsafe().getInt(addr + pos + 4);
-            pos += FRAME_HEADER_SIZE + payloadLen;
-            count++;
+    private static final class FrameScan {
+        private final long frameCount;
+        private final long lastGood;
+
+        private FrameScan(long lastGood, long frameCount) {
+            this.lastGood = lastGood;
+            this.frameCount = frameCount;
         }
-        return count;
     }
 }