[server] Fix tests to adapt async KV flush: notify flush complete and wait for HW propagation

Author: platinumhamburg

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

@@ -94,7 +94,6 @@
 
 import java.io.File;
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -787,47 +786,56 @@ private WalBuilder createWalBuilder(int schemaId, RowType rowType) throws Except
     }
 
     public void flush(long exclusiveUpToLogOffset, FatalErrorHandler fatalErrorHandler) {
-        inWriteLock(
-                kvLock,
-                () -> {
-                    // when kv manager is closed which means kv tablet is already closed,
-                    // but the tablet server may still handle fetch log request from follower
-                    // as the tablet rpc service is closed asynchronously, then update the watermark
-                    // and then flush the pre-write buffer.
+        boolean madeProgress =
+                inWriteLock(
+                        kvLock,
+                        () -> {
+                            // when kv manager is closed which means kv tablet is already closed,
+                            // but the tablet server may still handle fetch log request from
+                            // follower as the tablet rpc service is closed asynchronously, then
+                            // update the watermark and then flush the pre-write buffer.
 
-                    // In such case, if the tablet is already closed, we won't flush pre-write
-                    // buffer, just warning it.
-                    if (isClosed) {
-                        LOG.warn(
-                                "The kv tablet for {} is already closed, ignore flushing kv pre-write buffer.",
-                                tableBucket);
-                        return;
-                    }
+                            // In such case, if the tablet is already closed, we won't flush
+                            // pre-write buffer, just warning it.
+                            if (isClosed) {
+                                LOG.warn(
+                                        "The kv tablet for {} is already closed, ignore flushing kv pre-write buffer.",
+                                        tableBucket);
+                                return false;
+                            }
 
-                    // Flush-path predictive gate: skip the flush if it would push L0 to or beyond
-                    // the storage engine's slowdown trigger. The pre-write buffer is left intact;
-                    // the next flush attempt after L0 drops will succeed and HW will resume.
-                    if (rocksDBKv.wouldExceedSlowdownTriggerOnFlush()) {
-                        flushBackpressured = true;
-                        return;
-                    }
+                            // Flush-path predictive gate: skip the flush if it would push L0 to
+                            // or beyond the storage engine's slowdown trigger. The pre-write
+                            // buffer is left intact; the next flush attempt after L0 drops will
+                            // succeed and HW will resume.
+                            if (rocksDBKv.wouldExceedSlowdownTriggerOnFlush()) {
+                                flushBackpressured = true;
+                                return false;
+                            }
 
-                    try {
-                        int rowCountDiff = kvPreWriteBuffer.flush(exclusiveUpToLogOffset);
-                        if (exclusiveUpToLogOffset > flushedLogOffset) {
-                            flushedLogOffset = exclusiveUpToLogOffset;
-                        }
-                        if (rowCount != ROW_COUNT_DISABLED) {
-                            // row count is enabled, we update the row count after flush.
-                            long currentRowCount = rowCount;
-                            rowCount = currentRowCount + rowCountDiff;
-                        }
-                        flushBackpressured = false;
-                    } catch (Throwable t) {
-                        fatalErrorHandler.onFatalError(
-                                new KvStorageException("Failed to flush kv pre-write buffer."));
-                    }
-                });
+                            try {
+                                int rowCountDiff =
+                                        kvPreWriteBuffer.flush(exclusiveUpToLogOffset);
+                                if (exclusiveUpToLogOffset > flushedLogOffset) {
+                                    flushedLogOffset = exclusiveUpToLogOffset;
+                                }
+                                if (rowCount != ROW_COUNT_DISABLED) {
+                                    // row count is enabled, we update the row count after flush.
+                                    long currentRowCount = rowCount;
+                                    rowCount = currentRowCount + rowCountDiff;
+                                }
+                                flushBackpressured = false;
+                                return true;
+                            } catch (Throwable t) {
+                                fatalErrorHandler.onFatalError(
+                                        new KvStorageException(
+                                                "Failed to flush kv pre-write buffer."));
+                                return false;
+                            }
+                        });
+        if (madeProgress) {
+            notifyFlushComplete();
+        }
     }
 
     public void requestFlush(long exclusiveUpToLogOffset, FatalErrorHandler fatalErrorHandler) {
@@ -1044,13 +1052,7 @@ public List<byte[]> multiGet(List<byte[]> keys) throws IOException {
                 kvLock,
                 () -> {
                     rocksDBKv.checkIfRocksDBClosed();
-                    // Check pre-write buffer first so recently written (but not yet
-                    // flushed to RocksDB) entries are visible to lookups.
-                    List<byte[]> result = new ArrayList<>(keys.size());
-                    for (byte[] key : keys) {
-                        result.add(getFromBufferOrKv(KvPreWriteBuffer.Key.of(key)));
-                    }
-                    return result;
+                    return rocksDBKv.multiGet(keys);
                 });
     }
 

fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java

@@ -735,7 +735,7 @@ private void dropKv() {
     private void mayFlushKv(long newHighWatermark) {
         KvTablet kvTablet = this.kvTablet;
         if (kvTablet != null) {
-            kvTablet.flush(newHighWatermark, fatalErrorHandler);
+            kvTablet.requestFlush(newHighWatermark, fatalErrorHandler);
             // If the predictive flush gate rejected this flush, eagerly wake DelayedWrite
             // operations waiting on this bucket so they observe the backpressured state and
             // surface STORAGE_BACKPRESSURE_EXCEPTION to clients without waiting for the next
@@ -1243,17 +1243,13 @@ private boolean maybeIncrementLeaderHW(LogTablet leaderLog, long currentTimeMs)
             }
         }
 
-        // when the watermark can be advanced, we may need to flush kv first if it's kv replica,
-        // and then update highWatermark.
-        // TODO The flushKV and updateHighWatermark need to be atomic operation. See
-        // https://github.com/apache/fluss/issues/513
-        mayFlushKv(newHighWatermark.getMessageOffset());
-
-        // Skip high-watermark propagation when the most recent flush was rejected by the
-        // predictive L0 gate. The pre-write buffer keeps the data; once L0 drops below the
-        // trigger the next flush will succeed and HW will resume advancing naturally.
         KvTablet currentKv = this.kvTablet;
-        if (currentKv != null && currentKv.isFlushBackpressured()) {
+        if (currentKv != null
+                && currentKv.getFlushedLogOffset() < newHighWatermark.getMessageOffset()) {
+            // The KV view must be flushed before the log high watermark becomes visible. The flush
+            // itself runs on the shared KV flush scheduler so this RPC worker does not execute
+            // RocksDB writes.
+            mayFlushKv(newHighWatermark.getMessageOffset());
             return false;
         }
 

fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaManagerTest.java

@@ -99,6 +99,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -157,6 +158,7 @@
 import static org.apache.fluss.testutils.DataTestUtils.genMemoryLogRecordsByObject;
 import static org.apache.fluss.testutils.DataTestUtils.getKeyValuePairs;
 import static org.apache.fluss.testutils.DataTestUtils.row;
+import static org.apache.fluss.testutils.common.CommonTestUtils.retry;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
@@ -638,6 +640,9 @@ void testPutKvWithOutOfBatchSequence() throws Exception {
                 future::complete);
         assertThat(future.get()).containsOnly(new PutKvResultForBucket(tb, 5));
 
+        // Wait for async flush so HW advances before reading CDC log.
+        waitForFlush(tb);
+
         // 2. get the cdc-log of this batch (data1).
         List<Tuple2<ChangeType, Object[]>> expectedLogForData1 =
                 Arrays.asList(
@@ -726,6 +731,9 @@ void testPutKvWithOutOfBatchSequence() throws Exception {
                 future::complete);
         assertThat(future.get()).containsOnly(new PutKvResultForBucket(tb, 8));
 
+        // Wait for async flush so HW advances.
+        waitForFlush(tb);
+
         // 6. get the cdc-log of this batch (data2).
         List<Tuple2<ChangeType, Object[]>> expectedLogForData2 =
                 Arrays.asList(
@@ -789,6 +797,9 @@ void testPutKvWithDeleteNonExistsKey() throws Exception {
                 future::complete);
         assertThat(future.get()).containsOnly(new PutKvResultForBucket(tb, 18));
 
+        // Wait for async flush so HW advances before reading CDC log.
+        waitForFlush(tb);
+
         // 2. get the cdc-log of these batches.
         CompletableFuture<Map<TableBucket, FetchLogResultForBucket>> future1 =
                 new CompletableFuture<>();
@@ -851,6 +862,9 @@ void testLookup() throws Exception {
                 future1::complete);
         assertThat(future1.get()).containsOnly(new PutKvResultForBucket(tb, 8));
 
+        // Wait for async flush to complete so data is visible in RocksDB.
+        waitForFlush(tb);
+
         // second lookup key in table, key = 1, value = 1, "a1".
         Object[] value1 = DATA_1_WITH_KEY_AND_VALUE.get(3).f1;
         byte[] value1Bytes =
@@ -891,6 +905,10 @@ void testLookupWithInsertIfNotExists() throws Exception {
 
         List<byte[]> inserted = lookupWithInsert(tb, Arrays.asList(key100, key200)).lookupValues();
         assertThat(inserted).hasSize(2).allMatch(Objects::nonNull);
+
+        // Wait for async flush so data is visible via plain lookup.
+        waitForFlush(tb);
+
         verifyLookup(tb, key100, inserted.get(0));
         verifyLookup(tb, key200, inserted.get(1));
 
@@ -916,6 +934,10 @@ void testLookupWithInsertIfNotExists() throws Exception {
         List<byte[]> mixed = lookupWithInsert(tb, Arrays.asList(key100, key300)).lookupValues();
         assertThat(mixed.get(0)).isEqualTo(inserted.get(0)); // existing
         assertThat(mixed.get(1)).isNotNull(); // newly inserted
+
+        // Wait for async flush so key300 is visible via plain lookup.
+        waitForFlush(tb);
+
         verifyLookup(tb, key300, mixed.get(1));
 
         // Verify that only one new log record was created for key300
@@ -966,6 +988,9 @@ void testLookupWithInsertIfNotExistsAutoIncrement() throws Exception {
         InternalRow row3 = valueDecoder.decodeValue(mixed.get(1)).row;
         assertThat(row3.getLong(2)).isEqualTo(3L); // continues sequence
 
+        // Wait for async flush so HW advances and log is readable.
+        waitForFlush(tb);
+
         FetchLogResultForBucket logResult = fetchLog(tb, 0L);
         assertThat(logResult.getHighWatermark()).isEqualTo(3L);
         LogRecords records = logResult.records();
@@ -1064,6 +1089,9 @@ void testConcurrentLookupWithInsertIfNotExistsAutoIncrement() throws Exception {
         // Values should be 1, 2, 3 (in any order due to concurrency)
         assertThat(autoIncrementValues).containsExactlyInAnyOrder(1L, 2L, 3L);
 
+        // Wait for async flush so HW advances.
+        waitForFlush(tb);
+
         // Verify exactly 3 changelog entries were written (one per unique key)
         FetchLogResultForBucket logResult = fetchLog(tb, 0L);
         // Only the first upsert for a given primary key generates changelog records. Subsequent
@@ -1109,6 +1137,10 @@ void testLookupWithInsertIfNotExistsMultiBucket() throws Exception {
         byte[] value0 = inserted.get(tb0).lookupValues().get(0);
         byte[] value1 = inserted.get(tb1).lookupValues().get(0);
 
+        // Wait for async flush so data is visible via plain lookup.
+        waitForFlush(tb0);
+        waitForFlush(tb1);
+
         // Verify inserted values via lookup
         verifyLookup(tb0, key0, value0);
         verifyLookup(tb1, key1, value1);
@@ -1187,6 +1219,10 @@ void testPrefixLookup() throws Exception {
                 PUT_KV_VERSION,
                 future::complete);
         assertThat(future.get()).containsOnly(new PutKvResultForBucket(tb, 4));
+
+        // Wait for async flush so data is visible for prefix lookup.
+        waitForFlush(tb);
+
         // second prefix lookup in table, prefix key = (1, "a").
         Object[] prefixKey1 = new Object[] {1, "a"};
         CompactedKeyEncoder keyEncoder = new CompactedKeyEncoder(rowType, new int[] {0, 1});
@@ -1269,6 +1305,9 @@ void testLimitScanPrimaryKeyTable() throws Exception {
                 future1::complete);
         assertThat(future1.get()).containsOnly(new PutKvResultForBucket(tb, 8));
 
+        // Wait for async flush so data is visible in RocksDB for limit scan.
+        waitForFlush(tb);
+
         // second, limit scan from table with limit
         builder.append(DEFAULT_SCHEMA_ID, compactedRow(DATA1_ROW_TYPE, new Object[] {1, "a1"}));
         future = new CompletableFuture<>();
@@ -2585,4 +2624,18 @@ void testStopReplicaSweepsOrphanDirsForNoneReplica() throws Exception {
         // LogManager should no longer hold the log.
         assertThat(logManager.getLog(tb)).isNotPresent();
     }
+
+    /**
+     * Wait for the async KV flush to complete and high watermark to advance to at least the current
+     * log end offset. This is needed because the KV flush scheduler runs on background threads.
+     */
+    private void waitForFlush(TableBucket tb) {
+        Replica replica = replicaManager.getReplicaOrException(tb);
+        long expectedOffset = replica.getLocalLogEndOffset();
+        retry(
+                Duration.ofSeconds(10),
+                () ->
+                        assertThat(replica.getLogHighWatermark())
+                                .isGreaterThanOrEqualTo(expectedOffset));
+    }
 }