remove estimator bandaid

Author: jerrinot

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

@@ -65,7 +65,6 @@
  */
 public class QwpUdpSender implements Sender {
     private static final int VARINT_INT_UPPER_BOUND = 5;
-    private static final int SAFETY_MARGIN_BYTES = 8;
     private static final Logger LOG = LoggerFactory.getLogger(QwpUdpSender.class);
 
     private final UdpLineChannel channel;
@@ -85,15 +84,24 @@ public class QwpUdpSender implements Sender {
     private int inProgressRowValueCount;
     private QwpTableBuffer currentTableBuffer;
     private String currentTableName;
+
+    // prefix* arrays: per-column snapshots captured before the in-progress row,
+    // used to encode and flush only the committed prefix when a row is still being built.
+    // Indexed by column index. -1 means the column has no in-progress data.
     private int[] prefixArrayDataOffsetBefore = new int[8];
     private int[] prefixArrayShapeOffsetBefore = new int[8];
     private long[] prefixStringDataSizeBefore = new long[8];
     private int[] prefixSymbolDictionarySizeBefore = new int[8];
     private int[] prefixSizeBefore = new int[8];
     private int[] prefixValueCountBefore = new int[8];
+
+    // columns that need NULL/default fill for the current row (columns not yet written to)
     private QwpTableBuffer.ColumnBuffer[] rowFillColumns = new QwpTableBuffer.ColumnBuffer[8];
+    // maps column index -> 1-based position in rowFillColumns (0 means absent)
     private int[] rowFillColumnPositions = new int[8];
+    // per-column marks to detect duplicate writes within a single row; compared against currentRowMark
     private int[] stagedColumnMarks = new int[8];
+    // monotonically increasing mark; incremented per row to invalidate stagedColumnMarks without clearing
     private int currentRowMark = 1;
     private int rowFillColumnCount;
     private int inProgressColumnCount;
@@ -845,7 +853,6 @@ private long estimateBaseForCurrentSchema() {
                 estimate += 1;
             }
         }
-        estimate += SAFETY_MARGIN_BYTES;
         return estimate;
     }
 
@@ -1085,6 +1092,11 @@ public QwpTableBuffer currentTableBufferForTest() {
         return currentTableBuffer;
     }
 
+    @TestOnly
+    public long committedDatagramEstimateForTest() {
+        return committedDatagramEstimate;
+    }
+
     private void captureInProgressColumnPrefixState() {
         int columnCount = currentTableBuffer.getColumnCount();
         ensurePrefixColumnCapacity(columnCount);
@@ -1299,6 +1311,11 @@ private static int utf8Length(CharSequence s) {
         return len;
     }
 
+    /**
+     * Captures the state of a column buffer at the moment the in-progress row starts
+     * writing to it. The snapshot allows the sender to compute incremental datagram
+     * size estimates and to roll back the column to its pre-row state on error or cancel.
+     */
     private static final class InProgressColumnState {
         private int arrayDataOffsetBefore;
         private int arrayShapeOffsetBefore;

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

@@ -139,6 +139,14 @@ public void testBoundedSenderMixedNullablePaddingPreservesRowsAndPacketLimit() t
         });
     }
 
