[server] Keep async KV flushes on WAL batch boundaries

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;

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<PreparedFlush> 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++;

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.
+     *
+     * <p>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}.
-         *
-         * <p>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<PreparedFlush> 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<PreparedFlush> split(long targetBytesPerSegment, int targetEntriesPerSegment) {
+            checkArgument(targetEntriesPerSegment > 0, "targetEntriesPerSegment must be positive.");
             List<PreparedFlush> 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(