[server] add support of KvPreWriteBuffer backpressure

fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java

@@ -783,6 +783,31 @@ public class ConfigOptions {
                                     + ". "
                                     + "This option allows to allocate memory in batches to have better CPU-cached friendliness due to contiguous segments.");
 
+    public static final ConfigOption<Double>
+            SERVER_KV_PRE_WRITE_BUFFER_MEMORY_HIGH_WATERMARK_RATIO =
+                    key("server.kv.pre-write-buffer.memory-high-watermark-ratio")
+                            .doubleType()
+                            .defaultValue(0.60)
+                            .withDescription(
+                                    "The high watermark, as a ratio of the TabletServer JVM maximum heap size, "
+                                            + "for the memory shared by all KV pre-write buffers. KV writes "
+                                            + "are rejected with a retriable error when memory usage reaches this "
+                                            + "watermark. The valid range is "
+                                            + "(server.kv.pre-write-buffer.memory-low-watermark-ratio, 1.0]. "
+                                            + "The default value is 0.20.");
+
+    public static final ConfigOption<Double> SERVER_KV_PRE_WRITE_BUFFER_MEMORY_LOW_WATERMARK_RATIO =
+            key("server.kv.pre-write-buffer.memory-low-watermark-ratio")
+                    .doubleType()
+                    .defaultValue(0.50)
+                    .withDescription(
+                            "The low watermark, as a ratio of the TabletServer JVM maximum heap size, "
+                                    + "for the memory shared by all KV pre-write buffers. KV writes "
+                                    + "resume when memory usage reaches or drops below this watermark. "
+                                    + "The valid range is [0.0, "
+                                    + "server.kv.pre-write-buffer.memory-high-watermark-ratio). "
+                                    + "The default value is 0.15.");
+
     public static final ConfigOption<Duration> SERVER_BUFFER_POOL_WAIT_TIMEOUT =
             key("server.buffer.wait-timeout")
                     .durationType()

fluss-common/src/main/java/org/apache/fluss/config/FlussConfigUtils.java

@@ -133,6 +133,33 @@ public static void validateTabletConfigs(Configuration conf) {
         validMinValue(ConfigOptions.TABLET_SERVER_ID, serverId.get(), 0);
     }
 
+    private static void validateKvPreWriteBufferMemoryWatermarks(Configuration conf) {
+        double highWatermarkRatio =
+                conf.get(ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_HIGH_WATERMARK_RATIO);
+        double lowWatermarkRatio =
+                conf.get(ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_LOW_WATERMARK_RATIO);
+        if (!(highWatermarkRatio > 0.0 && highWatermarkRatio <= 1.0)) {
+            throw new IllegalConfigurationException(
+                    String.format(
+                            "Invalid configuration for %s: %s. It must be within (0.0, 1.0].",
+                            ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_HIGH_WATERMARK_RATIO
+                                    .key(),
+                            highWatermarkRatio));
+        }
+        if (!(lowWatermarkRatio >= 0.0 && lowWatermarkRatio < highWatermarkRatio)) {
+            throw new IllegalConfigurationException(
+                    String.format(
+                            "Invalid configuration for %s: %s. It must be within [0.0, %s), "
+                                    + "where %s is the configured high watermark ratio.",
+                            ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_LOW_WATERMARK_RATIO
+                                    .key(),
+                            lowWatermarkRatio,
+                            highWatermarkRatio,
+                            ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_HIGH_WATERMARK_RATIO
+                                    .key()));
+        }
+    }
+
     public static void validateRemoteDataDirs(Configuration conf) {
         // Validate remote.data.dirs
         List<String> remoteDataDirs = conf.get(ConfigOptions.REMOTE_DATA_DIRS);
@@ -228,6 +255,8 @@ protected static void validateServerConfigs(Configuration conf) {
                             "Invalid configuration for %s, it must be less than or equal %d bytes.",
                             ConfigOptions.LOG_SEGMENT_FILE_SIZE.key(), Integer.MAX_VALUE));
         }
