Repository: apache/kafka
File: coordinator-common/src/test/java/org/apache/kafka/coordinator/common/runtime/CoordinatorRuntimeTest.java (Java)

Code before:
        // Schedule the loading.
        runtime.scheduleLoadOperation(TP, 10);

        // Verify the initial state.
        CoordinatorRuntime<MockCoordinatorShard, String>.CoordinatorContext ctx = runtime.contextOrThrow(TP);
        assertEquals(0L, ctx.coordinator.lastWrittenOffset());
        assertEquals(0L, ctx.coordinator.lastCommittedOffset());
        assertEquals(List.of(0L), ctx.coordinator.snapshotRegistry().epochsList());

        List<String> records = new ArrayList<>();
        for (int i = 0; i < 1000; i++) {
            records.add("record-" + i);
        }

        // Write #1.
        CompletableFuture<String> write1 = runtime.scheduleWriteOperation("write#1", TP, DEFAULT_WRITE_TIMEOUT,
            state -> new CoordinatorResult<>(records, "response1")
        );

        // Verify that the write has not completed exceptionally.
        // This will catch any exceptions thrown including RecordTooLargeException.
        assertFalse(write1.isCompletedExceptionally());

        int batchSize = mockWriter.entries(TP).get(0).sizeInBytes();
        int maxBatchSize = mockWriter.config(TP).maxMessageSize();
        assertTrue(batchSize < INITIAL_BUFFER_SIZE && INITIAL_BUFFER_SIZE <= maxBatchSize);

        ByteBuffer cachedBuffer = ctx.bufferSupplier.get(1);
        assertEquals(INITIAL_BUFFER_SIZE, cachedBuffer.capacity());
        // ctx.bufferSupplier.get(1); will clear cachedBuffer in bufferSupplier. Use release to put it back to bufferSupplier
        ctx.bufferSupplier.release(cachedBuffer);

        // Reduce max message size below initial buffer size.
        mockWriter.updateConfig(new LogConfig(
            Map.of(TopicConfig.MAX_MESSAGE_BYTES_CONFIG, String.valueOf(INITIAL_BUFFER_SIZE - 66))));
        assertEquals(INITIAL_BUFFER_SIZE - 66, mockWriter.config(TP).maxMessageSize());

        // Write #2.
        CompletableFuture<String> write2 = runtime.scheduleWriteOperation("write#2", TP, DEFAULT_WRITE_TIMEOUT,
            state -> new CoordinatorResult<>(records, "response2")
        );
        assertFalse(write2.isCompletedExceptionally());

        // Verify that there is no cached buffer since the cached buffer size is greater than new maxMessageSize.
        assertEquals(1, ctx.bufferSupplier.get(1).capacity());

        // Write #3.
        CompletableFuture<String> write3 = runtime.scheduleWriteOperation("write#3", TP, DEFAULT_WRITE_TIMEOUT,
            state -> new CoordinatorResult<>(records, "response3")
        );
        assertFalse(write3.isCompletedExceptionally());

Review comment (performance):
nit:
```suggestion
        assertTrue(batchSize <= INITIAL_BUFFER_SIZE && INITIAL_BUFFER_SIZE <= maxBatchSize);
```

Code after:
        // Schedule the loading.
        runtime.scheduleLoadOperation(TP, 10);

        // Verify the initial state.
        CoordinatorRuntime<MockCoordinatorShard, String>.CoordinatorContext ctx = runtime.contextOrThrow(TP);
        assertEquals(0L, ctx.coordinator.lastWrittenOffset());
        assertEquals(0L, ctx.coordinator.lastCommittedOffset());
        assertEquals(List.of(0L), ctx.coordinator.snapshotRegistry().epochsList());

        List<String> records = new ArrayList<>();
        for (int i = 0; i < 1000; i++) {
            records.add("record-" + i);
        }

        // Write #1.
        CompletableFuture<String> write1 = runtime.scheduleWriteOperation("write#1", TP, DEFAULT_WRITE_TIMEOUT,
            state -> new CoordinatorResult<>(records, "response1")
        );

        // Verify that the write has not completed exceptionally.
        // This will catch any exceptions thrown including RecordTooLargeException.
        assertFalse(write1.isCompletedExceptionally());

        int batchSize = mockWriter.entries(TP).get(0).sizeInBytes();
        int maxBatchSize = mockWriter.config(TP).maxMessageSize();
        assertTrue(batchSize <= INITIAL_BUFFER_SIZE && INITIAL_BUFFER_SIZE <= maxBatchSize);

        ByteBuffer cachedBuffer = ctx.bufferSupplier.get(1);
        assertEquals(INITIAL_BUFFER_SIZE, cachedBuffer.capacity());
        // ctx.bufferSupplier.get(1); will clear cachedBuffer in bufferSupplier. Use release to put it back to bufferSupplier
        ctx.bufferSupplier.release(cachedBuffer);

        // Reduce max message size below initial buffer size.
        mockWriter.updateConfig(new LogConfig(
            Map.of(TopicConfig.MAX_MESSAGE_BYTES_CONFIG, String.valueOf(INITIAL_BUFFER_SIZE - 66))));
        assertEquals(INITIAL_BUFFER_SIZE - 66, mockWriter.config(TP).maxMessageSize());

        // Write #2.
        CompletableFuture<String> write2 = runtime.scheduleWriteOperation("write#2", TP, DEFAULT_WRITE_TIMEOUT,
            state -> new CoordinatorResult<>(records, "response2")
        );
        assertFalse(write2.isCompletedExceptionally());

        // Verify that there is no cached buffer since the cached buffer size is greater than new maxMessageSize.
        assertEquals(1, ctx.bufferSupplier.get(1).capacity());

        // Write #3.
        CompletableFuture<String> write3 = runtime.scheduleWriteOperation("write#3", TP, DEFAULT_WRITE_TIMEOUT,
            state -> new CoordinatorResult<>(records, "response3")
        );
        assertFalse(write3.isCompletedExceptionally());