Repository: apache/kafka
File: storage/src/test/java/org/apache/kafka/storage/internals/log/LogConcurrencyTest.java (Java)

Code before:

    @BeforeEach
    public void setup() {
        scheduler.startup();
    }

    @AfterEach
    public void teardown() throws Exception {
        scheduler.shutdown();
        if (log != null) {
            log.close();
        }
        Utils.delete(tmpDir);
    }

    @Test
    public void testUncommittedDataNotConsumed() throws Exception {
        testUncommittedDataNotConsumed(createLog());
    }

    @Test
    public void testUncommittedDataNotConsumedFrequentSegmentRolls() throws Exception {
        final Properties logProps = new Properties();
        logProps.put(LogConfig.INTERNAL_SEGMENT_BYTES_CONFIG, 237);
        final LogConfig logConfig = new LogConfig(logProps);
        testUncommittedDataNotConsumed(createLog(logConfig));
    }

    private void testUncommittedDataNotConsumed(UnifiedLog log) throws Exception {
        ExecutorService executor = Executors.newFixedThreadPool(2);
        try {
            final int maxOffset = 5000;
            final ConsumerTask consumer = new ConsumerTask(log, maxOffset);
            final LogAppendTask appendTask = new LogAppendTask(log, maxOffset);

            final Future<List<FetchedBatch>> consumerFuture = executor.submit(consumer);
            final Future<Void> fetcherTaskFuture = executor.submit(appendTask);

            fetcherTaskFuture.get();
            List<FetchedBatch> consumedBatches = consumerFuture.get();

            validateConsumedData(log, consumedBatches);
        } finally {
            executor.shutdownNow();
            assertTrue(executor.awaitTermination(5L, TimeUnit.SECONDS));
        }
    }

    private UnifiedLog createLog() throws IOException {
        return createLog(new LogConfig(Map.of()));
    }

Review comment (performance):
Could you please try to use `CompletableFuture`?
```java
private void testUncommittedDataNotConsumed(UnifiedLog log) {
        var maxOffset = 5000;
        var consumerFuture = CompletableFuture.supplyAsync(consumerTask(log, maxOffset));
        var fetcherTaskFuture = CompletableFuture.runAsync(logAppendTask(log, maxOffset));
        fetcherTaskFuture.join();
        validateConsumedData(log, consumerFuture.join());
    }

    private Supplier<List<FetchedBatch>> consumerTask(UnifiedLog log, int lastOffset) {
        return () -> assertDoesNotThrow(() -> {
            final List<FetchedBatch> consumedBatches = new ArrayList<>(); // Now a local variable
            long fetchOffset = 0L;
            while (log.highWatermark() < lastOffset) {
                final FetchDataInfo readInfo = log.read(fetchOffset, 1, FetchIsolation.HIGH_WATERMARK, true);
                for (RecordBatch batch : readInfo.records.batches()) {
                    consumedBatches.add(new FetchedBatch(batch.baseOffset(), batch.partitionLeaderEpoch()));
                    fetchOffset = batch.lastOffset() + 1;
                }
            }
            return consumedBatches;
        });
    }

    private Runnable logAppendTask(UnifiedLog log, int lastOffset) {
        return () -> assertDoesNotThrow(() -> {
            int leaderEpoch = 1;
            boolean isLeader = true;
            while (log.highWatermark() < lastOffset) {
                switch (TestUtils.RANDOM.nextInt(2)) {
                    case 0 -> {
                        final LogOffsetMetadata logEndOffsetMetadata = log.logEndOffsetMetadata();
                        final long logEndOffset = logEndOffsetMetadata.messageOffset;
                        final int batchSize = TestUtils.RANDOM.nextInt(9) + 1;
                        final SimpleRecord[] records = IntStream.rangeClosed(0, batchSize)
                                .mapToObj(i -> new SimpleRecord(String.valueOf(i).getBytes()))
                                .toArray(SimpleRecord[]::new);

                        if (isLeader) {
                            log.appendAsLeader(MemoryRecords.withRecords(Compression.NONE, records), leaderEpoch);
                            log.maybeIncrementHighWatermark(logEndOffsetMetadata);
                        } else {
                            log.appendAsFollower(
                                    MemoryRecords.withRecords(
                                            logEndOffset,
                                            Compression.NONE,
                                            leaderEpoch,
                                            records
                                    ),
                                    Integer.MAX_VALUE
                            );
                            log.updateHighWatermark(logEndOffset);
                        }
                    }
                    case 1 -> {
                        isLeader = !isLeader;
                        leaderEpoch += 1;

                        if (!isLeader) {
                            log.truncateTo(log.highWatermark());
                        }
                    }
                }
            }
        });
    }
```

Code after:
        scheduler.startup();
    }

    @AfterEach
    public void teardown() throws Exception {
        scheduler.shutdown();
        if (log != null) {
            log.close();
        }
        Utils.delete(tmpDir);
    }

    @Test
    public void testUncommittedDataNotConsumed() throws Exception {
        testUncommittedDataNotConsumed(createLog());
    }

    @Test
    public void testUncommittedDataNotConsumedFrequentSegmentRolls() throws Exception {
        final Properties logProps = new Properties();
        logProps.put(LogConfig.INTERNAL_SEGMENT_BYTES_CONFIG, 237);
        final LogConfig logConfig = new LogConfig(logProps);
        testUncommittedDataNotConsumed(createLog(logConfig));
    }

    private void testUncommittedDataNotConsumed(UnifiedLog log) throws Exception {
        var maxOffset = 5000;
        var consumerFuture = CompletableFuture.supplyAsync(consumerTask(log, maxOffset));
        var fetcherTaskFuture = CompletableFuture.runAsync(logAppendTask(log, maxOffset));

        fetcherTaskFuture.join();
        validateConsumedData(log, consumerFuture.join());
    }

    /**
     * Simple consumption task which reads the log in ascending order and collects
     * consumed batches for validation
     */
    private Supplier<Map<Long, Integer>> consumerTask(UnifiedLog log, int lastOffset) {
        return () -> assertDoesNotThrow(() -> {
            final Map<Long, Integer> consumedBatches = new LinkedHashMap<>();
            long fetchOffset = 0L;
            while (log.highWatermark() < lastOffset) {
                final FetchDataInfo readInfo = log.read(fetchOffset, 1, FetchIsolation.HIGH_WATERMARK, true);
                for (RecordBatch batch : readInfo.records.batches()) {
                    consumedBatches.put(batch.baseOffset(), batch.partitionLeaderEpoch());
                    fetchOffset = batch.lastOffset() + 1;
                }
            }
            return consumedBatches;
        });