fix(qwp): drop cached timestamp columns and repair accounting on buffer clear/rollback

Senders cache designated-timestamp ColumnBuffer refs and derive
bookkeeping from their table buffers. rollbackUncommittedColumns()
(failed at()) and QwpTableBuffer.clear() close those columns, leaving
dangling refs that NPE on the next at() and corrupt pending accounting.

- Invalidate WebSocket timestamp caches in rollbackRow(); cancelRow()
  now routes through it.
- Add QwpTableBuffer.Owner: clear() notifies the owning sender before
  releasing columns. WebSocket subtracts the buffer's exact accounted
  share (committedBytes - new baselineBytes, seed offsets excluded);
  UDP resets transient row state, datagram estimate, and headroom
  state (its base-estimate cache is keyed by column count alone).
- Unify flush snapshot anchoring at post-reset storage so per-table
  contributions are exact under any reset/switch/flush history.

Hot path unchanged. Regressions: clear between rows (both units, both
senders), non-current/mid-row clear, accounting locality, stale
headroom estimate.

Author: jerrinot

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

@@ -67,7 +67,7 @@
  * state until commit, so committed table data can be flushed without replaying
  * the row back into column storage.
  */
-public class QwpUdpSender implements Sender {
+public class QwpUdpSender implements Sender, QwpTableBuffer.Owner {
     private static final int ADAPTIVE_HEADROOM_EWMA_SHIFT = 2;
     private static final Logger LOG = LoggerFactory.getLogger(QwpUdpSender.class);
     private static final int MAX_TABLE_NAME_LENGTH = 127;
@@ -195,6 +195,27 @@ public void cancelRow() {
         rollbackCurrentRowToCommittedState();
     }
 
+    /**
+     * {@link QwpTableBuffer.Owner} callback: the buffer is about to close all
+     * of its column buffers, so drop cached column references, in-progress
+     * column states, and the committed datagram estimate derived from it.
+     * The table's adaptive headroom state must go too: its base-estimate
+     * cache and EWMA samples are keyed by column count alone, which is only
+     * sound while columns are add-only — clear() can replace the schema with
+     * a different one of the same size.
+     */
+    @Override
+    public void onTableBufferClear(QwpTableBuffer buffer) {
+        if (buffer == currentTableBuffer) {
+            clearTransientRowState();
+            resetCommittedDatagramEstimate();
+        }
+        TableHeadroomState headroomState = tableHeadroomStates.get(buffer.getTableName());
+        if (headroomState != null) {
+            headroomState.reset();
+        }
+    }
+
     @Override
     public void close() {
         if (!closed) {
@@ -580,7 +601,7 @@ public Sender table(CharSequence tableName) {
             currentTableName = currentTableBuffer.getTableName();
         } else {
             currentTableName = tableName.toString();
-            currentTableBuffer = new QwpTableBuffer(currentTableName);
+            currentTableBuffer = new QwpTableBuffer(currentTableName, this);
             tableBuffers.put(currentTableName, currentTableBuffer);
         }
         currentTableHeadroomState = tableHeadroomStates.get(currentTableName);
@@ -1490,6 +1511,15 @@ long predictNextRowGrowth() {
             return Math.max(lastRowDatagramGrowth, ewmaRowDatagramGrowth);
         }
 
+        void reset() {
+            cachedBaseEstimate = -1;
+            cachedBaseEstimateColumnCount = -1;
+            committedSampleCount = 0;
+            ewmaRowDatagramGrowth = 0;
+            lastRowDatagramGrowth = 0;
+            schemaColumnCount = -1;
+        }
+
         void recordCommittedRow(int currentSchemaColumnCount, long rowDatagramGrowth) {
             if (schemaColumnCount != currentSchemaColumnCount) {
                 committedSampleCount = 0;

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

@@ -120,7 +120,7 @@
  * Initial connection failures are not retained as terminal sender state; a later
  * operation may try to connect again.
  */
-public class QwpWebSocketSender implements Sender {
+public class QwpWebSocketSender implements Sender, QwpTableBuffer.Owner {
 
     public static final long DEFAULT_AUTH_TIMEOUT_MS = 15_000L;
     // Soft per-batch byte budget. Trips a flush when raw column-buffer bytes
@@ -1011,10 +1011,7 @@ public QwpWebSocketSender byteColumn(CharSequence columnName, byte value) {
     @Override
     public void cancelRow() {
         checkNotClosed();
-        if (currentTableBuffer != null) {
-            currentTableBuffer.cancelCurrentRow();
-            currentTableBuffer.rollbackUncommittedColumns();
-        }
+        rollbackRow();
     }
 
     /**
@@ -1811,6 +1808,15 @@ public long getPendingBytes() {
         return pendingBytes;
     }
 
+    /**
+     * Row counterpart of {@link #getPendingBytes()}: committed-but-unflushed
+     * rows across all tables, as tracked for the auto-flush row threshold.
+     */
+    @TestOnly
+    public int getPendingRowCount() {
+        return pendingRowCount;
+    }
+
     /**
      * Server-advertised cap on the per-batch raw byte size. Zero before the
      * first connect; updated by every successful reconnect via
@@ -1824,10 +1830,18 @@ public int getServerMaxBatchSize() {
     @TestOnly
     public QwpTableBuffer getTableBuffer(String tableName) {
         QwpTableBuffer buffer = tableBuffers.get(tableName);
+        if (buffer == currentTableBuffer && buffer != null) {
+            // Keep the timestamp caches and byte snapshot intact: re-snapping
+            // mid-row would fold in-progress bytes into the snapshot and break
+            // sendRow()'s delta accounting.
+            return buffer;
+        }
         if (buffer == null) {
             buffer = new QwpTableBuffer(tableName, this);
             tableBuffers.put(tableName, buffer);
         }
+        cachedTimestampColumn = null;
+        cachedTimestampNanosColumn = null;
         currentTableBuffer = buffer;
         currentTableBufferSnapshotBytes = buffer.getBufferedBytes();
         currentTableName = tableName;
@@ -2235,6 +2249,33 @@ public void reset() {
         cachedTimestampNanosColumn = null;
     }
 
+    /**
+     * {@link QwpTableBuffer.Owner} callback: the buffer is about to close all
+     * of its column buffers, so drop any cached column references and remove
+     * exactly the buffer's accounted share from the pending totals. That
+     * share is its committed bytes (the snapshot for the current table, which
+     * excludes any in-progress row; raw storage otherwise) minus its baseline
+     * bytes, which pendingBytes never included. Runs before the buffer wipes
+     * itself, so all three values still reflect the pre-clear state.
+     */
+    @Override
+    public void onTableBufferClear(QwpTableBuffer buffer) {
+        long committedBytes;
+        if (buffer == currentTableBuffer) {
+            cachedTimestampColumn = null;
+            cachedTimestampNanosColumn = null;
+            committedBytes = currentTableBufferSnapshotBytes;
+            currentTableBufferSnapshotBytes = 0;
+        } else {
+            committedBytes = buffer.getBufferedBytes();
+        }
+        pendingBytes -= committedBytes - buffer.getBaselineBytes();
+        pendingRowCount -= buffer.getRowCount();
+        if (pendingRowCount == 0) {
+            firstPendingRowTimeNanos = 0;
+        }
+    }
+
     /**
      * Register an async listener for connection-state transitions: initial
      * connect, primary failover, endpoint attempt failures, the full address
@@ -3823,7 +3864,11 @@ private void resetTableBuffersAfterFlush(ObjList<CharSequence> keys) {
         }
         currentBatchMaxSymbolId = -1;
         pendingBytes = 0;
-        currentTableBufferSnapshotBytes = 0;
+        // Anchor at the post-reset storage (baseline seed bytes of var-width
+        // columns), matching the empty-flush path above, so per-row deltas
+        // accumulated from here always equal committedBytes - baselineBytes.
+        currentTableBufferSnapshotBytes = currentTableBuffer == null
+                ? 0 : currentTableBuffer.getBufferedBytes();
         pendingRowCount = 0;
         firstPendingRowTimeNanos = 0;
     }
@@ -3860,6 +3905,10 @@ private void resetSymbolDictStateForNewConnection() {
     }
 
     private void rollbackRow() {
+        // rollbackUncommittedColumns() may close a newly created designated
+        // timestamp column, so cached references must not survive rollback.
+        cachedTimestampColumn = null;
+        cachedTimestampNanosColumn = null;
         if (currentTableBuffer != null) {
             currentTableBuffer.cancelCurrentRow();
             currentTableBuffer.rollbackUncommittedColumns();

core/src/main/java/io/questdb/client/cutlass/qwp/protocol/QwpTableBuffer.java

@@ -63,8 +63,10 @@ public class QwpTableBuffer implements QuietCloseable {
     private static final int MAX_COLUMN_NAME_LENGTH = 127;
     private final LowerCaseCharSequenceIntHashMap columnNameToIndex;
     private final ObjList<ColumnBuffer> columns;
+    private final Owner owner;
     private final QwpWebSocketSender sender;
     private final String tableName;
+    private long baselineBytes; // storage retained by empty columns after reset() (e.g. var-width seed offsets)
     private QwpColumnDef[] cachedColumnDefs;
     private int columnAccessCursor; // tracks expected next column index
     private boolean columnDefsCacheValid;
@@ -74,18 +76,33 @@ public class QwpTableBuffer implements QuietCloseable {
     private int rowCount;
 
     public QwpTableBuffer(String tableName) {
-        this(tableName, null);
+        this(tableName, (QwpWebSocketSender) null);
     }
 
     /**
      * Use this constructor overload to allow writing to a symbol column.
      * {@link ColumnBuffer#addSymbol(CharSequence)} needs the sender to
      * call {@link QwpWebSocketSender#getOrAddGlobalSymbol(CharSequence)}, registering
      * the symbol in the global dictionary shared with the server.
+     * The sender is also registered as the buffer's {@link Owner}.
      */
     public QwpTableBuffer(String tableName, QwpWebSocketSender sender) {
+        this(tableName, sender, sender);
+    }
+
+    /**
+     * Use this constructor overload for owners that derive state from the
+     * buffer but do not use the global symbol dictionary (symbol columns fall
+     * back to the per-buffer dictionary).
+     */
+    public QwpTableBuffer(String tableName, Owner owner) {
+        this(tableName, null, owner);
+    }
+
+    private QwpTableBuffer(String tableName, QwpWebSocketSender sender, Owner owner) {
         this.tableName = tableName;
         this.sender = sender;
+        this.owner = owner;
         this.columns = new ObjList<>();
         this.columnNameToIndex = new LowerCaseCharSequenceIntHashMap();
         this.rowCount = 0;
@@ -118,8 +135,15 @@ public void cancelCurrentRow() {
     /**
      * Clears the buffer completely, including column definitions.
      * Frees all off-heap memory.
+     * <p>
+     * The registered {@link Owner} is notified before any column state is
+     * released, so it can drop cached {@link ColumnBuffer} references and
+     * repair bookkeeping derived from the buffer's pre-clear contents.
      */
     public void clear() {
+        if (owner != null) {
+            owner.onTableBufferClear(this);
+        }
         for (int i = 0, n = columns.size(); i < n; i++) {
             columns.get(i).close();
         }
@@ -132,13 +156,26 @@ public void clear() {
         rowCount = 0;
         columnDefsCacheValid = false;
         cachedColumnDefs = null;
+        baselineBytes = 0;
     }
 
     @Override
     public void close() {
         clear();
     }
 
+    /**
+     * Bytes of storage the columns retained through the last {@link #reset()}
+     * (var-width columns re-seed a 4-byte initial offset entry, for example).
+     * These bytes are part of {@link #getBufferedBytes()} but predate any row
+     * committed since the reset, so an owner accounting for committed-row
+     * bytes must exclude them: contribution = committed bytes - baseline.
+     * Zero for a freshly created or {@link #clear()}ed buffer.
+     */
+    public long getBaselineBytes() {
+        return baselineBytes;
+    }
+
     /**
      * Returns the total bytes buffered across all columns.
      * This queries actual buffer sizes, not estimates.
@@ -312,6 +349,7 @@ public void reset() {
         committedColumnCount = columns.size();
         inProgressColumnCount = 0;
         rowCount = 0;
+        baselineBytes = getBufferedBytes();
     }
 
     public void retainInProgressRow(
@@ -489,6 +527,18 @@ static int elementSizeInBuffer(byte type) {
         }
     }
 
+    /**
+     * A sender that owns this buffer and derives state from it: cached
+     * {@link ColumnBuffer} references, pending-byte/row accounting, datagram
+     * size estimates. {@link #clear()} closes every column buffer at once,
+     * which invalidates all of that derived state, so it notifies the owner
+     * <em>before</em> releasing anything -- the callback may still read the
+     * buffer's pre-clear row count and byte sizes.
+     */
+    public interface Owner {
+        void onTableBufferClear(QwpTableBuffer buffer);
+    }
+
     /**
      * Helper class to capture array data from DoubleArray/LongArray.appendToBufPtr().
      */

core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpUdpSenderTest.java

@@ -269,6 +269,45 @@ public void testAtMicrosOversizeFailureRollsBackWithoutLeakingTimestampState() t
         });
     }
 
+    @Test
+    public void testAtMicrosRecoversAfterTableBufferClear() throws Exception {
+        assertUdpSenderRecoversAfterTableBufferClear(ChronoUnit.MICROS);
+    }
+
+    @Test
+    public void testAtNanosRecoversAfterTableBufferClear() throws Exception {
+        assertUdpSenderRecoversAfterTableBufferClear(ChronoUnit.NANOS);
+    }
+
+    @Test
+    public void testTableBufferClearInvalidatesHeadroomBaseEstimate() throws Exception {
+        assertMemoryLeak(() -> {
+            // Same column count before and after clear(), but a far longer
+            // column name: a stale base estimate cached for the old schema
+            // (keyed by column count alone) would under-estimate the new one.
+            String longName = repeat('n', 120);
+            CapturingNetworkFacade nf = new CapturingNetworkFacade();
+            try (QwpUdpSender sender = new QwpUdpSender(nf, 0, 0, 9000, 1, 1024 * 1024)) {
+                sender.table("t")
+                        .longColumn("a", 1)
+                        .atNow();
+                sender.currentTableBufferForTest().clear();
+
+                sender.table("t")
+                        .longColumn(longName, 2)
+                        .atNow();
+                long estimate = sender.committedDatagramEstimateForTest();
+                sender.flush();
+
+                int actual = nf.lengths.get(nf.lengths.size() - 1);
+                Assert.assertTrue(
+                        "estimate must cover the rebuilt schema [estimate=" + estimate + ", actual=" + actual + "]",
+                        estimate >= actual
+                );
+            }
+        });
+    }
+
     @Test
     public void testAtNanosOversizeFailureRollsBackWithoutLeakingTimestampState() throws Exception {
         assertMemoryLeak(() -> {
@@ -1955,6 +1994,32 @@ private static void assertEstimateAtLeastActual(List<ScenarioRow> rows) {
         }
     }
 
+    private static void assertUdpSenderRecoversAfterTableBufferClear(ChronoUnit unit) throws Exception {
+        assertMemoryLeak(() -> {
+            CapturingNetworkFacade nf = new CapturingNetworkFacade();
+            try (QwpUdpSender sender = new QwpUdpSender(nf, 0, 0, 9000, 1, 1024 * 1024)) {
+                sender.table("t")
+                        .longColumn("a", 1)
+                        .at(1_000_000L, unit);
+                Assert.assertTrue(sender.committedDatagramEstimateForTest() > 0);
+
+                // Raw clear() discards the buffered row; the sender must drop
+                // its cached designated-timestamp column and datagram estimate.
+                sender.currentTableBufferForTest().clear();
+                Assert.assertEquals(0, sender.committedDatagramEstimateForTest());
+
+                sender.table("t")
+                        .longColumn("a", 2)
+                        .at(2_000_000L, unit);
+                sender.flush();
+            }
+
+            assertRowsEqual(Collections.singletonList(
+                    decodedRow("t", "a", 2L, "", 2_000_000L)
+            ), decodeRows(nf.packets));
+        });
+    }
+
     private static void assertNullableArrayNullState(QwpTableBuffer.ColumnBuffer column) {
         Assert.assertNotNull(column);
         Assert.assertEquals(1, column.getSize());