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

Code before:
                        .withSerializer(new StringSerializer())
                        .withAppendLingerMs(10)
                        .withExecutorService(mock(ExecutorService.class))
                        .build();

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

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

        // Write #1.
        CompletableFuture<String> write1 = runtime.scheduleWriteOperation(
                "write#1", TP, Duration.ofMillis(20),
                state -> new CoordinatorResult<>(List.of("record1"), "response1")
        );
        assertFalse(write1.isDone());
        assertNotNull(ctx.currentBatch);
        assertEquals(0, writer.entries(TP).size());

        // Verify that the linger timeout task is created; there will also be a default write timeout task.
        assertEquals(2, schedulerTimer.size());

        // Advance past the linger time.
        clockTimer.advanceClock(11);

        // At this point, there are still two scheduled tasks; the linger task has not fired
        // because we did not advance the schedulerTimer.
        assertEquals(2, schedulerTimer.size());

        // Write #2.
        CompletableFuture<String> write2 = runtime.scheduleWriteOperation(
                "write#2", TP, Duration.ofMillis(20),
                state -> new CoordinatorResult<>(List.of("record2"), "response2")
        );

        // The batch should have been flushed.
        assertEquals(1, writer.entries(TP).size());

        // Because flushing the batch cancels the linger task, there should now be two write timeout tasks.
        assertEquals(2, schedulerTimer.size());

        // Verify batch contains both two records
        MemoryRecords batch = writer.entries(TP).get(0);
        RecordBatch recordBatch = batch.firstBatch();
        assertEquals(2, recordBatch.countOrNull());

        // Commit and verify that writes are completed.
        writer.commit(TP);

Review comment (performance):
Could you move this to line#4921? It makes more sense there, since the goal of advancing the `clockTimer` is to ensure `flushCurrentBatch` is executed during writing #2, right?

Code after:
                .withSerializer(new StringSerializer())
                .withAppendLingerMs(10)
                .withExecutorService(mock(ExecutorService.class))
                .build();

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

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

        // Write #1.
        CompletableFuture<String> write1 = runtime.scheduleWriteOperation("write#1", TP, Duration.ofMillis(20),
            state -> new CoordinatorResult<>(List.of("record1"), "response1")
        );
        assertFalse(write1.isDone());
        assertNotNull(ctx.currentBatch);
        assertEquals(0, writer.entries(TP).size());

        // Verify that the linger timeout task is created; there will also be a default write timeout task.
        assertEquals(2, schedulerTimer.size());

        // Advance past the linger time.
        clockTimer.advanceClock(11);

        // At this point, there are still two scheduled tasks; the linger task has not fired
        // because we did not advance the schedulerTimer.
        assertEquals(2, schedulerTimer.size());

        // Write #2.
        CompletableFuture<String> write2 = runtime.scheduleWriteOperation("write#2", TP, Duration.ofMillis(20),
            state -> new CoordinatorResult<>(List.of("record2"), "response2")
        );

        // The batch should have been flushed.
        assertEquals(1, writer.entries(TP).size());

        // Because flushing the batch cancels the linger task, there should now be two write timeout tasks.
        assertEquals(2, schedulerTimer.size());

        // Verify batch contains both two records
        MemoryRecords batch = writer.entries(TP).get(0);
        RecordBatch recordBatch = batch.firstBatch();
        assertEquals(2, recordBatch.countOrNull());

        // Commit and verify that writes are completed.
        writer.commit(TP);
        assertTrue(write1.isDone());
        assertTrue(write2.isDone());