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

Code before:
                .build();

        // 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());

        // Generate enough records to create a batch that has INITIAL_BUFFER_SIZE < batchSize < maxBatchSize
        List<String> records = new ArrayList<>();
        for (int i = 0; i < 1000000; 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());

        // Verify that the next buffer retrieved from the bufferSupplier is the expanded buffer.
        assertTrue(ctx.bufferSupplier.get(1).capacity() > INITIAL_BUFFER_SIZE);
    }

    @Test
    public void testBufferShrinkWhenMaxMessageSizeReducedBelowInitialBufferSize() {
        MockTimer timer = new MockTimer();
        var mockWriter = new InMemoryPartitionWriter(false) {
            private LogConfig config = new LogConfig(Map.of(
                TopicConfig.MAX_MESSAGE_BYTES_CONFIG, String.valueOf(1024 * 1024) // 1MB
            ));

            @Override
            public LogConfig config(TopicPartition tp) {
                return config;
            }

            public void updateConfig(LogConfig newConfig) {
                this.config = newConfig;
            }
        };
        StringSerializer serializer = new StringSerializer();

        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()

Review comment (performance):
It’s clearer to check if INITIAL_BUFFER_SIZE < batchSize < maxBatchSize first.

```suggestion
        int batchSize = writer.entries(TP).get(0).sizeInBytes();
        assertTrue(batchSize > INITIAL_BUFFER_SIZE && batchSize <= maxBatchSize);
```

Code after:
                .build();

        // 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());

        // Generate enough records to create a batch that has INITIAL_BUFFER_SIZE < batchSize < maxBatchSize
        List<String> records = new ArrayList<>();
        for (int i = 0; i < 1000000; 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(INITIAL_BUFFER_SIZE < batchSize && batchSize <= maxBatchSize);

        // Verify that the next buffer retrieved from the bufferSupplier is the expanded buffer.
        assertTrue(ctx.bufferSupplier.get(1).capacity() > INITIAL_BUFFER_SIZE);
    }

    @Test
    public void testBufferShrinkWhenMaxMessageSizeReducedBelowInitialBufferSize() {
        MockTimer timer = new MockTimer();
        var mockWriter = new InMemoryPartitionWriter(false) {
            private LogConfig config = new LogConfig(Map.of(
                TopicConfig.MAX_MESSAGE_BYTES_CONFIG, String.valueOf(1024 * 1024) // 1MB
            ));

            @Override
            public LogConfig config(TopicPartition tp) {
                return config;
            }

            public void updateConfig(LogConfig newConfig) {
                this.config = newConfig;
            }
        };