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

Code before:
            records(writeTimestamp, "record2")
        ), writer.entries(TP));

        // There is no pending high watermark.
        assertEquals(NO_OFFSET, runtime.contextOrThrow(TP).highWatermarklistener.lastHighWatermark());

        // Advance the clock then commit records from write#1.
        timer.advanceClock(700);
        writer.commit(TP, 1);

        // We should still have one pending event and the pending high watermark should be updated.
        assertEquals(1, processor.size());
        assertEquals(1, runtime.contextOrThrow(TP).highWatermarklistener.lastHighWatermark());

        // Poll once to process the high watermark update and complete the writes.
        processor.poll();
        long purgatoryTimeMs = timer.time().milliseconds() - writeTimestamp;

        // Advance the clock past write timeout. write#2 has now timed out.
        timer.advanceClock(300 + 1);
        processor.poll();

        assertEquals(NO_OFFSET, runtime.contextOrThrow(TP).highWatermarklistener.lastHighWatermark());
        assertEquals(1, runtime.contextOrThrow(TP).coordinator.lastCommittedOffset());
        assertTrue(write1.isDone());
        assertTrue(write2.isCompletedExceptionally());
        verify(runtimeMetrics, times(1)).recordEventPurgatoryTime(purgatoryTimeMs);
        verify(runtimeMetrics, times(1)).recordEventPurgatoryTime(writeTimeout.toMillis() + 1);
    }

    @Test
    public void testWriteEventCompletesOnlyOnce() throws Exception {
        // Completes once via timeout, then again with HWM update.
        Duration writeTimeout = Duration.ofMillis(1000L);
        MockTimer timer = new MockTimer();
        MockPartitionWriter writer = new MockPartitionWriter();
        ManualEventProcessor processor = new ManualEventProcessor();
        CoordinatorRuntimeMetrics runtimeMetrics = mock(CoordinatorRuntimeMetrics.class);

        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
                .withTime(timer.time())
                .withTimer(timer)
                .withDefaultWriteTimeOut(writeTimeout)
                .withLoader(new MockCoordinatorLoader())
                .withEventProcessor(processor)
                .withPartitionWriter(writer)
                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
                .withCoordinatorRuntimeMetrics(runtimeMetrics)
                .withCoordinatorMetrics(mock(CoordinatorMetrics.class))
                .withSerializer(new StringSerializer())

Review comment (bug):
Can we check that write2 fails with an exception?
```
        assertFutureThrows(RecordTooLargeException.class, write2);
```

Code after:
        // so the write should fail with RecordTooLargeException
        CompletableFuture<String> write2 = runtime.scheduleWriteOperation("write#2", TP, Duration.ofMillis(50),
            state -> new CoordinatorResult<>(largeRecord, "response2")
        );

        // Check that write2 fails with RecordTooLargeException
        assertFutureThrows(RecordTooLargeException.class, write2);

        // Verify the state. The first batch was flushed and the largeRecord
        // write failed.
        assertEquals(2L, ctx.coordinator.lastWrittenOffset());
        assertEquals(0L, ctx.coordinator.lastCommittedOffset());
        assertEquals(List.of(
            new MockCoordinatorShard.RecordAndMetadata(0, records.get(0)),
            new MockCoordinatorShard.RecordAndMetadata(1, records.get(1))
        ), ctx.coordinator.coordinator().fullRecords());
        assertEquals(List.of(
            records(firstBatchTimestamp, compression, records)
        ), writer.entries(TP));

        // Commit and verify that writes are completed.
        writer.commit(TP);
        assertTrue(write1.isDone());
        assertTrue(write2.isDone());
        assertEquals(2L, ctx.coordinator.lastCommittedOffset());
        assertEquals("response1", write1.get(5, TimeUnit.SECONDS));
    } 

    @Test
    public void testRecordEventPurgatoryTime() throws Exception {
        Duration writeTimeout = Duration.ofMillis(1000);
        MockTimer timer = new MockTimer();
        MockPartitionWriter writer = new MockPartitionWriter();
        ManualEventProcessor processor = new ManualEventProcessor();
        CoordinatorRuntimeMetrics runtimeMetrics = mock(CoordinatorRuntimeMetrics.class);

        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
                .withTime(timer.time())
                .withTimer(timer)
                .withDefaultWriteTimeOut(writeTimeout)
                .withLoader(new MockCoordinatorLoader())
                .withEventProcessor(processor)
                .withPartitionWriter(writer)
                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
                .withCoordinatorRuntimeMetrics(runtimeMetrics)
                .withCoordinatorMetrics(mock(CoordinatorMetrics.class))
                .withSerializer(new StringSerializer())
                .withExecutorService(mock(ExecutorService.class))
                .build();