[server] Add pre-write buffer memory metrics for primary key tables

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

@@ -256,6 +256,20 @@ public class MetricNames {
     public static final String ROCKSDB_SHARED_WRITE_BUFFER_CAPACITY =
             "rocksdbSharedWriteBufferCapacity";
 
+    // Server-level pre-write buffer metrics (aggregated from all KV tablets, Sum aggregation)
+    /**
+     * Estimated memory usage of the pre-write buffers across all KV tablets in this server (Sum
+     * aggregation).
+     */
+    public static final String KV_PRE_WRITE_BUFFER_MEMORY_USAGE_BYTES =
+            "kvPreWriteBufferMemoryUsageBytes";
+
+    /**
+     * Number of entries buffered in the pre-write buffers across all KV tablets in this server (Sum
+     * aggregation).
+     */
+    public static final String KV_PRE_WRITE_BUFFER_ENTRY_COUNT = "kvPreWriteBufferEntryCount";
+
     // Table-level RocksDB memory metrics (Sum aggregation)
     /** Total memtable memory usage across all buckets of this table. */
     public static final String ROCKSDB_MEMTABLE_MEMORY_USAGE_TOTAL =

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

@@ -227,6 +227,15 @@ private KvManager(
                     this::getSharedBlockCachePinnedUsage,
                     conf.get(ConfigOptions.KV_SHARED_BLOCK_CACHE_SIZE).getBytes());
         }
