KAFKA-19519: Introduce group.coordinator.append.max.buffer.size config

clients/src/main/java/org/apache/kafka/common/utils/BufferSupplier.java

@@ -22,6 +22,7 @@
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Simple non-threadsafe interface for caching byte buffers. This is suitable for simple cases like ensuring that
@@ -30,6 +31,7 @@
  * iterating over the records in the batch.
  */
 public abstract class BufferSupplier implements AutoCloseable {
+    protected AtomicLong cachedSize = new AtomicLong();
 
     public static final BufferSupplier NO_CACHING = new BufferSupplier() {
         @Override
@@ -40,6 +42,11 @@ public ByteBuffer get(int capacity) {
         @Override
         public void release(ByteBuffer buffer) {}
 
+        @Override
+        public long size() {
+            return 0;
+        }
+
         @Override
         public void close() {}
     };
@@ -58,6 +65,11 @@ public static BufferSupplier create() {
      */
     public abstract void release(ByteBuffer buffer);
 
+    /**
+     * Return total size in bytes of cached buffers.
+     */
+    public abstract long size();
+
     /**
      * Release all resources associated with this supplier.
      */
@@ -82,11 +94,18 @@ public void release(ByteBuffer buffer) {
             // We currently keep a single buffer in flight, so optimise for that case
             Deque<ByteBuffer> bufferQueue = bufferMap.computeIfAbsent(buffer.capacity(), k -> new ArrayDeque<>(1));
             bufferQueue.addLast(buffer);
+            cachedSize.addAndGet(buffer.capacity());
+        }
+
+        @Override
+        public long size() {
+            return cachedSize.get();
         }
 
         @Override
         public void close() {
             bufferMap.clear();
+            cachedSize.set(0);
         }
     }
 
@@ -113,12 +132,18 @@ public ByteBuffer get(int minCapacity) {
         public void release(ByteBuffer buffer) {
             buffer.clear();
             cachedBuffer = buffer;
+            cachedSize.set(buffer.capacity());
+        }
+
+        @Override
+        public long size() {
+            return cachedSize.get();
         }
 
         @Override
         public void close() {
             cachedBuffer = null;
+            cachedSize.set(0);
         }
     }
-
 }

coordinator-common/src/main/java/org/apache/kafka/coordinator/common/runtime/CoordinatorRuntime.java

@@ -65,6 +65,7 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 import static java.lang.Math.min;
@@ -119,6 +120,7 @@ public static class Builder<S extends CoordinatorShard<U>, U> {
         private Compression compression;
         private OptionalInt appendLingerMs;
         private ExecutorService executorService;
+        private Supplier<Integer> appendMaxBufferSizeSupplier;
 
         public Builder<S, U> withLogPrefix(String logPrefix) {
             this.logPrefix = logPrefix;
@@ -195,6 +197,11 @@ public Builder<S, U> withExecutorService(ExecutorService executorService) {
             return this;
         }
 
+        public Builder<S, U> withMaxBufferSizeSupplier(Supplier<Integer> maxBufferSizeSupplier) {
+            this.appendMaxBufferSizeSupplier = maxBufferSizeSupplier;
+            return this;
+        }
+
         @SuppressWarnings("checkstyle:CyclomaticComplexity")
         public CoordinatorRuntime<S, U> build() {
             if (logPrefix == null)
@@ -227,6 +234,8 @@ public CoordinatorRuntime<S, U> build() {
                 throw new IllegalArgumentException("AppendLinger must be empty or >= 0");
             if (executorService == null)
                 throw new IllegalArgumentException("ExecutorService must be set.");
+            if (appendMaxBufferSizeSupplier == null)
+                throw new IllegalArgumentException("Append max buffer size supplier must be set.");
 
             return new CoordinatorRuntime<>(
                 logPrefix,
@@ -243,7 +252,8 @@ public CoordinatorRuntime<S, U> build() {
                 serializer,
                 compression,
                 appendLingerMs,
-                executorService
+                executorService,
+                appendMaxBufferSizeSupplier
             );
         }
     }
@@ -771,13 +781,15 @@ private void freeCurrentBatch() {
             // Cancel the linger timeout.
             currentBatch.lingerTimeoutTask.ifPresent(TimerTask::cancel);
 
-            // Release the buffer only if it is not larger than the maxBatchSize.
-            int maxBatchSize = partitionWriter.config(tp).maxMessageSize();
+            // Release the buffer only if it is not larger than the max buffer size.
+            int maxBufferSize = appendMaxBufferSizeSupplier.get();
 
-            if (currentBatch.builder.buffer().capacity() <= maxBatchSize) {
+            if (currentBatch.builder.buffer().capacity() <= maxBufferSize) {
                 bufferSupplier.release(currentBatch.builder.buffer());
-            } else if (currentBatch.buffer.capacity() <= maxBatchSize) {
+            } else if (currentBatch.buffer.capacity() <= maxBufferSize) {
                 bufferSupplier.release(currentBatch.buffer);
+            } else {
+                runtimeMetrics.recordAppendBufferDiscarded();
             }
 
             currentBatch = null;
@@ -2050,6 +2062,11 @@ public void onHighWatermarkUpdated(
      */
     private final ExecutorService executorService;
 
+    /**
+     * The maximum buffer size that the coordinator can cache.
+     */
+    private final Supplier<Integer> appendMaxBufferSizeSupplier;
+
     /**
      * Atomic boolean indicating whether the runtime is running.
      */
@@ -2078,6 +2095,7 @@ public void onHighWatermarkUpdated(
      * @param compression                       The compression codec.
      * @param appendLingerMs                    The append linger time in ms.
      * @param executorService                   The executor service.
+     * @param appendMaxBufferSizeSupplier       The append max buffer size supplier.
      */
     @SuppressWarnings("checkstyle:ParameterNumber")
     private CoordinatorRuntime(
@@ -2096,6 +2114,7 @@ private CoordinatorRuntime(
         Compression compression,
         OptionalInt appendLingerMs,
         ExecutorService executorService
+        Supplier<Integer> appendMaxBufferSizeSupplier
     ) {
         this.logPrefix = logPrefix;
         this.log = logContext.logger(CoordinatorRuntime.class);
@@ -2113,6 +2132,10 @@ private CoordinatorRuntime(
         this.compression = compression;
         this.appendLingerMs = appendLingerMs;
         this.executorService = executorService;
+        this.appendMaxBufferSizeSupplier = appendMaxBufferSizeSupplier;
+        this.runtimeMetrics.registerAppendBufferSizeGauge(
+            () -> coordinators.values().stream().mapToLong(c -> c.bufferSupplier.size()).sum()
+        );
     }
 
     /**

coordinator-common/src/main/java/org/apache/kafka/coordinator/common/runtime/CoordinatorRuntimeMetrics.java

@@ -86,4 +86,16 @@ public interface CoordinatorRuntimeMetrics extends AutoCloseable {
      * @param sizeSupplier The size supplier.
      */
     void registerEventQueueSizeGauge(Supplier<Integer> sizeSupplier);
+
+    /**
+     * Register the coordinator append buffer size gauge.
+     *
+     * @param sizeSupplier The size supplier.
+     */
+    void registerAppendBufferSizeGauge(Supplier<Long> sizeSupplier);
+
+    /**
+     * Called when a buffer is discarded upon release instead of being cached.
+     */
+    void recordAppendBufferDiscarded();
 }

coordinator-common/src/main/java/org/apache/kafka/coordinator/common/runtime/CoordinatorRuntimeMetricsImpl.java

@@ -69,6 +69,16 @@ public class CoordinatorRuntimeMetricsImpl implements CoordinatorRuntimeMetrics
      */
     public static final String BATCH_FLUSH_TIME_METRIC_NAME = "batch-flush-time-ms";
 
+    /**
+     * The coordinator append buffer size metric name.
+     */
+    public static final String APPEND_BUFFER_SIZE_METRIC_NAME = "coordinator-append-buffer-size-bytes";
+
+    /**
+     * The coordinator append buffer skip cache count metric name.
+     */
+    public static final String APPEND_BUFFER_SKIP_CACHE_COUNT_METRIC_NAME = "coordinator-append-buffer-skip-cache-count";
+
     /**
      * Metric to count the number of partitions in Loading state.
      */
@@ -92,6 +102,17 @@ public class CoordinatorRuntimeMetricsImpl implements CoordinatorRuntimeMetrics
      */
     private final MetricName eventQueueSize;
 
+    /**
+     * Metric to count the size of the coordinator append buffer.
+     */
+    private final MetricName appendBufferSize;
+
+    /**
+     * Metric to count the number of over-sized append buffers that were discarded.
+     */
+    private final MetricName appendBufferSkipCacheCount;
+    private final AtomicLong appendBufferSkipCacheCounter = new AtomicLong(0);
+
     /**
      * The Kafka metrics registry.
      */
@@ -156,6 +177,17 @@ public CoordinatorRuntimeMetricsImpl(Metrics metrics, String metricsGroup) {
 
         this.eventQueueSize = kafkaMetricName("event-queue-size", "The event accumulator queue size.");
 
+        this.appendBufferSize = kafkaMetricName(
+            APPEND_BUFFER_SIZE_METRIC_NAME,
+            "The current total size in bytes of the append buffers being held in the coordinator's cache."
+        );
+
+        this.appendBufferSkipCacheCount = kafkaMetricName(
+            APPEND_BUFFER_SKIP_CACHE_COUNT_METRIC_NAME,
+            "The count of over-sized append buffers that were discarded instead of being cached upon release."
+        );
+
+        metrics.addMetric(appendBufferSkipCacheCount, (Gauge<Long>) (config, now) -> appendBufferSkipCacheCounter.get());
         metrics.addMetric(numPartitionsLoading, (Gauge<Long>) (config, now) -> numPartitionsLoadingCounter.get());
         metrics.addMetric(numPartitionsActive, (Gauge<Long>) (config, now) -> numPartitionsActiveCounter.get());
         metrics.addMetric(numPartitionsFailed, (Gauge<Long>) (config, now) -> numPartitionsFailedCounter.get());
@@ -252,7 +284,9 @@ public void close() {
             numPartitionsLoading,
             numPartitionsActive,
             numPartitionsFailed,
-            eventQueueSize
+            eventQueueSize,
+            appendBufferSize,
+            appendBufferSkipCacheCount
         ).forEach(metrics::removeMetric);
 
         metrics.removeSensor(partitionLoadSensor.name());
@@ -340,4 +374,14 @@ public void recordThreadIdleTime(double idleTimeMs) {
     public void registerEventQueueSizeGauge(Supplier<Integer> sizeSupplier) {
         metrics.addMetric(eventQueueSize, (Gauge<Long>) (config, now) -> (long) sizeSupplier.get());
     }
+
+    @Override
+    public void registerAppendBufferSizeGauge(Supplier<Long> sizeSupplier) {
+        metrics.addMetric(appendBufferSize, (Gauge<Long>) (config, now) -> sizeSupplier.get());
+    }
+
+    @Override
+    public void recordAppendBufferDiscarded() {
+        appendBufferSkipCacheCounter.incrementAndGet();
+    }
 }

coordinator-common/src/test/java/org/apache/kafka/coordinator/common/runtime/CoordinatorRuntimeMetricsImplTest.java

@@ -31,6 +31,8 @@
 import java.util.Set;
 import java.util.stream.IntStream;
 
+import static org.apache.kafka.coordinator.common.runtime.CoordinatorRuntimeMetricsImpl.APPEND_BUFFER_SIZE_METRIC_NAME;
+import static org.apache.kafka.coordinator.common.runtime.CoordinatorRuntimeMetricsImpl.APPEND_BUFFER_SKIP_CACHE_COUNT_METRIC_NAME;
 import static org.apache.kafka.coordinator.common.runtime.CoordinatorRuntimeMetricsImpl.BATCH_FLUSH_TIME_METRIC_NAME;
 import static org.apache.kafka.coordinator.common.runtime.CoordinatorRuntimeMetricsImpl.BATCH_LINGER_TIME_METRIC_NAME;
 import static org.apache.kafka.coordinator.common.runtime.CoordinatorRuntimeMetricsImpl.EVENT_PROCESSING_TIME_METRIC_NAME;
@@ -56,6 +58,8 @@ public void testMetricNames() {
             kafkaMetricName(metrics, NUM_PARTITIONS_METRIC_NAME, "state", "loading"),
             kafkaMetricName(metrics, NUM_PARTITIONS_METRIC_NAME, "state", "active"),
             kafkaMetricName(metrics, NUM_PARTITIONS_METRIC_NAME, "state", "failed"),
+            kafkaMetricName(metrics, APPEND_BUFFER_SIZE_METRIC_NAME),
+            kafkaMetricName(metrics, APPEND_BUFFER_SKIP_CACHE_COUNT_METRIC_NAME),
             kafkaMetricName(metrics, "event-queue-size"),
             kafkaMetricName(metrics, "partition-load-time-max"),
             kafkaMetricName(metrics, "partition-load-time-avg"),
@@ -90,6 +94,7 @@ public void testMetricNames() {
 
         try (CoordinatorRuntimeMetricsImpl runtimeMetrics = new CoordinatorRuntimeMetricsImpl(metrics, METRICS_GROUP)) {
             runtimeMetrics.registerEventQueueSizeGauge(() -> 0);
+            runtimeMetrics.registerAppendBufferSizeGauge(() -> 0L);
             expectedMetrics.forEach(metricName -> assertTrue(metrics.metrics().containsKey(metricName)));
         }
 
@@ -238,6 +243,55 @@ public void testEventQueueSizeMetricsGroupIsolation() {
         }
     }
 
+    @Test
+    public void testAppendBufferSize() {
+        Metrics metrics = new Metrics();
+
+        try (CoordinatorRuntimeMetricsImpl runtimeMetrics = new CoordinatorRuntimeMetricsImpl(metrics, METRICS_GROUP)) {
+            runtimeMetrics.registerAppendBufferSizeGauge(() -> 5L);
+            assertMetricGauge(metrics, kafkaMetricName(metrics, APPEND_BUFFER_SIZE_METRIC_NAME), 5);
+        }
+    }
+
+    @Test
+    public void testAppendBufferSizeMetricsGroupIsolation() {
+        Time time = new MockTime();
+        Metrics metrics = new Metrics(time);
+
+        try (CoordinatorRuntimeMetricsImpl runtimeMetrics = new CoordinatorRuntimeMetricsImpl(metrics, METRICS_GROUP);
+            CoordinatorRuntimeMetricsImpl otherRuntimeMetrics = new CoordinatorRuntimeMetricsImpl(metrics, OTHER_METRICS_GROUP)) {
+            runtimeMetrics.registerAppendBufferSizeGauge(() -> 5L);
+            otherRuntimeMetrics.registerAppendBufferSizeGauge(() -> 0L);
+
+            assertMetricGauge(metrics, kafkaMetricName(metrics, APPEND_BUFFER_SIZE_METRIC_NAME), 5);
+            assertMetricGauge(metrics, otherGroupKafkaMetricName(metrics, APPEND_BUFFER_SIZE_METRIC_NAME), 0);
+        }
+    }
+
+    @Test
+    public void testAppendBufferSkipCacheCount() {
+        Metrics metrics = new Metrics();
+
+        try (CoordinatorRuntimeMetricsImpl runtimeMetrics = new CoordinatorRuntimeMetricsImpl(metrics, METRICS_GROUP)) {
+            runtimeMetrics.recordAppendBufferDiscarded();
+            assertMetricGauge(metrics, kafkaMetricName(metrics, APPEND_BUFFER_SKIP_CACHE_COUNT_METRIC_NAME), 1);
+        }
+    }
+
+    @Test
+    public void testAppendBufferSkipCacheCountMetricsGroupIsolation() {
+        Time time = new MockTime();
+        Metrics metrics = new Metrics(time);
+
+        try (CoordinatorRuntimeMetricsImpl runtimeMetrics = new CoordinatorRuntimeMetricsImpl(metrics, METRICS_GROUP);
+            CoordinatorRuntimeMetricsImpl otherRuntimeMetrics = new CoordinatorRuntimeMetricsImpl(metrics, OTHER_METRICS_GROUP)) {
+            runtimeMetrics.recordAppendBufferDiscarded();
+
+            assertMetricGauge(metrics, kafkaMetricName(metrics, APPEND_BUFFER_SKIP_CACHE_COUNT_METRIC_NAME), 1);
+            assertMetricGauge(metrics, otherGroupKafkaMetricName(metrics, APPEND_BUFFER_SKIP_CACHE_COUNT_METRIC_NAME), 0);
+        }
+    }
+
     @ParameterizedTest
     @ValueSource(strings = {
         EVENT_QUEUE_TIME_METRIC_NAME,