KAFKA-19427: Allow the coordinator to grow its buffer dynamically

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

@@ -600,6 +600,16 @@ class CoordinatorContext {
          */
         CoordinatorBatch currentBatch;
 
+        /**
+         * Used for initial buffer size for write operations.
+         */
+        final int minBufferSize;
+
+        /**
+         * The max batch size
+         */
+        final int maxBatchSize;
+
         /**
          * Constructor.
          *
@@ -627,6 +637,8 @@ private CoordinatorContext(
                 defaultWriteTimeout
             );
             this.bufferSupplier = new BufferSupplier.GrowableBufferSupplier();
+            this.minBufferSize = 16384; // 16KB
+            this.maxBatchSize = partitionWriter.config(tp).maxMessageSize();
         }
 
         /**
@@ -759,7 +771,13 @@ private void freeCurrentBatch() {
             currentBatch.lingerTimeoutTask.ifPresent(TimerTask::cancel);
 
             // Release the buffer.
-            bufferSupplier.release(currentBatch.buffer);
+            if (currentBatch.builder.buffer().capacity() > maxBatchSize) {
+                // If the buffer exceeds the maxBatchSize, we should revert to using the original smaller buffer to avoid retaining the larger one.
+                bufferSupplier.release(currentBatch.buffer);
+            } else {
+                bufferSupplier.release(currentBatch.builder.buffer());
+            }
+
 
             currentBatch = null;
         }
@@ -856,10 +874,8 @@ private void maybeAllocateNewBatch(
             long currentTimeMs
         ) {
             if (currentBatch == null) {
-                LogConfig logConfig = partitionWriter.config(tp);
-                int maxBatchSize = logConfig.maxMessageSize();
                 long prevLastWrittenOffset = coordinator.lastWrittenOffset();
-                ByteBuffer buffer = bufferSupplier.get(maxBatchSize);
+                ByteBuffer buffer = bufferSupplier.get(minBufferSize);
 
                 MemoryRecordsBuilder builder = new MemoryRecordsBuilder(
                     buffer,
@@ -1908,11 +1924,6 @@ public void onHighWatermarkUpdated(
         }
     }
 
-    /**
-     * 16KB. Used for initial buffer size for write operations.
-     */
-    static final int MIN_BUFFER_SIZE = 16384;
-
     /**
      * The log prefix.
      */