+
+        validateKvPreWriteBufferMemoryWatermarks(conf);
     }
 
     private static void validMinValue(

fluss-common/src/main/java/org/apache/fluss/exception/InsufficientKvPreWriteBufferException.java

@@ -0,0 +1,43 @@
+/*
+ * 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.exception;
+
+import org.apache.fluss.annotation.PublicEvolving;
+
+/**
+ * The tablet server rejected a KV write because its global pre-write buffer high watermark was
+ * reached.
+ *
+ * <p>The client may retry the write after a backoff. The rejected write has not been appended to
+ * the log.
+ *
+ * @since 1.0
+ */
+@PublicEvolving
+public class InsufficientKvPreWriteBufferException extends RetriableException {
+
+    private static final long serialVersionUID = 1L;
+
+    public InsufficientKvPreWriteBufferException(String message) {
+        super(message);
+    }
+
+    public InsufficientKvPreWriteBufferException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java

@@ -158,6 +158,8 @@ public class MetricNames {
             "preWriteBufferTruncateAsDuplicatedPerSecond";
     public static final String KV_PRE_WRITE_BUFFER_TRUNCATE_AS_ERROR_RATE =
             "preWriteBufferTruncateAsErrorPerSecond";
+    public static final String KV_PRE_WRITE_BUFFER_MEMORY_USAGE = "preWriteBufferMemoryUsage";
+    public static final String KV_PRE_WRITE_BUFFER_UNDER_PRESSURE = "preWriteBufferUnderPressure";
 
     // --------------------------------------------------------------------------------------------
     // RocksDB metrics

fluss-common/src/test/java/org/apache/fluss/config/FlussConfigUtilsTest.java

@@ -191,6 +191,32 @@ void testValidateTabletConfigs() {
                 .hasMessageContaining("it must be greater than or equal 0");
     }
 
+    @Test
+    void testValidateKvPreWriteBufferMemoryWatermarks() {
+        Configuration conf = new Configuration();
+        conf.set(ConfigOptions.REMOTE_DATA_DIR, "s3://bucket/path");
+        conf.set(ConfigOptions.TABLET_SERVER_ID, 0);
+        validateTabletConfigs(conf);
+
+        conf.set(ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_HIGH_WATERMARK_RATIO, 0.0);
+        assertThatThrownBy(() -> validateTabletConfigs(conf))
+                .isInstanceOf(IllegalConfigurationException.class)
+                .hasMessageContaining(
+                        ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_HIGH_WATERMARK_RATIO.key())
+                .hasMessageContaining("(0.0, 1.0]");
+
+        conf.set(ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_HIGH_WATERMARK_RATIO, 0.5);
+        conf.set(ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_LOW_WATERMARK_RATIO, 0.5);
+        assertThatThrownBy(() -> validateTabletConfigs(conf))
+                .isInstanceOf(IllegalConfigurationException.class)
+                .hasMessageContaining(
+                        ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_LOW_WATERMARK_RATIO.key())
+                .hasMessageContaining("[0.0, 0.5)");
+
+        conf.set(ConfigOptions.SERVER_KV_PRE_WRITE_BUFFER_MEMORY_LOW_WATERMARK_RATIO, 0.4);
+        validateTabletConfigs(conf);
+    }
+
     @Test
     void testValidateLogRetentionCheckInterval() {
         assertThat(ConfigOptions.LOG_RETENTION_CHECK_INTERVAL.key())

fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/Errors.java

@@ -33,6 +33,7 @@
 import org.apache.fluss.exception.FencedTieringEpochException;
 import org.apache.fluss.exception.IneligibleReplicaException;
 import org.apache.fluss.exception.InsufficientKvLeaderReplicaCapacityException;
+import org.apache.fluss.exception.InsufficientKvPreWriteBufferException;
 import org.apache.fluss.exception.InvalidAlterTableException;
 import org.apache.fluss.exception.InvalidColumnProjectionException;
 import org.apache.fluss.exception.InvalidConfigException;
@@ -275,7 +276,11 @@ public enum Errors {
     INSUFFICIENT_KV_LEADER_REPLICA_CAPACITY(
             71,
             "The cluster does not have enough KV leader replica capacity.",
-            InsufficientKvLeaderReplicaCapacityException::new);
+            InsufficientKvLeaderReplicaCapacityException::new),
+    INSUFFICIENT_KV_PRE_WRITE_BUFFER(
+            72,
+            "The tablet server KV pre-write buffer high watermark has been reached.",
+            InsufficientKvPreWriteBufferException::new);
 
     private static final Logger LOG = LoggerFactory.getLogger(Errors.class);
 

fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiErrorTest.java

@@ -17,6 +17,7 @@
 
 package org.apache.fluss.rpc.protocol;
 
+import org.apache.fluss.exception.InsufficientKvPreWriteBufferException;
 import org.apache.fluss.exception.NotEnoughReplicasException;
 import org.apache.fluss.exception.TimeoutException;
 import org.apache.fluss.exception.UnknownTableOrBucketException;
@@ -63,6 +64,13 @@ private static Collection<Arguments> parameters() {
                         Errors.NOT_ENOUGH_REPLICAS_EXCEPTION,
                         notEnoughReplicasErrorMsg));
 
+        String preWriteBufferFullMessage = "KV pre-write buffer is full.";
+        arguments.add(
+                Arguments.of(
+                        new InsufficientKvPreWriteBufferException(preWriteBufferFullMessage),
+                        Errors.INSUFFICIENT_KV_PRE_WRITE_BUFFER,
+                        preWriteBufferFullMessage));
+
         // avoid populating the error message if it's a generic one
         arguments.add(
                 Arguments.of(

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

@@ -38,6 +38,7 @@
 import org.apache.fluss.server.TabletManagerBase;
 import org.apache.fluss.server.kv.autoinc.AutoIncrementManager;
 import org.apache.fluss.server.kv.autoinc.ZkSequenceGeneratorFactory;
+import org.apache.fluss.server.kv.prewrite.KvPreWriteBufferMemoryManager;
 import org.apache.fluss.server.kv.rowmerger.RowMerger;
 import org.apache.fluss.server.log.LogManager;
 import org.apache.fluss.server.log.LogTablet;
@@ -115,6 +116,9 @@ public static RateLimiter getDefaultRateLimiter() {
 
     private final TabletServerMetricGroup serverMetricGroup;
 
+    /** Tablet-server-wide memory quota shared by all KV pre-write buffers. */
+    private final KvPreWriteBufferMemoryManager preWriteBufferMemoryManager;
+
     private final ZooKeeperClient zkClient;
 
     private final Map<TableBucket, KvTablet> currentKvs = new ConcurrentHashMap<>();
@@ -159,6 +163,24 @@ private KvManager(
         this.remoteKvDir = FlussPaths.remoteKvDir(conf);
         this.remoteFileSystem = remoteKvDir.getFileSystem();
         this.serverMetricGroup = tabletServerMetricGroup;
+        long maxHeapMemory = Runtime.getRuntime().maxMemory();
+        long highWatermarkBytes =
+                new MemorySize(maxHeapMemory)
+                        .multiply(
+                                conf.get(
+                                        ConfigOptions
+                                                .SERVER_KV_PRE_WRITE_BUFFER_MEMORY_HIGH_WATERMARK_RATIO))
+                        .getBytes();
+        long lowWatermarkBytes =
+                new MemorySize(maxHeapMemory)
+                        .multiply(
+                                conf.get(
+                                        ConfigOptions
+                                                .SERVER_KV_PRE_WRITE_BUFFER_MEMORY_LOW_WATERMARK_RATIO))
+                        .getBytes();
+        this.preWriteBufferMemoryManager =
+                new KvPreWriteBufferMemoryManager(highWatermarkBytes, lowWatermarkBytes);
+        tabletServerMetricGroup.registerKvPreWriteBufferMemoryManager(preWriteBufferMemoryManager);
         this.sharedRocksDBRateLimiter = createSharedRateLimiter(conf);
         this.currentSharedRateLimitBytesPerSec =
                 conf.get(ConfigOptions.KV_SHARED_RATE_LIMITER_BYTES_PER_SEC).getBytes();
@@ -278,7 +300,8 @@ public KvTablet getOrCreateKv(
                                     schemaGetter,
                                     tableConfig.getChangelogImage(),
                                     sharedRocksDBRateLimiter,
-                                    autoIncrementManager);
+                                    autoIncrementManager,
+                                    preWriteBufferMemoryManager);
                     currentKvs.put(tableBucket, tablet);
 
                     LOG.info(
@@ -396,7 +419,8 @@ public KvTablet loadKv(File tabletDir, SchemaGetter schemaGetter) throws Excepti
                         schemaGetter,
                         tableConfig.getChangelogImage(),
                         sharedRocksDBRateLimiter,
-                        autoIncrementManager);
+                        autoIncrementManager,
+                        preWriteBufferMemoryManager);
         if (this.currentKvs.containsKey(tableBucket)) {
             throw new IllegalStateException(
                     String.format(

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

@@ -52,6 +52,7 @@
 import org.apache.fluss.server.kv.autoinc.AutoIncrementUpdater;
 import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer;
 import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.TruncateReason;
+import org.apache.fluss.server.kv.prewrite.KvPreWriteBufferMemoryManager;
 import org.apache.fluss.server.kv.rocksdb.RocksDBKv;
 import org.apache.fluss.server.kv.rocksdb.RocksDBKvBuilder;
 import org.apache.fluss.server.kv.rocksdb.RocksDBResourceContainer;
@@ -169,15 +170,18 @@ private KvTablet(
             SchemaGetter schemaGetter,
             ChangelogImage changelogImage,
             @Nullable RocksDBStatistics rocksDBStatistics,
-            AutoIncrementManager autoIncrementManager) {
+            AutoIncrementManager autoIncrementManager,
+            KvPreWriteBufferMemoryManager preWriteBufferMemoryManager) {
         this.physicalPath = physicalPath;
         this.tableBucket = tableBucket;
         this.logTablet = logTablet;
         this.kvTabletDir = kvTabletDir;
         this.rocksDBKv = rocksDBKv;
         this.writeBatchSize = writeBatchSize;
         this.serverMetricGroup = serverMetricGroup;
-        this.kvPreWriteBuffer = new KvPreWriteBuffer(createKvBatchWriter(), serverMetricGroup);
+        this.kvPreWriteBuffer =
+                new KvPreWriteBuffer(
+                        createKvBatchWriter(), serverMetricGroup, preWriteBufferMemoryManager);
         this.logFormat = logFormat;
         this.arrowWriterProvider = new ArrowWriterPool(arrowBufferAllocator);
         this.memorySegmentPool = memorySegmentPool;
@@ -195,6 +199,11 @@ private KvTablet(
         this.rowCount = changelogImage == ChangelogImage.WAL ? ROW_COUNT_DISABLED : 0L;
     }
 
+    /**
+     * Creates a KV tablet backed by the tablet-server-wide pre-write buffer memory manager.
+     *
+     * @param preWriteBufferMemoryManager global pre-write buffer memory quota
+     */
     public static KvTablet create(
             PhysicalTablePath tablePath,
             TableBucket tableBucket,
@@ -210,7 +219,8 @@ public static KvTablet create(
             SchemaGetter schemaGetter,
             ChangelogImage changelogImage,
             RateLimiter sharedRateLimiter,
-            AutoIncrementManager autoIncrementManager)
+            AutoIncrementManager autoIncrementManager,
+            KvPreWriteBufferMemoryManager preWriteBufferMemoryManager)
             throws IOException {
         RocksDBKv kv = buildRocksDBKv(serverConf, kvTabletDir, sharedRateLimiter);
 
@@ -243,7 +253,8 @@ public static KvTablet create(
                 schemaGetter,
                 changelogImage,
                 rocksDBStatistics,
-                autoIncrementManager);
+                autoIncrementManager,
+                preWriteBufferMemoryManager);
     }
 
     private static RocksDBKv buildRocksDBKv(
@@ -668,8 +679,6 @@ private WalBuilder createWalBuilder(int schemaId, RowType rowType) throws Except
     }
 
     public void flush(long exclusiveUpToLogOffset, FatalErrorHandler fatalErrorHandler) {
-        // todo: need to introduce a backpressure mechanism
-        // to avoid too much records in kvPreWriteBuffer
         inWriteLock(
                 kvLock,
                 () -> {
@@ -851,10 +860,14 @@ public void close() throws Exception {
                     }
                     // Note: RocksDB metrics lifecycle is managed by TableMetricGroup
                     // No need to close it here
-                    if (rocksDBKv != null) {
-                        rocksDBKv.close();
+                    try {
+                        kvPreWriteBuffer.close();
+                    } finally {
+                        if (rocksDBKv != null) {
+                            rocksDBKv.close();
+                        }
+                        isClosed = true;
                     }
-                    isClosed = true;
                 });
     }