[kv] Batch old value lookups in write path

fluss-server/src/main/java/org/apache/fluss/server/kv/KvTablet.java

@@ -96,8 +96,11 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
@@ -638,18 +641,28 @@ private void processKvRecords(
         KvRecordBatch.ReadContext readContext =
                 KvRecordReadContext.createReadContext(kvFormat, schemaGetter);
         ValueDecoder valueDecoder = new ValueDecoder(schemaGetter, kvFormat);
+        List<PendingKvRecord> pendingRecords = new ArrayList<>(kvRecords.getRecordCount());
 
         for (KvRecord kvRecord : kvRecords.records(readContext)) {
             byte[] keyBytes = BytesUtils.toArray(kvRecord.getKey());
             KvPreWriteBuffer.Key key = KvPreWriteBuffer.Key.of(keyBytes);
             BinaryRow row = kvRecord.getRow();
             BinaryValue currentValue = row == null ? null : new BinaryValue(schemaIdOfNewData, row);
+            pendingRecords.add(new PendingKvRecord(key, currentValue));
+        }
+
+        Map<KvPreWriteBuffer.Key, byte[]> rocksDbOldValues =
+                bulkLoadOldValues(pendingRecords, currentMerger, autoIncrementUpdater);
 
+        for (PendingKvRecord pendingRecord : pendingRecords) {
+            KvPreWriteBuffer.Key key = pendingRecord.key;
+            BinaryValue currentValue = pendingRecord.currentValue;
             if (currentValue == null) {
                 logOffset =
                         processDeletion(
                                 key,
                                 currentMerger,
+                                rocksDbOldValues,
                                 valueDecoder,
                                 walBuilder,
                                 latestSchemaRow,
@@ -661,6 +674,7 @@ private void processKvRecords(
                                 currentValue,
                                 currentMerger,
                                 autoIncrementUpdater,
+                                rocksDbOldValues,
                                 valueDecoder,
                                 walBuilder,
                                 latestSchemaRow,
@@ -669,9 +683,72 @@ private void processKvRecords(
         }
     }
 
+    /**
+     * Loads the initial old values needed to process one KV record batch.
+     *
+     * <p>This method runs while {@code putAsLeader} holds the tablet write lock. It first checks
+     * the pre-write buffer, whose entries are newer than RocksDB, and sends only the remaining
+     * distinct keys through one RocksDB multi-get. A {@code null} value in the returned map means
+     * that RocksDB did not contain the key; map absence means that no RocksDB lookup was needed.
+     *
+     * <p>Records are still applied one by one in their original order. Before using an initial
+     * RocksDB value, {@link #getFromBufferOrOldValues} checks the pre-write buffer again so that a
+     * repeated key observes an earlier mutation from the same batch.
+     */
+    private Map<KvPreWriteBuffer.Key, byte[]> bulkLoadOldValues(
+            List<PendingKvRecord> pendingRecords,
+            RowMerger currentMerger,
+            AutoIncrementUpdater autoIncrementUpdater)
+            throws IOException {
+        Set<KvPreWriteBuffer.Key> keysRequiringOldValue = new LinkedHashSet<>();
+        for (PendingKvRecord pendingRecord : pendingRecords) {
+            if (requiresOldValue(pendingRecord, currentMerger, autoIncrementUpdater)) {
+                keysRequiringOldValue.add(pendingRecord.key);
+            }
+        }
+
+        List<KvPreWriteBuffer.Key> rocksDbKeys = new ArrayList<>();
+        List<byte[]> rocksDbKeyBytes = new ArrayList<>();
+        for (KvPreWriteBuffer.Key key : keysRequiringOldValue) {
+            if (kvPreWriteBuffer.get(key) == null) {
+                rocksDbKeys.add(key);
+                rocksDbKeyBytes.add(key.get());
+            }
+        }
+
+        Map<KvPreWriteBuffer.Key, byte[]> oldValues = new HashMap<>();
+        if (rocksDbKeys.isEmpty()) {
+            return oldValues;
+        }
+
+        List<byte[]> values = rocksDBKv.multiGet(rocksDbKeyBytes);
+        checkState(
+                values.size() == rocksDbKeys.size(),
+                "RocksDB multi-get returned %s values for %s keys.",
+                values.size(),
+                rocksDbKeys.size());
+        for (int i = 0; i < rocksDbKeys.size(); i++) {
+            oldValues.put(rocksDbKeys.get(i), values.get(i));
+        }
+        return oldValues;
+    }
+
+    private boolean requiresOldValue(
+            PendingKvRecord pendingRecord,
+            RowMerger currentMerger,
+            AutoIncrementUpdater autoIncrementUpdater) {
+        if (pendingRecord.currentValue == null) {
+            return currentMerger.deleteBehavior() == DeleteBehavior.ALLOW;
+        }
+        return changelogImage != ChangelogImage.WAL
+                || autoIncrementUpdater.hasAutoIncrement()
+                || !(currentMerger instanceof DefaultRowMerger);
+    }
+
     private long processDeletion(
             KvPreWriteBuffer.Key key,
             RowMerger currentMerger,
+            Map<KvPreWriteBuffer.Key, byte[]> rocksDbOldValues,
             ValueDecoder valueDecoder,
             WalBuilder walBuilder,
             PaddingRow latestSchemaRow,
@@ -687,7 +764,7 @@ private long processDeletion(
                             + "The table.delete.behavior is set to 'disable'.");
         }
 
-        byte[] oldValueBytes = getFromBufferOrKv(key);
+        byte[] oldValueBytes = getFromBufferOrOldValues(key, rocksDbOldValues);
         if (oldValueBytes == null) {
             LOG.debug(
                     "The specific key can't be found in kv tablet although the kv record is for deletion, "
@@ -711,6 +788,7 @@ private long processUpsert(
             BinaryValue currentValue,
             RowMerger currentMerger,
             AutoIncrementUpdater autoIncrementUpdater,
+            Map<KvPreWriteBuffer.Key, byte[]> rocksDbOldValues,
             ValueDecoder valueDecoder,
             WalBuilder walBuilder,
             PaddingRow latestSchemaRow,
@@ -726,7 +804,7 @@ private long processUpsert(
             return applyUpdate(key, null, currentValue, walBuilder, latestSchemaRow, logOffset);
         }
 
-        byte[] oldValueBytes = getFromBufferOrKv(key);
+        byte[] oldValueBytes = getFromBufferOrOldValues(key, rocksDbOldValues);
         if (oldValueBytes == null) {
             BinaryValue valueToInsert = currentMerger.merge(null, currentValue);
             return applyInsert(
@@ -1157,6 +1235,26 @@ private byte[] getFromBufferOrKv(KvPreWriteBuffer.Key key) throws IOException {
         return value.get();
     }
 
+    private byte[] getFromBufferOrOldValues(
+            KvPreWriteBuffer.Key key, Map<KvPreWriteBuffer.Key, byte[]> rocksDbOldValues) {
+        KvPreWriteBuffer.Value value = kvPreWriteBuffer.get(key);
+        if (value != null) {
+            return value.get();
+        }
+        checkState(rocksDbOldValues.containsKey(key), "Old value was not loaded for key %s.", key);
+        return rocksDbOldValues.get(key);
+    }
+
+    private static final class PendingKvRecord {
+        private final KvPreWriteBuffer.Key key;
+        private final @Nullable BinaryValue currentValue;
+
+        private PendingKvRecord(KvPreWriteBuffer.Key key, @Nullable BinaryValue currentValue) {
+            this.key = key;
+            this.currentValue = currentValue;
+        }
+    }
+
     public List<byte[]> multiGet(List<byte[]> keys) throws IOException {
         return inReadLock(
                 kvLock,

fluss-server/src/main/java/org/apache/fluss/server/kv/rocksdb/RocksDBStatistics.java

@@ -135,14 +135,16 @@ public long getCompactionBytesWritten() {
     /**
      * Get get operation latency in microseconds (P99).
      *
-     * <p>This uses RocksDB Statistics histogram data to get the P99 latency of get operations. P99
-     * is used instead of average because it better reflects tail latency issues, which are more
-     * critical for monitoring database performance.
+     * <p>This uses RocksDB Statistics histogram data to get the higher P99 latency of point-get and
+     * multi-get operations. P99 is used instead of average because it better reflects tail latency
+     * issues, which are more critical for monitoring database performance.
      *
      * @return P99 get latency in microseconds, or 0 if not available
      */
     public long getGetLatencyMicros() {
-        return getHistogramValue(HistogramType.DB_GET);
+        return Math.max(
+                getHistogramValue(HistogramType.DB_GET),
+                getHistogramValue(HistogramType.DB_MULTIGET));
     }
 
     /**

fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletTest.java

@@ -124,6 +124,12 @@
 import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.assertj.core.api.Fail.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 /** Test for {@link KvTablet}. */
 class KvTabletTest {
@@ -578,6 +584,46 @@ void testPartialUpdateFirstInsertThenUpdate() throws Exception {
         checkEqual(actualLogRecords, expectedLogs, rowType);
     }
 
+    @Test
+    void testBatchOldValueLookupUsesSingleDeduplicatedRocksDbMultiGet() throws Exception {
+        initLogTabletAndKvTablet(DATA2_SCHEMA, new HashMap<>());
+        KvRecordTestUtils.KvRecordFactory recordFactory =
+                KvRecordTestUtils.KvRecordFactory.of(DATA2_ROW_TYPE);
+
+        kvTablet.putAsLeader(
+                kvRecordBatchFactory.ofRecords(
+                        recordFactory.ofRecord("k1", new Object[] {1, "v1", null}),
+                        recordFactory.ofRecord("k2", new Object[] {2, "v2", null})),
+                null);
+        flushAndWait(kvTablet, Long.MAX_VALUE);
+
+        // Keep k3 in the pre-write buffer. Its old value must not be loaded from RocksDB.
+        kvTablet.putAsLeader(
+                kvRecordBatchFactory.ofRecords(
+                        recordFactory.ofRecord("k3", new Object[] {3, "v3", null})),
+                null);
+
+        RocksDBKv rocksDBKv = spy(kvTablet.getRocksDBKv());
+        Field rocksDBKvField = KvTablet.class.getDeclaredField("rocksDBKv");
+        rocksDBKvField.setAccessible(true);
+        rocksDBKvField.set(kvTablet, rocksDBKv);
+
+        kvTablet.putAsLeader(
+                kvRecordBatchFactory.ofRecords(
+                        recordFactory.ofRecord("k1", new Object[] {1, "v11", null}),
+                        recordFactory.ofRecord("k1", new Object[] {1, "v12", null}),
+                        recordFactory.ofRecord("k2", new Object[] {2, "v22", null}),
+                        recordFactory.ofRecord("k3", new Object[] {3, "v33", null})),
+                new int[] {0, 1});
+
+        verify(rocksDBKv, times(1)).multiGet(argThat(keys -> keys.size() == 2));
+        verify(rocksDBKv, never()).get(any(byte[].class));
+        assertThat(kvTablet.getKvPreWriteBuffer().get(Key.of("k1".getBytes())))
+                .isEqualTo(valueOf(compactedRow(DATA2_ROW_TYPE, new Object[] {1, "v12", null})));
+        assertThat(kvTablet.getKvPreWriteBuffer().get(Key.of("k3".getBytes())))
+                .isEqualTo(valueOf(compactedRow(DATA2_ROW_TYPE, new Object[] {3, "v33", null})));
+    }
+
     @Test
     void testPutWithMultiThread() throws Exception {
         initLogTabletAndKvTablet(DATA1_SCHEMA_PK, new HashMap<>());