diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvRecoverHelper.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvRecoverHelper.java index a3afeec4bbf..109ccb21f4f 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvRecoverHelper.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvRecoverHelper.java @@ -253,6 +253,9 @@ private long readLogRecordsAndApply( rowCountUpdater, autoIncIdRangeUpdater, resumeRecordConsumer); + if (fetchIsolation == FetchIsolation.LOG_END) { + kvTablet.getKvPreWriteBuffer().markWalBatchEnd(nextFetchOffset); + } } } return nextFetchOffset; diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvTablet.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvTablet.java index 76c603c8ec0..2312b6fd4e9 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvTablet.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvTablet.java @@ -123,11 +123,11 @@ public final class KvTablet { private static final long ROW_COUNT_DISABLED = -1; /** - * Max records per native write of the asynchronous flush; mirrors the batching capacity of + * KV entry budget per native write of the asynchronous flush; mirrors the batching capacity of * {@code RocksDBWriteBatchWrapper} (hundreds of keys per write batch is RocksDB best practice). - * Together with {@code writeBatchSize} this bounds one atomic native write. + * The budget is checked after each complete WAL batch to keep its KV mutations atomic. */ - private static final int MAX_RECORDS_PER_NATIVE_WRITE = 500; + private static final int TARGET_ENTRIES_PER_NATIVE_WRITE = 500; private final PhysicalTablePath physicalPath; private final TableBucket tableBucket; @@ -178,6 +178,8 @@ public final class KvTablet { private volatile long rowCount; + @Nullable private Runnable beforeNativeWrite; + @GuardedBy("kvLock") private volatile boolean isClosed = false; @@ -804,16 +806,27 @@ private LogAppendInfo putAsLeader( tableBucket)); } - return kvWriteProcessor.putAsLeader( - kvRecords, - targetColumns, - mergeMode, - kvStateAccessor, - originalPartitionName, - memoizedLakeLookup); + LogAppendInfo appendInfo = + kvWriteProcessor.putAsLeader( + kvRecords, + targetColumns, + mergeMode, + kvStateAccessor, + originalPartitionName, + memoizedLakeLookup); + if (!appendInfo.duplicated()) { + // KvWriteProcessor appends one WAL batch for each accepted KV batch. + kvPreWriteBuffer.markWalBatchEnd(appendInfo.lastOffset() + 1); + } + return appendInfo; }); } + @VisibleForTesting + void setBeforeNativeWrite(@Nullable Runnable beforeNativeWrite) { + this.beforeNativeWrite = beforeNativeWrite; + } + @VisibleForTesting long localLogEndOffset() { return logTablet.localLogEndOffset(); @@ -958,16 +971,17 @@ private boolean tryAcquireScheduledFlush() { } /** - * Writes the prepared entries to RocksDB in segments of at most {@code - * MAX_RECORDS_PER_NATIVE_WRITE} records / {@code writeBatchSize} bytes. Each segment forms - * exactly one atomic native write (the writer has implicit flushes disabled) and is completed - * immediately after it lands, so {@code flushedLogOffset}/{@code rowCount} stay consistent with - * the RocksDB content even if a later segment is rejected by the no-slowdown gate. + * Writes the prepared entries to RocksDB in complete WAL batch groups targeting {@code + * TARGET_ENTRIES_PER_NATIVE_WRITE} entries / {@code writeBatchSize} bytes. Budgets are checked + * after each complete WAL batch. Each segment forms exactly one atomic native write (the writer + * has implicit flushes disabled) and is completed immediately after it lands, so {@code + * flushedLogOffset}/{@code rowCount} stay consistent with the RocksDB content even if a later + * segment is rejected by the no-slowdown gate. */ @GuardedBy("kvLock") private void writePreparedFlush(PreparedFlush preparedFlush) throws Exception { List segments = - preparedFlush.split(writeBatchSize, MAX_RECORDS_PER_NATIVE_WRITE); + preparedFlush.split(writeBatchSize, TARGET_ENTRIES_PER_NATIVE_WRITE); int nextSegment = 0; try (ResourceGuard.Lease lease = rocksDBKv.getResourceGuard().acquireResource(); KvBatchWriter kvBatchWriter = createNoSlowdownKvBatchWriter()) { @@ -987,6 +1001,9 @@ private void writePreparedFlush(PreparedFlush preparedFlush) throws Exception { kvBatchWriter.put(entry.getKey().get(), value.get()); } } + if (beforeNativeWrite != null) { + beforeNativeWrite.run(); + } kvBatchWriter.flush(); completeFlushedSegment(segment); nextSegment++; diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBuffer.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBuffer.java index aae74e19b7f..18e922cf3f0 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBuffer.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBuffer.java @@ -110,6 +110,23 @@ public KvPreWriteBuffer(TabletServerMetricGroup serverMetricGroup) { truncateAsErrorCount = serverMetricGroup.kvTruncateAsErrorCount(); } + /** + * Marks the last KV mutation when it ends a successfully appended or recovered WAL batch. + * + *

Call after applying each batch and before staging any mutation from the next batch. Empty + * batches do not mark an entry; their offsets are covered by the flush target or a later batch. + * + * @param batchEndOffset the exclusive end offset of the completed WAL batch + */ + public void markWalBatchEnd(long batchEndOffset) { + KvEntry last = allKvEntries.peekLast(); + // Empty batches have no mutation to mark. Their offsets are covered by the flush target + // or a later batch; never infer a boundary from a gap between KV mutations. + if (last != null && last.logSequenceNumber + 1 == batchEndOffset) { + last.endOfWalBatch = true; + } + } + /** * Delete a key-value pair with the given key. * @@ -363,6 +380,7 @@ public static class KvEntry { private final Key key; private final Value value; private final long logSequenceNumber; + private boolean endOfWalBatch; // the previous mapped value in the buffer before this key-value put; null once the // referenced entry has been flushed (see completeFlush) @@ -496,57 +514,42 @@ public boolean isEmpty() { } /** - * Splits this prepared flush into consecutive segments so that each segment can be written - * to the kv storage as one native write and completed via {@link #completeFlush} - * independently. Segments preserve the entry order, so completing them in order keeps the - * list-prefix invariant checked by {@link #completeFlush}. - * - *

The upper log sequence number of every segment except the last one is the log sequence - * number of the first entry of the next segment, so advancing {@code flushedLogOffset} to a - * segment boundary never claims an entry that has not been written yet. The last segment - * keeps the original target so an empty tail still publishes the full flush range. + * Splits the prepared prefix only between complete WAL batches. Each segment is one atomic + * native write and can therefore publish its end independently after success. Budgets are + * checked after adding each complete WAL batch. * - * @param maxBytesPerSegment max key/value payload bytes per segment, {@code <= 0} means - * unlimited; a single entry larger than the limit is kept as an oversized singleton - * @param maxRecordsPerSegment max record count per segment + * @param targetBytesPerSegment key/value payload budget, or non-positive for unlimited + * @param targetEntriesPerSegment KV entry budget; the final batch may exceed either budget */ - public List split(long maxBytesPerSegment, int maxRecordsPerSegment) { - checkArgument(maxRecordsPerSegment > 0, "maxRecordsPerSegment must be positive."); - // Single pass: only boundary detection needs to visit every entry (byte sizes and - // row-count deltas). Segments are zero-copy subList views of the entry list. + public List split(long targetBytesPerSegment, int targetEntriesPerSegment) { + checkArgument(targetEntriesPerSegment > 0, "targetEntriesPerSegment must be positive."); List segments = null; int segmentStart = 0; long segmentBytes = 0; int segmentRowCountDiff = 0; for (int i = 0; i < entries.size(); i++) { KvEntry entry = entries.get(i); - long currentEntryBytes = entryBytes(entry.getKey(), entry.getValue()); - boolean hasEntries = i > segmentStart; - boolean recordLimitReached = hasEntries && i - segmentStart >= maxRecordsPerSegment; - boolean byteLimitExceeded = - hasEntries - && maxBytesPerSegment > 0 - && currentEntryBytes > maxBytesPerSegment - segmentBytes; - if (recordLimitReached || byteLimitExceeded) { - // Seal the current segment right before this entry: all entries below this - // entry's log sequence number are contained in the sealed segments. + segmentBytes += entryBytes(entry.getKey(), entry.getValue()); + segmentRowCountDiff += rowCountDelta(entry); + if (entry.endOfWalBatch + && i + 1 < entries.size() + && (i + 1 - segmentStart >= targetEntriesPerSegment + || (targetBytesPerSegment > 0 + && segmentBytes >= targetBytesPerSegment))) { if (segments == null) { segments = new ArrayList<>(); } segments.add( new PreparedFlush( - entry.getLogSequenceNumber(), - entries.subList(segmentStart, i), + entry.logSequenceNumber + 1, + entries.subList(segmentStart, i + 1), segmentRowCountDiff)); - segmentStart = i; + segmentStart = i + 1; segmentBytes = 0; segmentRowCountDiff = 0; } - segmentBytes += currentEntryBytes; - segmentRowCountDiff += rowCountDelta(entry); } if (segments == null) { - // Everything fits into a single segment: reuse this prepared flush as-is. return Collections.singletonList(this); } segments.add( diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/KvFlushBatchBoundaryTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/KvFlushBatchBoundaryTest.java new file mode 100644 index 00000000000..a2c53dec426 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/KvFlushBatchBoundaryTest.java @@ -0,0 +1,218 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.server.kv; + +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.config.MemorySize; +import org.apache.fluss.exception.StorageBackpressureException; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.rpc.protocol.MergeMode; +import org.apache.fluss.server.entity.NotifyLeaderAndIsrData; +import org.apache.fluss.server.log.FetchIsolation; +import org.apache.fluss.server.replica.Replica; +import org.apache.fluss.server.replica.ReplicaTestBase; +import org.apache.fluss.server.zk.data.LeaderAndIsr; +import org.apache.fluss.utils.types.Tuple2; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.apache.fluss.record.TestData.DATA1_PHYSICAL_TABLE_PATH_PK; +import static org.apache.fluss.record.TestData.DATA1_TABLE_ID_PK; +import static org.apache.fluss.server.coordinator.CoordinatorContext.INITIAL_COORDINATOR_EPOCH; +import static org.apache.fluss.testutils.DataTestUtils.genKvRecordBatch; +import static org.apache.fluss.testutils.DataTestUtils.genKvRecords; +import static org.apache.fluss.testutils.DataTestUtils.getKeyValuePairs; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests asynchronous KV flush boundaries through leader HW publication and WAL recovery. */ +class KvFlushBatchBoundaryTest extends ReplicaTestBase { + @Override + protected KvFlushScheduler createTestKvFlushScheduler(Configuration conf) { + return new KvFlushScheduler() { + @Override + public void enqueue(KvTablet tablet) {} + + @Override + public void retryLater(KvTablet tablet, long delayMs) {} + + @Override + public void close() {} + }; + } + + @Test + void testBackpressurePublishesWholeBatchAndRetryCompletes() throws Exception { + Replica replica = createLeader(); + try { + appendBatches(replica); + long target = replica.getLocalLogEndOffset(); + assertThat(replica.getLogHighWatermark()).isZero(); + KvTablet tablet = replica.getKvTablet(); + rejectSecondWrite(tablet); + tablet.runScheduledFlush(); + assertCommittedPrefix(replica); + + tablet.setBeforeNativeWrite(null); + tablet.requestFlushRetry(); + tablet.runScheduledFlush(); + assertThat(tablet.getFlushedLogOffset()).isEqualTo(target); + assertThat(replica.getLogHighWatermark()).isEqualTo(target); + assertThat(replica.checkEnoughReplicasReachOffset(target).f0).isTrue(); + assertThat(tablet.getRowCount()).isEqualTo(601); + assertThat(tablet.getKvPreWriteBuffer().pendingFlushBytes()).isZero(); + assertValue(tablet, 0, "before"); + assertValue(tablet, 600, "tail"); + assertThat( + replica.getLogTablet() + .read(0, Integer.MAX_VALUE, FetchIsolation.HIGH_WATERMARK, true) + .getRecords() + .batches()) + .extracting(batch -> batch.nextLogOffset()) + .containsExactly(600L, 601L); + } finally { + deleteReplica(replica); + } + } + + @Test + void testRecoveredBatchBoundariesSurviveBackpressureAndTruncation() throws Exception { + conf.set(ConfigOptions.KV_WRITE_BATCH_SIZE, new MemorySize(64)); + Replica replica = createLeader(); + try { + appendBatches(replica); + replica.makeFollower(state(TABLET_SERVER_ID + 1, 1)); + replica.makeLeader(state(TABLET_SERVER_ID, 2)); + KvTablet tablet = replica.getKvTablet(); + tablet.requestFlush( + replica.getLocalLogEndOffset(), + failure -> { + throw new AssertionError(failure); + }); + rejectSecondWrite(tablet); + tablet.runScheduledFlush(); + assertCommittedPrefix(replica); + long hw = replica.getLogHighWatermark(); + replica.makeFollower(state(TABLET_SERVER_ID + 1, 3)); + assertThat(replica.isLeader()).isFalse(); + replica.truncateTo(hw); + assertThat(replica.getLocalLogEndOffset()).isEqualTo(hw); + assertThat( + replica.getLogTablet() + .read(0, Integer.MAX_VALUE, FetchIsolation.LOG_END, true) + .getRecords() + .batches()) + .singleElement() + .satisfies(batch -> assertThat(batch.nextLogOffset()).isEqualTo(hw)); + replica.makeLeader(state(TABLET_SERVER_ID, 4)); + assertThat(replica.getKvTablet().getRowCount()).isEqualTo(600); + assertValue(replica.getKvTablet(), 0, "before"); + replica.putRecordsToLeader( + genKvRecordBatch(new Object[] {600, "tail"}), null, MergeMode.DEFAULT, 0); + replica.getKvTablet().runScheduledFlush(); + assertThat(replica.getLogHighWatermark()).isEqualTo(601); + assertValue(replica.getKvTablet(), 600, "tail"); + } finally { + deleteReplica(replica); + } + } + + private void deleteReplica(Replica replica) throws Exception { + try { + replica.delete(); + } finally { + replicaManager + .getServerMetricGroup() + .removeTableBucketMetricGroup( + DATA1_PHYSICAL_TABLE_PATH_PK.getTablePath(), + new TableBucket(DATA1_TABLE_ID_PK, 1)); + } + } + + private Replica createLeader() throws Exception { + Replica replica = + makeKvReplica(DATA1_PHYSICAL_TABLE_PATH_PK, new TableBucket(DATA1_TABLE_ID_PK, 1)); + replica.makeLeader(state(TABLET_SERVER_ID, 0)); + return replica; + } + + private NotifyLeaderAndIsrData state(int leader, int epoch) { + return new NotifyLeaderAndIsrData( + DATA1_PHYSICAL_TABLE_PATH_PK, + new TableBucket(DATA1_TABLE_ID_PK, 1), + Arrays.asList(TABLET_SERVER_ID, TABLET_SERVER_ID + 1), + new LeaderAndIsr( + leader, + epoch, + Collections.singletonList(leader), + Collections.emptyList(), + INITIAL_COORDINATOR_EPOCH, + epoch)); + } + + private void appendBatches(Replica replica) throws Exception { + Object[][] rows = new Object[600][]; + for (int i = 0; i < rows.length; i++) { + rows[i] = new Object[] {i, "before"}; + } + replica.putRecordsToLeader(genKvRecordBatch(rows), null, MergeMode.DEFAULT, 0); + replica.putRecordsToLeader( + genKvRecordBatch(new Object[] {600, "tail"}), null, MergeMode.DEFAULT, 0); + assertThat(replica.getLocalLogEndOffset()).isEqualTo(601); + } + + private void rejectSecondWrite(KvTablet tablet) { + AtomicInteger attempts = new AtomicInteger(); + tablet.setBeforeNativeWrite( + () -> { + if (attempts.incrementAndGet() == 2) { + throw new StorageBackpressureException("Reject second native write"); + } + }); + } + + private void assertCommittedPrefix(Replica replica) throws Exception { + KvTablet tablet = replica.getKvTablet(); + assertThat(tablet.getFlushState()).isEqualTo(KvTablet.FlushState.STORAGE_BLOCKED); + assertThat(tablet.getFlushedLogOffset()).isEqualTo(600); + assertThat(replica.getLogHighWatermark()).isEqualTo(600); + assertThat(replica.checkEnoughReplicasReachOffset(600).f0).isTrue(); + assertThat(replica.checkEnoughReplicasReachOffset(601).f0).isFalse(); + assertThat(tablet.getRowCount()).isEqualTo(600); + assertValue(tablet, 0, "before"); + assertValue(tablet, 599, "before"); + assertThat( + replica.getLogTablet() + .read(0, Integer.MAX_VALUE, FetchIsolation.HIGH_WATERMARK, true) + .getRecords() + .batches()) + .singleElement() + .satisfies(batch -> assertThat(batch.nextLogOffset()).isEqualTo(600)); + } + + private void assertValue(KvTablet tablet, int key, String value) throws Exception { + Tuple2 expected = + getKeyValuePairs(genKvRecords(new Object[] {key, value})).get(0); + assertThat(tablet.multiGet(Collections.singletonList(expected.f0)).get(0).toByteArray()) + .containsExactly(expected.f1); + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletTest.java index 4bd4e58e5b2..568d9f9fa57 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletTest.java @@ -113,6 +113,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -2005,19 +2006,24 @@ void testScheduledFlushWritesLargePreparedRangeInSegments() throws Exception { new HashMap<>(), manualScheduler); - // 1200 records force the prepared range to be written as multiple native segments - // (500 records each), exercising the per-segment completion of the scheduled flush. + // The 500-entry budget groups complete 400-entry batches into writes of 800 and 400. int recordCount = 1200; List records = new ArrayList<>(recordCount); for (int i = 0; i < recordCount; i++) { records.add(kvRecordFactory.ofRecord("key" + i, new Object[] {i, "v" + i})); } - kvTablet.putAsLeader(kvRecordBatchFactory.ofRecords(records), null); + for (int start = 0; start < recordCount; start += 400) { + kvTablet.putAsLeader( + kvRecordBatchFactory.ofRecords(records.subList(start, start + 400)), null); + } + AtomicInteger nativeWrites = new AtomicInteger(); + kvTablet.setBeforeNativeWrite(nativeWrites::incrementAndGet); long flushOffset = logTablet.localLogEndOffset(); kvTablet.requestFlush(flushOffset, NOPErrorHandler.INSTANCE); kvTablet.runScheduledFlush(); + assertThat(nativeWrites.get()).isEqualTo(2); assertThat(kvTablet.getFlushedLogOffset()).isEqualTo(flushOffset); assertThat(kvTablet.getRowCount()).isEqualTo(recordCount); assertThat(kvTablet.getKvPreWriteBuffer().pendingFlushBytes()).isEqualTo(0); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferTest.java index 26d49807bb0..45d983133a1 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferTest.java @@ -232,15 +232,17 @@ void testSplitPreparedFlushByRecordCount() { // +key0(lsn 0), +key1(lsn 1), +key2(lsn 2), +key3(lsn 3), -key2(lsn 4) for (int i = 0; i < 4; i++) { bufferInsert(buffer, "key" + i, "value" + i, i); + buffer.markWalBatchEnd(i + 1); } bufferDelete(buffer, "key2", 4); + buffer.markWalBatchEnd(5); + buffer.markWalBatchEnd(10); PreparedFlush preparedFlush = buffer.prepareFlush(10); List segments = preparedFlush.split(0, 2); assertThat(segments).hasSize(3); - // Every segment boundary is the lsn of the first entry of the next segment; the last - // segment keeps the original target so the full flush range gets published. + // Each entry here is a complete WAL batch. The last segment includes the empty tail. assertThat(segments.get(0).entries()).hasSize(2); assertThat(segments.get(0).exclusiveUpToLogSequenceNumber()).isEqualTo(2); assertThat(segments.get(1).entries()).hasSize(2); @@ -265,20 +267,23 @@ void testSplitPreparedFlushByByteSize() { KvPreWriteBuffer buffer = new KvPreWriteBuffer(TestingMetricGroups.TABLET_SERVER_METRICS); // Entry payload sizes are 6, 5, and 5 bytes. bufferInsert(buffer, "a", "12345", 0); + buffer.markWalBatchEnd(1); bufferInsert(buffer, "b", "1234", 1); + buffer.markWalBatchEnd(2); bufferInsert(buffer, "c", "1234", 2); + buffer.markWalBatchEnd(3); PreparedFlush preparedFlush = buffer.prepareFlush(3); List segments = preparedFlush.split(10, Integer.MAX_VALUE); - // Adding the second entry would exceed the limit, while the last two entries exactly fit. + // The first two complete batches cross the byte budget and close the first segment. assertThat(segments).hasSize(2); - assertThat(segments.get(0).entries()).hasSize(1); - assertThat(segments.get(0).exclusiveUpToLogSequenceNumber()).isEqualTo(1); - assertThat(segments.get(0).rowCountDiff()).isEqualTo(1); - assertThat(segments.get(1).entries()).hasSize(2); + assertThat(segments.get(0).entries()).hasSize(2); + assertThat(segments.get(0).exclusiveUpToLogSequenceNumber()).isEqualTo(2); + assertThat(segments.get(0).rowCountDiff()).isEqualTo(2); + assertThat(segments.get(1).entries()).hasSize(1); assertThat(segments.get(1).exclusiveUpToLogSequenceNumber()).isEqualTo(3); - assertThat(segments.get(1).rowCountDiff()).isEqualTo(2); + assertThat(segments.get(1).rowCountDiff()).isEqualTo(1); } @Test @@ -286,7 +291,9 @@ void testSplitPreparedFlushWithOversizedEntry() { KvPreWriteBuffer buffer = new KvPreWriteBuffer(TestingMetricGroups.TABLET_SERVER_METRICS); // The first entry is larger than the byte limit and must remain a non-empty singleton. bufferInsert(buffer, "a", "1234567890", 0); + buffer.markWalBatchEnd(1); bufferInsert(buffer, "b", "123", 1); + buffer.markWalBatchEnd(2); PreparedFlush preparedFlush = buffer.prepareFlush(2); List segments = preparedFlush.split(10, Integer.MAX_VALUE); @@ -301,11 +308,15 @@ void testSplitPreparedFlushWithOversizedEntry() { @Test void testSplitPreparedFlushUsesFirstReachedLimit() { KvPreWriteBuffer buffer = new KvPreWriteBuffer(TestingMetricGroups.TABLET_SERVER_METRICS); - // Entry payload sizes are 2, 2, 6, and 5 bytes. + // Entry payload sizes are 2, 2, 11, and 5 bytes. bufferInsert(buffer, "a", "x", 0); + buffer.markWalBatchEnd(1); bufferInsert(buffer, "b", "y", 1); - bufferInsert(buffer, "c", "12345", 2); + buffer.markWalBatchEnd(2); + bufferInsert(buffer, "c", "1234567890", 2); + buffer.markWalBatchEnd(3); bufferInsert(buffer, "d", "1234", 3); + buffer.markWalBatchEnd(4); PreparedFlush preparedFlush = buffer.prepareFlush(4); List segments = preparedFlush.split(10, 2); @@ -325,28 +336,31 @@ void testCompletePrefixSegmentsAndAbortRest() { KvPreWriteBuffer buffer = new KvPreWriteBuffer(TestingMetricGroups.TABLET_SERVER_METRICS); // Entry payload sizes are 6, 5, 5, and 6 bytes. bufferInsert(buffer, "a", "12345", 0); + buffer.markWalBatchEnd(1); bufferInsert(buffer, "b", "1234", 1); + buffer.markWalBatchEnd(2); bufferInsert(buffer, "c", "1234", 2); + buffer.markWalBatchEnd(3); bufferInsert(buffer, "d", "12345", 3); + buffer.markWalBatchEnd(4); PreparedFlush preparedFlush = buffer.prepareFlush(4); List segments = preparedFlush.split(10, Integer.MAX_VALUE); - assertThat(segments).hasSize(3); + assertThat(segments).hasSize(2); // Model a storage rejection after the first segment landed: complete the written prefix, // abort the rest. - assertThat(buffer.completeFlush(segments.get(0))).isEqualTo(1); + assertThat(buffer.completeFlush(segments.get(0))).isEqualTo(2); buffer.abortFlush(segments.get(1)); - buffer.abortFlush(segments.get(2)); // The completed entries are gone; the aborted ones are ACTIVE again and can be prepared // by the retry, which must cover exactly the remaining range. - assertThat(buffer.pendingFlushBytes()).isEqualTo(16); - assertThat(buffer.getAllKvEntries()).hasSize(3); + assertThat(buffer.pendingFlushBytes()).isEqualTo(11); + assertThat(buffer.getAllKvEntries()).hasSize(2); PreparedFlush retry = buffer.prepareFlush(4); - assertThat(retry.entries()).hasSize(3); - assertThat(retry.entries().get(0).getLogSequenceNumber()).isEqualTo(1); - assertThat(retry.rowCountDiff()).isEqualTo(3); + assertThat(retry.entries()).hasSize(2); + assertThat(retry.entries().get(0).getLogSequenceNumber()).isEqualTo(2); + assertThat(retry.rowCountDiff()).isEqualTo(2); } private static void bufferInsert( diff --git a/fluss-server/src/test/java/org/apache/fluss/server/replica/KvReplicaRestoreITCase.java b/fluss-server/src/test/java/org/apache/fluss/server/replica/KvReplicaRestoreITCase.java index c38d39a834a..ed938ea5a4d 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/replica/KvReplicaRestoreITCase.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/replica/KvReplicaRestoreITCase.java @@ -135,9 +135,14 @@ void testRestore() throws Exception { // we delete the kv dir to make flush fail FileUtils.deleteDirectory(replica.getKvTablet().getKvTabletDir()); - // should fail + // Append the batch that must be recovered from WAL. putRecordBatch(tableBucket, leaderServer, toKvRecordBatch(records)); + // A complete WAL batch is one native write and may remain in the active memtable. + // Further writes rotate the memtable and surface the deleted directory as an I/O error. + KvRecordBatch triggerBatch = + genKvRecordBatch(new Object[] {recordsNum - 1, "k" + (recordsNum - 1)}); + // wait for the replica to restore in another server AtomicInteger newLeaderServer = new AtomicInteger(-1); waitUntil( @@ -147,6 +152,7 @@ void testRestore() throws Exception { newLeaderServer.set(restoreServer); return true; } + putRecordBatch(tableBucket, leaderServer, triggerBatch); return false; }, Duration.ofMinutes(2),