+        tabletServerMetricGroup.setPreWriteBufferMetrics(
+                () ->
+                        currentKvs.values().stream()
+                                .mapToLong(KvTablet::kvPreWriteBufferMemoryUsageBytes)
+                                .sum(),
+                () ->
+                        currentKvs.values().stream()
+                                .mapToInt(KvTablet::kvPreWriteBufferEntryCount)
+                                .sum());
     }
 
     private static RateLimiter createSharedRateLimiter(Configuration conf) {

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

@@ -839,6 +839,16 @@ FlushState getFlushState() {
         return inReadLock(kvLock, () -> flushState);
     }
 
+    /** Returns an instantaneous estimate of this kv tablet's pre-write buffer memory usage. */
+    public long kvPreWriteBufferMemoryUsageBytes() {
+        return kvPreWriteBuffer.memoryUsageBytes();
+    }
+
+    /** Returns the number of entries held in this kv tablet's pre-write buffer. */
+    public int kvPreWriteBufferEntryCount() {
+        return kvPreWriteBuffer.entryCount();
+    }
+
     @VisibleForTesting
     void setFlushState(FlushState state) {
         inWriteLock(kvLock, () -> flushState = state);

fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBuffer.java

@@ -89,6 +89,19 @@
 @NotThreadSafe
 public class KvPreWriteBuffer {
 
+    /**
+     * Estimated JVM heap overhead of a single buffered entry besides its key/value payload bytes,
+     * covering the {@link KvEntry} object, its {@link Key} and {@link Value} wrappers, the byte
+     * array headers and the linked list node.
+     */
+    private static final long PER_ENTRY_OVERHEAD_BYTES = 144;
+
+    /**
+     * Estimated JVM heap overhead of a hash map node for an entry that is the latest version of its
+     * key in the buffer.
+     */
+    private static final long PER_MAP_NODE_OVERHEAD_BYTES = 32;
+
     // a mapping from the key to the kv-entry
     private final Map<Key, KvEntry> kvEntryMap = new HashMap<>();
 
@@ -105,6 +118,13 @@ public class KvPreWriteBuffer {
     // Accumulated byte size of entries not yet completed by a flush.
     private long pendingFlushBytes = 0;
 
+    // Estimated total memory footprint of the held entries, updated incrementally on the write
+    // path.
+    private volatile long memoryUsageBytes = 0;
+
+    // Number of held entries.
+    private volatile int entryCount = 0;
+
     public KvPreWriteBuffer(TabletServerMetricGroup serverMetricGroup) {
         truncateAsDuplicatedCount = serverMetricGroup.kvTruncateAsDuplicatedCount();
         truncateAsErrorCount = serverMetricGroup.kvTruncateAsErrorCount();
@@ -164,8 +184,8 @@ private void doPut(ChangeType changeType, Key key, Value value, long lsn) {
         allKvEntries.addLast(kvEntry);
         // update the max lsn
         maxLogSequenceNumber = lsn;
-        // track accumulated bytes for flush budget gating
-        pendingFlushBytes += entryBytes(key, value);
+        // update the accounting for flush budget gating and metrics
+        addToAccounting(kvEntry);
     }
 
     /**
@@ -209,8 +229,7 @@ public void truncateTo(long targetLogSequenceNumber, TruncateReason truncateReas
                                 + ", targetLogSequenceNumber="
                                 + targetLogSequenceNumber);
             }
-            pendingFlushBytes -= entryBytes(entry.getKey(), entry.getValue());
-            boolean removed = kvEntryMap.remove(entry.getKey(), entry);
+            boolean removed = removeFromMapAndAccounting(entry);
             // the removed entry is no longer the successor of its previous version; clear the
             // forward link so the truncated entry does not stay reachable through it
             if (entry.previousEntry != null) {
@@ -234,6 +253,21 @@ public long pendingFlushBytes() {
         return pendingFlushBytes;
     }
 
+    /** Returns the number of entries currently held in this buffer. */
+    public int entryCount() {
+        return entryCount;
+    }
+
+    /**
+     * Returns an estimation of the total memory footprint of the entries currently held in this
+     * buffer, including the key/value payload bytes tracked by {@link #pendingFlushBytes()} and the
+     * per-entry JVM object overhead. This is an approximation for observability purposes, not an
+     * exact measurement.
+     */
+    public long memoryUsageBytes() {
+        return memoryUsageBytes;
+    }
+
     /**
      * Prepares a prefix of entries for asynchronous flush without removing them from the buffer.
      *
@@ -269,8 +303,7 @@ public int completeFlush(PreparedFlush preparedFlush) {
                 throw new IllegalStateException("Prepared flush entry is not in PREPARED state.");
             }
             entry.state = EntryState.FLUSHED;
-            pendingFlushBytes -= entryBytes(entry.getKey(), entry.getValue());
-            kvEntryMap.remove(entry.getKey(), entry);
+            removeFromMapAndAccounting(entry);
             // the immediate successor is the only live referencer of a flushed entry; clearing
             // its reference makes the flushed entry (and, transitively, its older versions)
             // unreachable instead of being retained while no longer counted by pendingFlushBytes
@@ -305,6 +338,37 @@ public void abortAllPrepared() {
         }
     }
 
+    /**
+     * Adds an entry to the incrementally maintained memory estimate and entry count. An entry
+     * without a previous version is the latest version of a new key and thus adds one map node.
+     */
+    private void addToAccounting(KvEntry entry) {
+        long bytes = entryBytes(entry.getKey(), entry.getValue());
+        pendingFlushBytes += bytes;
+        memoryUsageBytes +=
+                bytes
+                        + PER_ENTRY_OVERHEAD_BYTES
+                        + (entry.previousEntry == null ? PER_MAP_NODE_OVERHEAD_BYTES : 0L);
+        entryCount++;
+    }
+
+    /**
+     * Removes an entry from the key map and deducts it from the incrementally maintained memory
+     * estimate and entry count. Returns whether the entry was the latest version of its key and
+     * thus removed from the map.
+     */
+    private boolean removeFromMapAndAccounting(KvEntry entry) {
+        long bytes = entryBytes(entry.getKey(), entry.getValue());
+        pendingFlushBytes -= bytes;
+        boolean removedFromMap = kvEntryMap.remove(entry.getKey(), entry);
+        memoryUsageBytes -=
+                bytes
+                        + PER_ENTRY_OVERHEAD_BYTES
+                        + (removedFromMap ? PER_MAP_NODE_OVERHEAD_BYTES : 0L);
+        entryCount--;
+        return removedFromMap;
+    }
+
     private static long entryBytes(Key key, Value value) {
         return (long) key.key.length + (value.value != null ? value.value.length : 0L);
     }

fluss-server/src/main/java/org/apache/fluss/server/metrics/group/TabletServerMetricGroup.java

@@ -90,6 +90,12 @@ public class TabletServerMetricGroup extends AbstractMetricGroup {
 
     private volatile long sharedWriteBufferCapacity;
 
+    /** Supplier for aggregated pre-write buffer memory usage, set by KvManager. */
+    private volatile LongSupplier preWriteBufferMemoryUsageSupplier = () -> 0L;
+
+    /** Supplier for aggregated pre-write buffer entry count, set by KvManager. */
+    private volatile LongSupplier preWriteBufferEntryCountSupplier = () -> 0L;
+
     public TabletServerMetricGroup(
             MetricRegistry registry, String clusterId, String rack, String hostname, int serverId) {
         super(registry, new String[] {clusterId, hostname, NAME}, null);
@@ -151,6 +157,9 @@ public TabletServerMetricGroup(
 
         // Register server-level RocksDB aggregated metrics
         registerServerRocksDBMetrics();
+
+        // Register server-level pre-write buffer aggregated metrics
+        registerServerKvPreWriteBufferMetrics();
     }
 
     /**
@@ -215,6 +224,34 @@ public void setSharedWriteBufferMetrics(LongSupplier usageSupplier, long capacit
         this.sharedWriteBufferCapacity = capacity;
     }
 
+    /**
+     * Register server-level pre-write buffer aggregated metrics. These metrics aggregate the memory
+     * usage of the pre-write buffers of all KV tablets in this server.
+     */
+    private void registerServerKvPreWriteBufferMetrics() {
+        gauge(
+                MetricNames.KV_PRE_WRITE_BUFFER_MEMORY_USAGE_BYTES,
+                () -> preWriteBufferMemoryUsageSupplier.getAsLong());
+        gauge(
+                MetricNames.KV_PRE_WRITE_BUFFER_ENTRY_COUNT,
+                () -> preWriteBufferEntryCountSupplier.getAsLong());
+    }
+
+    /**
+     * Sets the aggregated pre-write buffer metrics. Called by KvManager at construction time.
+     *
+     * @param memoryUsageSupplier supplier for the estimated memory usage of all pre-write buffers
+     * @param entryCountSupplier supplier for the total number of buffered entries across all
+     *     pre-write buffers
+     */
+    public void setPreWriteBufferMetrics(
+            LongSupplier memoryUsageSupplier, LongSupplier entryCountSupplier) {
+        this.preWriteBufferMemoryUsageSupplier =
+                checkNotNull(memoryUsageSupplier, "memoryUsageSupplier must not be null");
+        this.preWriteBufferEntryCountSupplier =
+                checkNotNull(entryCountSupplier, "entryCountSupplier must not be null");
+    }
+
     @Override
     protected final void putVariables(Map<String, String> variables) {
         variables.put("cluster_id", clusterId);

fluss-server/src/test/java/org/apache/fluss/server/kv/KvManagerTest.java

@@ -282,6 +282,54 @@ void testSharedWriteBufferConfiguredThroughKvManagerCreateAndLoad() throws Excep
                 .isEqualTo(capacity.getBytes());
     }
 
+    @Test
+    void testPreWriteBufferServerLevelMetrics() throws Exception {
+        initTableBuckets(null);
+        assertThat(
+                        gaugeValue(
+                                TestingMetricGroups.TABLET_SERVER_METRICS,
+                                MetricNames.KV_PRE_WRITE_BUFFER_MEMORY_USAGE_BYTES))
+                .isEqualTo(0L);
+        assertThat(
+                        gaugeValue(
+                                TestingMetricGroups.TABLET_SERVER_METRICS,
+                                MetricNames.KV_PRE_WRITE_BUFFER_ENTRY_COUNT))
+                .isEqualTo(0L);
+
+        // write one kv record without flushing so that it stays in the pre-write buffer
+        KvTablet kvTablet = getOrCreateKv(tablePath1, null, tableBucket1);
+        KvRecordBatch kvRecordBatch =
+                kvRecordBatchFactory.ofRecords(
+                        Collections.singletonList(
+                                kvRecordFactory.ofRecord(
+                                        "key1".getBytes(), new Object[] {1, "a"})));
+        kvTablet.putAsLeader(kvRecordBatch, null);
+
+        assertThat(
+                        gaugeValue(
+                                TestingMetricGroups.TABLET_SERVER_METRICS,
+                                MetricNames.KV_PRE_WRITE_BUFFER_ENTRY_COUNT))
+                .isEqualTo(1L);
+        assertThat(
+                        gaugeValue(
+                                TestingMetricGroups.TABLET_SERVER_METRICS,
+                                MetricNames.KV_PRE_WRITE_BUFFER_MEMORY_USAGE_BYTES))
+                .isPositive();
+
+        // the usage drops to zero once the buffered entries are flushed
+        flushAndWait(kvTablet, Long.MAX_VALUE);
+        assertThat(
+                        gaugeValue(
+                                TestingMetricGroups.TABLET_SERVER_METRICS,
+                                MetricNames.KV_PRE_WRITE_BUFFER_ENTRY_COUNT))
+                .isEqualTo(0L);
+        assertThat(
+                        gaugeValue(
+                                TestingMetricGroups.TABLET_SERVER_METRICS,
+                                MetricNames.KV_PRE_WRITE_BUFFER_MEMORY_USAGE_BYTES))
+                .isEqualTo(0L);
+    }
+
     @ParameterizedTest
     @MethodSource("partitionProvider")
     void testCreateKv(String partitionName) throws Exception {

fluss-server/src/test/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBufferTest.java

@@ -454,6 +454,36 @@ void testPendingFlushBytesTracking() {
         assertThat(buffer.pendingFlushBytes()).isEqualTo(0);
     }
 
+    @Test
+    void testEstimatedMemoryUsage() {
+        KvPreWriteBuffer buffer = new KvPreWriteBuffer(TestingMetricGroups.TABLET_SERVER_METRICS);
+
+        assertThat(buffer.memoryUsageBytes()).isEqualTo(0L);
+        assertThat(buffer.entryCount()).isEqualTo(0);
+
+        // +key1(10 bytes), +key2(11 bytes), -key3(4 bytes): 25 payload bytes in total
+        bufferInsert(buffer, "key1", "value1", 1);
+        bufferInsert(buffer, "key2", "value22", 2);
+        bufferDelete(buffer, "key3", 3);
+        long payloadBytes = 25;
+
+        // the estimation covers the payload bytes plus the per-entry object overhead
+        assertThat(buffer.memoryUsageBytes()).isGreaterThan(payloadBytes);
+        assertThat(buffer.entryCount()).isEqualTo(3);
+
+        // flushing all entries releases the whole accounted usage
+        flushBuffer(buffer, Long.MAX_VALUE);
+        assertThat(buffer.memoryUsageBytes()).isEqualTo(0L);
+        assertThat(buffer.entryCount()).isEqualTo(0);
+
+        // truncating entries also releases their accounted usage
+        bufferInsert(buffer, "key1", "value1", 4);
+        assertThat(buffer.memoryUsageBytes()).isPositive();
+        buffer.truncateTo(4, TruncateReason.ERROR);
+        assertThat(buffer.memoryUsageBytes()).isEqualTo(0L);
+        assertThat(buffer.entryCount()).isEqualTo(0);
+    }
+
     @Test
     void testCompleteFlushDetachesFlushedEntriesFromPreviousChain() {
         KvPreWriteBuffer buffer = new KvPreWriteBuffer(TestingMetricGroups.TABLET_SERVER_METRICS);

website/docs/maintenance/observability/monitor-metrics.md

@@ -463,8 +463,8 @@ Some metrics might not be exposed when using other JVM implementations (e.g. IBM
   </thead>
   <tbody>
     <tr>
-      <th rowspan="37"><strong>tabletserver</strong></th>
-      <td style={{textAlign: 'center', verticalAlign: 'middle' }} rowspan="25">-</td>
+      <th rowspan="39"><strong>tabletserver</strong></th>
+      <td style={{textAlign: 'center', verticalAlign: 'middle' }} rowspan="27">-</td>
       <td>messagesInPerSecond</td>
       <td>The number of messages written per second to this server.</td>
       <td>Meter</td>
@@ -588,6 +588,16 @@ Some metrics might not be exposed when using other JVM implementations (e.g. IBM
       <td>preWriteBufferTruncateAsErrorPerSecond</td>
       <td>The number of kv pre-write buffer truncate due to the error happened when writing cdc to log per second.</td>
       <td>Meter</td>
+    </tr>
+     <tr>
+      <td>kvPreWriteBufferMemoryUsageBytes</td>
+      <td>Estimated total memory usage of the KV pre-write buffers across all KV tablets in this server (in bytes), including the key/value payload bytes and the per-entry object overhead. It is an approximation for observability, not an exact measurement.</td>
+      <td>Gauge</td>
+    </tr>
+     <tr>
+      <td>kvPreWriteBufferEntryCount</td>
+      <td>The number of entries buffered in the KV pre-write buffers across all KV tablets in this server.</td>
+      <td>Gauge</td>
     </tr>
     <tr>
       <td rowspan="4">historical</td>