+    @Test
+    public void testEstimateMatchesActualEncodedSize() throws Exception {
+        assertMemoryLeak(() -> {
+            auditEstimateWithStableSchemaAndNullableValues();
+            auditEstimateAcrossSymbolDictionaryVarintBoundary();
+        });
+    }
+
     @Test
     public void testBoundedSenderNullableStringNullAcrossOverflowBoundaryPreservesRowsAndPacketLimit() throws Exception {
         assertMemoryLeak(() -> {
@@ -1702,6 +1710,91 @@ private static DoubleArrayValue doubleArrayValue(int[] shape, double... values)
         return new DoubleArrayValue(dims, elems);
     }
 
+    private static void auditEstimateAcrossSymbolDictionaryVarintBoundary() throws Exception {
+        ArrayList<ScenarioRow> rows = new ArrayList<>();
+        for (int i = 0; i < 160; i++) {
+            final int rowId = i;
+            rows.add(row(
+                    "sym_audit",
+                    sender -> sender.table("sym_audit")
+                            .longColumn("x", rowId)
+                            .symbol("sym", "sym-" + rowId)
+                            .atNow(),
+                    "x", (long) rowId,
+                    "sym", "sym-" + rowId
+            ));
+        }
+        assertEstimateAtLeastActual(rows);
+    }
+
+    private static void auditEstimateWithStableSchemaAndNullableValues() throws Exception {
+        ArrayList<ScenarioRow> rows = new ArrayList<>();
+        for (int i = 0; i < 96; i++) {
+            final int rowId = i;
+            final String stringValue = (i & 1) == 0 ? "tokyo-" + i + "-" + repeat('x', (i % 31) + 1) : null;
+            final long[] longArray = i % 3 == 0 ? new long[]{i, i + 1L, i + 2L} : null;
+            final double[][] doubleArray = i % 5 == 0 ? new double[][]{{i + 0.5, i + 1.5}, {i + 2.5, i + 3.5}} : null;
+            final Decimal64 decimal64 = i % 7 == 0 ? Decimal64.fromLong(i * 100L + 7, 2) : null;
+            final Decimal128 decimal128 = i % 11 == 0 ? Decimal128.fromLong(i * 1000L + 11, 4) : null;
+            final Decimal256 decimal256 = i % 13 == 0 ? Decimal256.fromLong(i * 10000L + 13, 3) : null;
+
+            rows.add(row(
+                    "audit",
+                    sender -> {
+                        sender.table("audit")
+                                .longColumn("l", rowId)
+                                .doubleColumn("d", rowId + 0.25)
+                                .symbol("sym", "stable");
+                        if (stringValue != null) {
+                            sender.stringColumn("s", stringValue);
+                        }
+                        if (longArray != null) {
+                            sender.longArray("la", longArray);
+                        }
+                        if (doubleArray != null) {
+                            sender.doubleArray("da", doubleArray);
+                        }
+                        if (decimal64 != null) {
+                            sender.decimalColumn("d64", decimal64);
+                        }
+                        if (decimal128 != null) {
+                            sender.decimalColumn("d128", decimal128);
+                        }
+                        if (decimal256 != null) {
+                            sender.decimalColumn("d256", decimal256);
+                        }
+                        sender.at(rowId + 1L, ChronoUnit.MICROS);
+                    },
+                    "l", (long) rowId,
+                    "d", rowId + 0.25,
+                    "sym", "stable",
+                    "s", stringValue,
+                    "la", longArray == null ? null : longArrayValue(shape(longArray.length), longArray),
+                    "da", doubleArray == null ? null : doubleArrayValue(shape(doubleArray.length, doubleArray[0].length), flatten(doubleArray)),
+                    "d64", decimal64 == null ? null : decimal(i * 100L + 7, 2),
+                    "d128", decimal128 == null ? null : decimal(i * 1000L + 11, 4),
+                    "d256", decimal256 == null ? null : decimal(i * 10000L + 13, 3),
+                    "", rowId + 1L
+            ));
+        }
+        assertEstimateAtLeastActual(rows);
+    }
+
+    private static void assertEstimateAtLeastActual(List<ScenarioRow> rows) throws Exception {
+        CapturingNetworkFacade nf = new CapturingNetworkFacade();
+        try (QwpUdpSender sender = new QwpUdpSender(nf, 0, 0, 9000, 1, 1024 * 1024)) {
+            for (int i = 0; i < rows.size(); i++) {
+                rows.get(i).writer.accept(sender);
+                long estimate = sender.committedDatagramEstimateForTest();
+                long actual = fullPacketSize(rows.subList(0, i + 1));
+                Assert.assertTrue(
+                        "row " + i + " estimate underflow: estimate=" + estimate + ", actual=" + actual,
+                        estimate >= actual
+                );
+            }
+        }
+    }
+
     private static List<DecodedRow> expectedRows(List<ScenarioRow> rows) {
         ArrayList<DecodedRow> expected = new ArrayList<>(rows.size());
         for (ScenarioRow row : rows) {