diff --git a/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java b/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java index d9d79502d8..a8c3826191 100644 --- a/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java +++ b/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 + 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 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 SERVER_BUFFER_POOL_WAIT_TIMEOUT = key("server.buffer.wait-timeout") .durationType() diff --git a/fluss-common/src/main/java/org/apache/fluss/config/FlussConfigUtils.java b/fluss-common/src/main/java/org/apache/fluss/config/FlussConfigUtils.java index 71bd9702eb..f280e19979 100644 --- a/fluss-common/src/main/java/org/apache/fluss/config/FlussConfigUtils.java +++ b/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 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( diff --git a/fluss-common/src/main/java/org/apache/fluss/exception/InsufficientKvPreWriteBufferException.java b/fluss-common/src/main/java/org/apache/fluss/exception/InsufficientKvPreWriteBufferException.java new file mode 100644 index 0000000000..1221dd4538 --- /dev/null +++ b/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. + * + *

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); + } +} diff --git a/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java b/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java index 5b24b5b673..f81c84eca1 100644 --- a/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java +++ b/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 diff --git a/fluss-common/src/test/java/org/apache/fluss/config/FlussConfigUtilsTest.java b/fluss-common/src/test/java/org/apache/fluss/config/FlussConfigUtilsTest.java index 2b9e815ee7..d2e183d5d3 100644 --- a/fluss-common/src/test/java/org/apache/fluss/config/FlussConfigUtilsTest.java +++ b/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()) diff --git a/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/Errors.java b/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/Errors.java index 9a46b04604..34dbe4ccd0 100644 --- a/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/Errors.java +++ b/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); diff --git a/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiErrorTest.java b/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiErrorTest.java index f2b4cfd3a3..73b3fac705 100644 --- a/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiErrorTest.java +++ b/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 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( diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvManager.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvManager.java index 6a5c6079c9..152369c2aa 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvManager.java +++ b/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 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( 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 da6b518266..16bbd08433 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 @@ -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,7 +170,8 @@ private KvTablet( SchemaGetter schemaGetter, ChangelogImage changelogImage, @Nullable RocksDBStatistics rocksDBStatistics, - AutoIncrementManager autoIncrementManager) { + AutoIncrementManager autoIncrementManager, + KvPreWriteBufferMemoryManager preWriteBufferMemoryManager) { this.physicalPath = physicalPath; this.tableBucket = tableBucket; this.logTablet = logTablet; @@ -177,7 +179,9 @@ private KvTablet( 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; }); } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvMemoryConsumer.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvMemoryConsumer.java new file mode 100644 index 0000000000..94f65a2808 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvMemoryConsumer.java @@ -0,0 +1,78 @@ +/* + * 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.prewrite; + +import org.apache.fluss.annotation.Internal; + +import javax.annotation.concurrent.NotThreadSafe; + +import static org.apache.fluss.utils.Preconditions.checkArgument; +import static org.apache.fluss.utils.Preconditions.checkNotNull; +import static org.apache.fluss.utils.Preconditions.checkState; + +/** + * A consumer of the TabletServer-wide KV memory pool. + * + *

The global manager owns admission across consumers, while each consumer tracks its own + * acquired bytes. This prevents one consumer from releasing memory acquired by another and provides + * an idempotent way to release all memory during cleanup. + */ +@Internal +@NotThreadSafe +abstract class KvMemoryConsumer { + + private final KvPreWriteBufferMemoryManager memoryManager; + private long acquiredBytes; + + protected KvMemoryConsumer(KvPreWriteBufferMemoryManager memoryManager) { + this.memoryManager = checkNotNull(memoryManager, "Memory manager must not be null."); + } + + protected final boolean tryAcquireMemory(long bytes) { + checkArgument(bytes >= 0, "The number of bytes to acquire must not be negative."); + if (!memoryManager.tryReserve(bytes)) { + return false; + } + acquiredBytes += bytes; + return true; + } + + protected final void freeMemory(long bytes) { + checkArgument(bytes >= 0, "The number of bytes to free must not be negative."); + checkState( + bytes <= acquiredBytes, + "Cannot free %s bytes when this consumer has only acquired %s bytes.", + bytes, + acquiredBytes); + if (bytes > 0) { + memoryManager.release(bytes); + acquiredBytes -= bytes; + } + } + + protected final void freeAllMemory() { + if (acquiredBytes > 0) { + memoryManager.release(acquiredBytes); + acquiredBytes = 0; + } + } + + protected final KvPreWriteBufferMemoryManager memoryManager() { + return memoryManager; + } +} 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 e14d09eacf..895cb3c297 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 @@ -18,6 +18,8 @@ package org.apache.fluss.server.kv.prewrite; import org.apache.fluss.annotation.VisibleForTesting; +import org.apache.fluss.exception.InsufficientKvPreWriteBufferException; +import org.apache.fluss.exception.RecordTooLargeException; import org.apache.fluss.memory.MemorySegment; import org.apache.fluss.metrics.Counter; import org.apache.fluss.record.ChangeType; @@ -85,7 +87,13 @@ * head to tail, it will stop flush. */ @NotThreadSafe -public class KvPreWriteBuffer implements AutoCloseable { +public class KvPreWriteBuffer extends KvMemoryConsumer implements AutoCloseable { + + // A conservative estimate of the heap retained by a KvEntry, its Key and Value wrappers, the + // linked-list node, and associated object/array headers. The byte-array payloads are accounted + // for separately. + private static final long ENTRY_OVERHEAD_BYTES = 128L; + private final KvBatchWriter kvBatchWriter; // a mapping from the key to the kv-entry @@ -102,7 +110,10 @@ public class KvPreWriteBuffer implements AutoCloseable { private long maxLogSequenceNumber = -1; public KvPreWriteBuffer( - KvBatchWriter kvBatchWriter, TabletServerMetricGroup serverMetricGroup) { + KvBatchWriter kvBatchWriter, + TabletServerMetricGroup serverMetricGroup, + KvPreWriteBufferMemoryManager memoryManager) { + super(memoryManager); this.kvBatchWriter = kvBatchWriter; truncateAsDuplicatedCount = serverMetricGroup.kvTruncateAsDuplicatedCount(); @@ -146,6 +157,25 @@ private void doPut(ChangeType changeType, Key key, Value value, long lsn) { + lsn); } + long memorySize = estimateMemorySize(key, value); + if (memorySize > memoryManager().highWatermarkBytes()) { + throw new RecordTooLargeException( + String.format( + "The KV entry requires %s bytes, which exceeds the pre-write buffer " + + "high watermark of %s bytes.", + memorySize, memoryManager().highWatermarkBytes())); + } + + if (!tryAcquireMemory(memorySize)) { + throw new InsufficientKvPreWriteBufferException( + String.format( + "Failed to reserve %s bytes in the KV pre-write buffer " + + "(used: %s, high watermark: %s).", + memorySize, + memoryManager().usedBytes(), + memoryManager().highWatermarkBytes())); + } + // create the kv entry with previous pointer if exists, and put the new entry to the map KvEntry kvEntry = kvEntryMap.compute( @@ -194,6 +224,7 @@ public void truncateTo(long targetLogSequenceNumber, TruncateReason truncateReas break; } descIter.remove(); + freeMemory(estimateMemorySize(entry.key, entry.value)); boolean removed = kvEntryMap.remove(entry.getKey(), entry); // if the latest entry is removed, we need to rollback the previous entry to the map if (removed && entry.previousEntry != null) { @@ -247,6 +278,7 @@ public int flush(long exclusiveUpToLogSequenceNumber) throws IOException { // can remove it from the map. Although it's not a must to remove from the map, // we remove it to reduce the memory usage kvEntryMap.remove(entry.getKey(), entry); + freeMemory(estimateMemorySize(entry.key, entry.value)); } // flush to underlying kv tablet if (flushedCount > 0) { @@ -273,8 +305,12 @@ public long getMaxLSN() { @Override public void close() throws Exception { - if (kvBatchWriter != null) { - kvBatchWriter.close(); + try { + if (kvBatchWriter != null) { + kvBatchWriter.close(); + } + } finally { + freeAllMemory(); } } @@ -286,6 +322,13 @@ public Counter getTruncateAsErrorCount() { return truncateAsErrorCount; } + private static long estimateMemorySize(Key key, Value value) { + byte[] valueBytes = value.get(); + return ENTRY_OVERHEAD_BYTES + + key.get().length + + (valueBytes == null ? 0 : valueBytes.length); + } + // ------------------------------------------------------------------------------------------- // Inner classes // ------------------------------------------------------------------------------------------- diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferMemoryManager.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferMemoryManager.java new file mode 100644 index 0000000000..ded2416274 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferMemoryManager.java @@ -0,0 +1,132 @@ +/* + * 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.prewrite; + +import org.apache.fluss.annotation.Internal; + +import javax.annotation.concurrent.ThreadSafe; + +import static org.apache.fluss.utils.Preconditions.checkArgument; +import static org.apache.fluss.utils.Preconditions.checkState; + +/** + * Tablet-server-wide memory watermarks shared by all KV pre-write buffers. + * + *

Writes are admitted until memory usage reaches the high watermark. They remain under pressure + * until usage falls to the low watermark. + * + *

The watermarks are intentionally global rather than per bucket. All reservations and releases + * are synchronized and safe to call concurrently from different KV tablets. + */ +@Internal +@ThreadSafe +public final class KvPreWriteBufferMemoryManager { + + private final long highWatermarkBytes; + private final long lowWatermarkBytes; + + private long usedBytes; + private boolean underPressure; + + /** + * Creates global pre-write buffer memory watermarks. + * + * @param highWatermarkBytes usage at which writes are rejected + * @param lowWatermarkBytes usage at or below which writes resume + */ + public KvPreWriteBufferMemoryManager(long highWatermarkBytes, long lowWatermarkBytes) { + checkArgument(highWatermarkBytes > 0, "High watermark must be greater than 0."); + checkArgument(lowWatermarkBytes >= 0, "Low watermark must not be negative."); + checkArgument( + lowWatermarkBytes < highWatermarkBytes, + "Low watermark must be less than high watermark."); + this.highWatermarkBytes = highWatermarkBytes; + this.lowWatermarkBytes = lowWatermarkBytes; + } + + /** + * Attempts to reserve bytes from the global memory pool. + * + *

Reservations are rejected while the manager is under pressure or when the reservation + * would exceed the high watermark. + * + * @param bytes number of bytes to reserve + * @return whether the reservation succeeded + */ + public synchronized boolean tryReserve(long bytes) { + checkArgument(bytes >= 0, "The number of bytes to reserve must not be negative."); + if (bytes == 0) { + return true; + } + + if (underPressure) { + return false; + } + + if (bytes > highWatermarkBytes - usedBytes) { + underPressure = usedBytes > lowWatermarkBytes; + return false; + } + + usedBytes += bytes; + underPressure = usedBytes >= highWatermarkBytes; + return true; + } + + /** + * Releases bytes previously reserved from this manager. + * + * @param bytes number of reserved bytes to release + */ + public synchronized void release(long bytes) { + checkArgument(bytes >= 0, "The number of bytes to release must not be negative."); + if (bytes == 0) { + return; + } + + checkState( + bytes <= usedBytes, + "Cannot release %s bytes when only %s bytes are reserved.", + bytes, + usedBytes); + usedBytes -= bytes; + if (usedBytes <= lowWatermarkBytes) { + underPressure = false; + } + } + + /** Returns the currently reserved bytes across all KV pre-write buffers. */ + public synchronized long usedBytes() { + return usedBytes; + } + + /** Returns the usage at which writes are rejected. */ + public long highWatermarkBytes() { + return highWatermarkBytes; + } + + /** Returns the usage at or below which writes resume. */ + public long lowWatermarkBytes() { + return lowWatermarkBytes; + } + + /** Returns whether new reservations are currently under pressure. */ + public synchronized boolean isUnderPressure() { + return underPressure; + } +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/metrics/group/TabletServerMetricGroup.java b/fluss-server/src/main/java/org/apache/fluss/server/metrics/group/TabletServerMetricGroup.java index 22215bc6de..dfad7d84c8 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/metrics/group/TabletServerMetricGroup.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/metrics/group/TabletServerMetricGroup.java @@ -30,6 +30,7 @@ import org.apache.fluss.metrics.ThreadSafeSimpleCounter; import org.apache.fluss.metrics.groups.AbstractMetricGroup; import org.apache.fluss.metrics.registry.MetricRegistry; +import org.apache.fluss.server.kv.prewrite.KvPreWriteBufferMemoryManager; import org.apache.fluss.server.kv.rocksdb.RocksDBStatistics; import java.util.Map; @@ -227,6 +228,18 @@ public Counter kvTruncateAsErrorCount() { return kvTruncateAsErrorCount; } + /** + * Registers TabletServer-wide metrics for the global KV pre-write buffer memory quota. + * + * @param memoryManager global memory manager shared by all KV pre-write buffers + */ + public void registerKvPreWriteBufferMemoryManager(KvPreWriteBufferMemoryManager memoryManager) { + gauge(MetricNames.KV_PRE_WRITE_BUFFER_MEMORY_USAGE, memoryManager::usedBytes); + gauge( + MetricNames.KV_PRE_WRITE_BUFFER_UNDER_PRESSURE, + () -> memoryManager.isUnderPressure() ? 1 : 0); + } + public Counter isrShrinks() { return isrShrinks; } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletMergeModeTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletMergeModeTest.java index 97ae1d9cc9..578dc3b409 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletMergeModeTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletMergeModeTest.java @@ -38,6 +38,7 @@ import org.apache.fluss.rpc.protocol.MergeMode; import org.apache.fluss.server.kv.autoinc.AutoIncrementManager; import org.apache.fluss.server.kv.autoinc.TestingSequenceGeneratorFactory; +import org.apache.fluss.server.kv.prewrite.KvPreWriteBufferMemoryManager; import org.apache.fluss.server.kv.rowmerger.RowMerger; import org.apache.fluss.server.log.FetchIsolation; import org.apache.fluss.server.log.LogTablet; @@ -160,7 +161,8 @@ void setUp() throws Exception { schemaGetter, tableConf.getChangelogImage(), KvManager.getDefaultRateLimiter(), - autoIncrementManager); + autoIncrementManager, + new KvPreWriteBufferMemoryManager(Long.MAX_VALUE, Long.MAX_VALUE - 1)); } @AfterEach diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletSchemaEvolutionTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletSchemaEvolutionTest.java index 1303521708..abcb830fcf 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletSchemaEvolutionTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/KvTabletSchemaEvolutionTest.java @@ -35,6 +35,7 @@ import org.apache.fluss.record.TestingSchemaGetter; import org.apache.fluss.server.kv.autoinc.AutoIncrementManager; import org.apache.fluss.server.kv.autoinc.TestingSequenceGeneratorFactory; +import org.apache.fluss.server.kv.prewrite.KvPreWriteBufferMemoryManager; import org.apache.fluss.server.kv.rowmerger.RowMerger; import org.apache.fluss.server.log.FetchIsolation; import org.apache.fluss.server.log.LogTablet; @@ -145,7 +146,8 @@ void setUp() throws Exception { schemaGetter, tableConf.getChangelogImage(), KvManager.getDefaultRateLimiter(), - autoIncrementManager); + autoIncrementManager, + new KvPreWriteBufferMemoryManager(Long.MAX_VALUE, Long.MAX_VALUE - 1)); } @AfterEach 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 3425fa75cf..5dd556d52f 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 @@ -20,6 +20,7 @@ import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; import org.apache.fluss.config.TableConfig; +import org.apache.fluss.exception.InsufficientKvPreWriteBufferException; import org.apache.fluss.exception.InvalidTableException; import org.apache.fluss.exception.InvalidTargetColumnException; import org.apache.fluss.exception.OutOfOrderSequenceException; @@ -55,6 +56,7 @@ import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.Key; import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.KvEntry; import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.Value; +import org.apache.fluss.server.kv.prewrite.KvPreWriteBufferMemoryManager; import org.apache.fluss.server.kv.rocksdb.RocksDBStatistics; import org.apache.fluss.server.kv.rowmerger.RowMerger; import org.apache.fluss.server.kv.scan.OpenScanResult; @@ -151,6 +153,15 @@ private void initLogTabletAndKvTablet(Schema schema, Map tableCo private void initLogTabletAndKvTablet( TablePath tablePath, Schema schema, Map tableConfig) throws Exception { + initLogTabletAndKvTablet(tablePath, schema, tableConfig, null); + } + + private void initLogTabletAndKvTablet( + TablePath tablePath, + Schema schema, + Map tableConfig, + @Nullable KvPreWriteBufferMemoryManager preWriteBufferMemoryManager) + throws Exception { PhysicalTablePath physicalTablePath = PhysicalTablePath.of(tablePath); schemaGetter = new TestingSchemaGetter(new SchemaInfo(schema, schemaId)); logTablet = createLogTablet(tempLogDir, 0L, physicalTablePath); @@ -162,7 +173,8 @@ private void initLogTabletAndKvTablet( logTablet, tmpKvDir, schemaGetter, - tableConfig); + tableConfig, + preWriteBufferMemoryManager); } private LogTablet createLogTablet(File tempLogDir, long tableId, PhysicalTablePath tablePath) @@ -193,6 +205,19 @@ private KvTablet createKvTablet( SchemaGetter schemaGetter, Map tableConfig) throws Exception { + return createKvTablet( + tablePath, tableBucket, logTablet, tmpKvDir, schemaGetter, tableConfig, null); + } + + private KvTablet createKvTablet( + PhysicalTablePath tablePath, + TableBucket tableBucket, + LogTablet logTablet, + File tmpKvDir, + SchemaGetter schemaGetter, + Map tableConfig, + @Nullable KvPreWriteBufferMemoryManager preWriteBufferMemoryManager) + throws Exception { TableConfig tableConf = new TableConfig(Configuration.fromMap(tableConfig)); RowMerger rowMerger = RowMerger.create(tableConf, KvFormat.COMPACTED, schemaGetter); AutoIncrementManager autoIncrementManager = @@ -201,7 +226,10 @@ private KvTablet createKvTablet( tablePath.getTablePath(), new TableConfig(new Configuration()), new TestingSequenceGeneratorFactory()); - + if (preWriteBufferMemoryManager == null) { + preWriteBufferMemoryManager = + new KvPreWriteBufferMemoryManager(Long.MAX_VALUE, Long.MAX_VALUE - 1); + } return KvTablet.create( tablePath, tableBucket, @@ -217,7 +245,35 @@ private KvTablet createKvTablet( schemaGetter, tableConf.getChangelogImage(), KvManager.getDefaultRateLimiter(), - autoIncrementManager); + autoIncrementManager, + preWriteBufferMemoryManager); + } + + @Test + void testClientWriteIsRejectedWhenPreWriteBufferIsFull() throws Exception { + KvPreWriteBufferMemoryManager memoryManager = + new KvPreWriteBufferMemoryManager(1_000L, 800L); + initLogTabletAndKvTablet( + TablePath.of("testDb", "quota_test"), + DATA1_SCHEMA_PK, + Collections.emptyMap(), + memoryManager); + + assertThat(memoryManager.tryReserve(950L)).isTrue(); + long logEndOffsetBeforePut = logTablet.localLogEndOffset(); + KvRecord record = kvRecordFactory.ofRecord("k1".getBytes(), new Object[] {1, "a"}); + KvRecordBatch batch = kvRecordBatchFactory.ofRecords(Collections.singletonList(record)); + + assertThatThrownBy(() -> kvTablet.putAsLeader(batch, null)) + .isInstanceOf(InsufficientKvPreWriteBufferException.class); + assertThat(logTablet.localLogEndOffset()).isEqualTo(logEndOffsetBeforePut); + assertThat(kvTablet.getKvPreWriteBuffer().getAllKvEntries()).isEmpty(); + assertThat(memoryManager.usedBytes()).isEqualTo(950L); + + kvTablet.close(); + assertThat(memoryManager.usedBytes()).isEqualTo(950L); + memoryManager.release(950L); + assertThat(memoryManager.usedBytes()).isZero(); } @Test diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferMemoryManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferMemoryManagerTest.java new file mode 100644 index 0000000000..0310d52ec3 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferMemoryManagerTest.java @@ -0,0 +1,131 @@ +/* + * 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.prewrite; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link KvPreWriteBufferMemoryManager}. */ +class KvPreWriteBufferMemoryManagerTest { + + @Test + void testClientPressureAndResumeHysteresis() { + KvPreWriteBufferMemoryManager manager = new KvPreWriteBufferMemoryManager(100, 80); + + assertThat(manager.highWatermarkBytes()).isEqualTo(100); + assertThat(manager.lowWatermarkBytes()).isEqualTo(80); + assertThat(manager.tryReserve(100)).isTrue(); + assertThat(manager.isUnderPressure()).isTrue(); + assertThat(manager.tryReserve(1)).isFalse(); + + manager.release(19); + assertThat(manager.usedBytes()).isEqualTo(81); + assertThat(manager.isUnderPressure()).isTrue(); + assertThat(manager.tryReserve(1)).isFalse(); + + manager.release(1); + assertThat(manager.isUnderPressure()).isFalse(); + assertThat(manager.tryReserve(20)).isTrue(); + assertThat(manager.usedBytes()).isEqualTo(100); + assertThat(manager.isUnderPressure()).isTrue(); + } + + @Test + void testClientReservationFailureDoesNotChangeUsage() { + KvPreWriteBufferMemoryManager manager = new KvPreWriteBufferMemoryManager(100, 80); + + assertThat(manager.tryReserve(101)).isFalse(); + assertThat(manager.usedBytes()).isZero(); + assertThat(manager.isUnderPressure()).isFalse(); + } + + @Test + void testRejectedLargeReservationEntersPressureOnlyAboveLowWatermark() { + KvPreWriteBufferMemoryManager manager = new KvPreWriteBufferMemoryManager(100, 80); + + assertThat(manager.tryReserve(80)).isTrue(); + assertThat(manager.tryReserve(21)).isFalse(); + assertThat(manager.isUnderPressure()).isFalse(); + + assertThat(manager.tryReserve(1)).isTrue(); + assertThat(manager.tryReserve(20)).isFalse(); + assertThat(manager.isUnderPressure()).isTrue(); + + manager.release(1); + assertThat(manager.usedBytes()).isEqualTo(80); + assertThat(manager.isUnderPressure()).isFalse(); + } + + @Test + void testInvalidReservationAndRelease() { + assertThatThrownBy(() -> new KvPreWriteBufferMemoryManager(0, 0)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> new KvPreWriteBufferMemoryManager(100, -1)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> new KvPreWriteBufferMemoryManager(100, 100)) + .isInstanceOf(IllegalArgumentException.class); + + KvPreWriteBufferMemoryManager manager = new KvPreWriteBufferMemoryManager(100, 80); + assertThatThrownBy(() -> manager.tryReserve(-1)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> manager.release(-1)).isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> manager.release(1)).isInstanceOf(IllegalStateException.class); + } + + @Test + void testConcurrentReservationsDoNotExceedLimit() throws Exception { + final int threadCount = 8; + final KvPreWriteBufferMemoryManager manager = + new KvPreWriteBufferMemoryManager(10_000, 8_000); + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + try { + List> reservationTasks = new ArrayList<>(); + for (int i = 0; i < threadCount; i++) { + reservationTasks.add( + () -> { + int reservations = 0; + while (manager.tryReserve(1)) { + reservations++; + } + return reservations; + }); + } + + List> results = executor.invokeAll(reservationTasks); + int successfulReservations = 0; + for (Future result : results) { + successfulReservations += result.get(); + } + + assertThat(successfulReservations).isEqualTo(10_000); + assertThat(manager.usedBytes()).isEqualTo(10_000); + assertThat(manager.isUnderPressure()).isTrue(); + } finally { + executor.shutdownNow(); + } + } +} 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 cb07c2d66b..00ca903cd0 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 @@ -17,6 +17,8 @@ package org.apache.fluss.server.kv.prewrite; +import org.apache.fluss.exception.InsufficientKvPreWriteBufferException; +import org.apache.fluss.exception.RecordTooLargeException; import org.apache.fluss.server.kv.KvBatchWriter; import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.TruncateReason; import org.apache.fluss.server.metrics.group.TestingMetricGroups; @@ -35,9 +37,7 @@ class KvPreWriteBufferTest { @Test void testIllegalLSN() { - KvPreWriteBuffer buffer = - new KvPreWriteBuffer( - new NopKvBatchWriter(), TestingMetricGroups.TABLET_SERVER_METRICS); + KvPreWriteBuffer buffer = createBuffer(new NopKvBatchWriter()); bufferInsert(buffer, "key1", "value1", 1); bufferDelete(buffer, "key1", 3); @@ -56,9 +56,7 @@ void testIllegalLSN() { @Test void testWriteAndFlush() throws Exception { - KvPreWriteBuffer buffer = - new KvPreWriteBuffer( - new NopKvBatchWriter(), TestingMetricGroups.TABLET_SERVER_METRICS); + KvPreWriteBuffer buffer = createBuffer(new NopKvBatchWriter()); int elementCount = 0; // put a series of kv entries @@ -139,9 +137,7 @@ void testWriteAndFlush() throws Exception { @Test void testTruncate() { - KvPreWriteBuffer buffer = - new KvPreWriteBuffer( - new NopKvBatchWriter(), TestingMetricGroups.TABLET_SERVER_METRICS); + KvPreWriteBuffer buffer = createBuffer(new NopKvBatchWriter()); int elementCount = 0; // put a series of kv entries @@ -194,9 +190,7 @@ void testTruncate() { @Test void testRowCount() throws IOException { - KvPreWriteBuffer buffer = - new KvPreWriteBuffer( - new NopKvBatchWriter(), TestingMetricGroups.TABLET_SERVER_METRICS); + KvPreWriteBuffer buffer = createBuffer(new NopKvBatchWriter()); int elementCount = 0; // put a series of kv entries @@ -238,11 +232,133 @@ void testRowCount() throws IOException { assertThat(buffer.flush(Long.MAX_VALUE)).isEqualTo(7); } + @Test + void testMemoryAccountingForMultipleVersions() throws Exception { + KvPreWriteBufferMemoryManager memoryManager = + new KvPreWriteBufferMemoryManager(10_000, 8_000); + KvPreWriteBuffer buffer = + new KvPreWriteBuffer( + new NopKvBatchWriter(), + TestingMetricGroups.TABLET_SERVER_METRICS, + memoryManager); + + buffer.insert(toKey("key"), "value-1".getBytes(), 0); + long firstEntryBytes = memoryManager.usedBytes(); + buffer.update(toKey("key"), "value-2".getBytes(), 1); + long secondEntryBytes = memoryManager.usedBytes() - firstEntryBytes; + buffer.delete(toKey("key"), 2); + long thirdEntryBytes = memoryManager.usedBytes() - firstEntryBytes - secondEntryBytes; + assertThat(memoryManager.usedBytes()) + .isEqualTo(firstEntryBytes + secondEntryBytes + thirdEntryBytes); + + buffer.flush(1); + assertThat(memoryManager.usedBytes()).isEqualTo(secondEntryBytes + thirdEntryBytes); + assertThat(buffer.getAllKvEntries()).hasSize(2); + + buffer.truncateTo(2, TruncateReason.ERROR); + assertThat(memoryManager.usedBytes()).isEqualTo(secondEntryBytes); + assertThat(getValue(buffer, "key")).isEqualTo("value-2"); + + buffer.flush(Long.MAX_VALUE); + assertThat(memoryManager.usedBytes()).isZero(); + assertThat(buffer.getAllKvEntries()).isEmpty(); + + buffer.close(); + assertThat(memoryManager.usedBytes()).isZero(); + } + + @Test + void testRejectedReservationDoesNotMutateBuffer() { + KvPreWriteBufferMemoryManager memoryManager = new KvPreWriteBufferMemoryManager(200, 160); + KvPreWriteBuffer buffer = + new KvPreWriteBuffer( + new NopKvBatchWriter(), + TestingMetricGroups.TABLET_SERVER_METRICS, + memoryManager); + buffer.insert(toKey("key"), "value".getBytes(), 0); + long retainedBytes = memoryManager.usedBytes(); + + assertThatThrownBy(() -> buffer.insert(toKey("key-2"), "value-2".getBytes(), 1)) + .isInstanceOf(InsufficientKvPreWriteBufferException.class); + assertThat(memoryManager.usedBytes()).isEqualTo(retainedBytes); + assertThat(buffer.getAllKvEntries()).hasSize(1); + assertThat(buffer.getMaxLSN()).isZero(); + assertThat(buffer.get(toKey("key-2"))).isNull(); + + buffer.truncateTo(0, TruncateReason.ERROR); + assertThat(memoryManager.usedBytes()).isZero(); + assertThat(memoryManager.isUnderPressure()).isFalse(); + } + + @Test + void testEntryLargerThanHighWatermarkIsNotRetriable() { + KvPreWriteBufferMemoryManager memoryManager = new KvPreWriteBufferMemoryManager(100, 80); + KvPreWriteBuffer buffer = + new KvPreWriteBuffer( + new NopKvBatchWriter(), + TestingMetricGroups.TABLET_SERVER_METRICS, + memoryManager); + + assertThatThrownBy(() -> buffer.insert(toKey("key"), "value".getBytes(), 0)) + .isInstanceOf(RecordTooLargeException.class); + assertThat(memoryManager.usedBytes()).isZero(); + assertThat(buffer.getAllKvEntries()).isEmpty(); + assertThat(buffer.getKvEntryMap()).isEmpty(); + } + + @Test + void testMemoryQuotaIsSharedAcrossBuffers() throws Exception { + KvPreWriteBufferMemoryManager memoryManager = new KvPreWriteBufferMemoryManager(200, 160); + KvPreWriteBuffer firstBuffer = + new KvPreWriteBuffer( + new NopKvBatchWriter(), + TestingMetricGroups.TABLET_SERVER_METRICS, + memoryManager); + KvPreWriteBuffer secondBuffer = + new KvPreWriteBuffer( + new NopKvBatchWriter(), + TestingMetricGroups.TABLET_SERVER_METRICS, + memoryManager); + + firstBuffer.insert(toKey("key"), "value".getBytes(), 0); + assertThatThrownBy(() -> secondBuffer.insert(toKey("key"), "value".getBytes(), 0)) + .isInstanceOf(InsufficientKvPreWriteBufferException.class); + + firstBuffer.close(); + secondBuffer.insert(toKey("key"), "value".getBytes(), 0); + assertThat(secondBuffer.getAllKvEntries()).hasSize(1); + secondBuffer.close(); + assertThat(memoryManager.usedBytes()).isZero(); + } + + @Test + void testCloseReleasesMemory() throws Exception { + KvPreWriteBufferMemoryManager memoryManager = + new KvPreWriteBufferMemoryManager(10_000, 8_000); + KvPreWriteBuffer buffer = + new KvPreWriteBuffer( + new NopKvBatchWriter(), + TestingMetricGroups.TABLET_SERVER_METRICS, + memoryManager); + buffer.insert(toKey("key"), "value".getBytes(), 0); + + buffer.close(); + + assertThat(memoryManager.usedBytes()).isZero(); + } + private static void bufferInsert( KvPreWriteBuffer kvPreWriteBuffer, String key, String value, int elementCount) { kvPreWriteBuffer.insert(toKey(key), value.getBytes(), elementCount); } + private static KvPreWriteBuffer createBuffer(KvBatchWriter kvBatchWriter) { + return new KvPreWriteBuffer( + kvBatchWriter, + TestingMetricGroups.TABLET_SERVER_METRICS, + new KvPreWriteBufferMemoryManager(10_000, 8_000)); + } + private static void bufferUpdate( KvPreWriteBuffer kvPreWriteBuffer, String key, String value, int elementCount) { kvPreWriteBuffer.update(toKey(key), value.getBytes(), elementCount); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/scan/ScannerManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/scan/ScannerManagerTest.java index 3843290f5d..262ffa3593 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/kv/scan/ScannerManagerTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/scan/ScannerManagerTest.java @@ -36,6 +36,7 @@ import org.apache.fluss.server.kv.KvTablet; import org.apache.fluss.server.kv.autoinc.AutoIncrementManager; import org.apache.fluss.server.kv.autoinc.TestingSequenceGeneratorFactory; +import org.apache.fluss.server.kv.prewrite.KvPreWriteBufferMemoryManager; import org.apache.fluss.server.kv.rowmerger.RowMerger; import org.apache.fluss.server.log.LogTablet; import org.apache.fluss.server.log.LogTestUtils; @@ -138,7 +139,8 @@ void setUp() throws Exception { schemaGetter, tableConf.getChangelogImage(), KvManager.getDefaultRateLimiter(), - autoIncrementManager); + autoIncrementManager, + new KvPreWriteBufferMemoryManager(Long.MAX_VALUE, Long.MAX_VALUE - 1)); } @AfterEach diff --git a/fluss-server/src/test/java/org/apache/fluss/server/metrics/group/TabletServerMetricGroupTest.java b/fluss-server/src/test/java/org/apache/fluss/server/metrics/group/TabletServerMetricGroupTest.java new file mode 100644 index 0000000000..810c32d109 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/metrics/group/TabletServerMetricGroupTest.java @@ -0,0 +1,57 @@ +/* + * 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.metrics.group; + +import org.apache.fluss.metrics.Gauge; +import org.apache.fluss.metrics.MetricNames; +import org.apache.fluss.metrics.registry.NOPMetricRegistry; +import org.apache.fluss.server.kv.prewrite.KvPreWriteBufferMemoryManager; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link TabletServerMetricGroup}. */ +class TabletServerMetricGroupTest { + + @Test + void testKvPreWriteBufferMemoryMetrics() { + TabletServerMetricGroup metricGroup = + new TabletServerMetricGroup( + NOPMetricRegistry.INSTANCE, "cluster", "rack", "host", 1); + KvPreWriteBufferMemoryManager manager = new KvPreWriteBufferMemoryManager(100, 80); + metricGroup.registerKvPreWriteBufferMemoryManager(manager); + + assertThat(gaugeValue(metricGroup, MetricNames.KV_PRE_WRITE_BUFFER_MEMORY_USAGE)) + .isEqualTo(0L); + assertThat(gaugeValue(metricGroup, MetricNames.KV_PRE_WRITE_BUFFER_UNDER_PRESSURE)) + .isEqualTo(0); + + assertThat(manager.tryReserve(100)).isTrue(); + assertThat(manager.tryReserve(1)).isFalse(); + + assertThat(gaugeValue(metricGroup, MetricNames.KV_PRE_WRITE_BUFFER_MEMORY_USAGE)) + .isEqualTo(100L); + assertThat(gaugeValue(metricGroup, MetricNames.KV_PRE_WRITE_BUFFER_UNDER_PRESSURE)) + .isEqualTo(1); + } + + private static Object gaugeValue(TabletServerMetricGroup metricGroup, String metricName) { + return ((Gauge) metricGroup.getMetrics().get(metricName)).getValue(); + } +}