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

Code before:
        assertEquals(List.of(0L), ctx.coordinator.snapshotRegistry().epochsList());

        int maxBatchSize = writer.config(TP).maxMessageSize();
        assertTrue(maxBatchSize > INITIAL_BUFFER_SIZE);

        // Generate enough records to create a batch that has INITIAL_BUFFER_SIZE < batchSize < maxBatchSize
        List<String> records = new ArrayList<>();
        for (int i = 0; i < 50000; 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 = writer.entries(TP).get(0).sizeInBytes();
        assertTrue(batchSize > INITIAL_BUFFER_SIZE && batchSize < maxBatchSize);
    }

    @Test
    public void testCoordinatorDoNotRetainBufferLargeThanMaxBufferSize() {
        MockTimer timer = new MockTimer();
        InMemoryPartitionWriter mockWriter = new InMemoryPartitionWriter(false);
        StringSerializer serializer = new StringSerializer();

        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
                .withTime(timer.time())
                .withTimer(timer)
                .withDefaultWriteTimeOut(DEFAULT_WRITE_TIMEOUT)
                .withLoader(new MockCoordinatorLoader())
                .withEventProcessor(new DirectEventProcessor())
                .withPartitionWriter(mockWriter)
                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
                .withCoordinatorRuntimeMetrics(mock(CoordinatorRuntimeMetrics.class))
                .withCoordinatorMetrics(mock(CoordinatorMetrics.class))
                .withSerializer(serializer)
                .withExecutorService(mock(ExecutorService.class))
                .withMaxBufferSizeSupplier(() -> MAX_BUFFER_SIZE)
                .build();

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

        // Verify the initial state.
        CoordinatorRuntime<MockCoordinatorShard, String>.CoordinatorContext ctx = runtime.contextOrThrow(TP);

Review comment (performance):
Could we assert that `recordAppendBufferDiscarded` is called on a mock `CoordinatorRuntimeMetrics`? The same for `testBufferShrinkWhenMaxBufferSizeReducedBelowBatchSize`.

Code after:

        int maxBatchSize = writer.config(TP).maxMessageSize();
        assertTrue(maxBatchSize > INITIAL_BUFFER_SIZE);

        // Generate enough records to create a batch that has INITIAL_BUFFER_SIZE < batchSize < maxBatchSize
        List<String> records = new ArrayList<>();
        for (int i = 0; i < 50000; 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 = writer.entries(TP).get(0).sizeInBytes();
        assertTrue(batchSize > INITIAL_BUFFER_SIZE && batchSize < maxBatchSize);
    }

    @Test
    public void testCoordinatorDoNotRetainBufferLargeThanCachedBufferMaxBytes() {
        MockTimer timer = new MockTimer();
        InMemoryPartitionWriter mockWriter = new InMemoryPartitionWriter(false);
        StringSerializer serializer = new StringSerializer();
        CoordinatorRuntimeMetrics runtimeMetrics = mock(CoordinatorRuntimeMetrics.class);

        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
                .withTime(timer.time())
                .withTimer(timer)
                .withDefaultWriteTimeOut(DEFAULT_WRITE_TIMEOUT)
                .withLoader(new MockCoordinatorLoader())
                .withEventProcessor(new DirectEventProcessor())
                .withPartitionWriter(mockWriter)
                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
                .withCoordinatorRuntimeMetrics(runtimeMetrics)
                .withCoordinatorMetrics(mock(CoordinatorMetrics.class))
                .withSerializer(serializer)
                .withExecutorService(mock(ExecutorService.class))
                .withCachedBufferMaxBytesSupplier(() -> CACHED_BUFFER_MAX_BYTES)
                .build();

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

        // Verify the initial state.
        CoordinatorRuntime<MockCoordinatorShard, String>.CoordinatorContext ctx = runtime.contextOrThrow(TP);