Repository: apache/kafka
File: clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java (Java)

Code before:
        assertFalse(result.readyNodes.isEmpty(), "The batch should be ready");
        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
        assertEquals(1, drained.size(), "Only node1 should be drained");
        assertEquals(1, drained.get(node1.id()).size(), "Only one batch should be drained");
        // Split and reenqueue the batch.
        accum.splitAndReenqueue(drained.get(node1.id()).get(0));
        time.sleep(101L);

        drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
        assertFalse(drained.isEmpty());
        assertFalse(drained.get(node1.id()).isEmpty());
        drained.get(node1.id()).get(0).complete(acked.get(), 100L);
        assertEquals(1, acked.get(), "The first message should have been acked.");
        assertTrue(future1.isDone());
        assertEquals(0, future1.get().offset());

        drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
        assertFalse(drained.isEmpty());
        assertFalse(drained.get(node1.id()).isEmpty());
        drained.get(node1.id()).get(0).complete(acked.get(), 100L);
        assertEquals(2, acked.get(), "Both message should have been acked.");
        assertTrue(future2.isDone());
        assertEquals(1, future2.get().offset());
    }

    // This test confirms us that splitting a single large record
    // creates an unsplittable batch (does not really split it)
    // that will continue to fail with MESSAGE_TOO_LARGE,
    // causing infinite retry loops
    @Test
    public void testSplitAndReenqueueWithSingleLargeRecord() throws ExecutionException, InterruptedException {
        long now = time.milliseconds();
        int smallBatchSize = 1024;
        RecordAccumulator accum = createTestRecordAccumulator(smallBatchSize, 10 * 1024, Compression.NONE, 10);

        // create a single record that is much larger than the batch size limit
        // we are trying to mimic by send a record larger than broker's message.max.bytes
        byte[] largeValue = new byte[4 * 1024]; // 4KB > 1KB
        
        // Create a buffer with enough space for the large record
        ByteBuffer buffer = ByteBuffer.allocate(8192);
        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, Compression.NONE, TimestampType.CREATE_TIME, 0L);
        ProducerBatch batch = new ProducerBatch(tp1, builder, now, true);
        
        final AtomicInteger acked = new AtomicInteger(0);
        Callback cb = (metadata, exception) -> acked.incrementAndGet();

        // create a large batch but only with one single record
        Future<RecordMetadata> future = batch.tryAppend(now, key, largeValue, Record.EMPTY_HEADERS, cb, now);
        assertNotNull(future, "Should be able to append the large record to batch");
        assertEquals(1, batch.recordCount, "Batch should contain exactly one record");

Review comment (performance):
We have the following code in Sender. It only tries to split the batch if it has more than 1 record. So, the splitting shouldn't occur forever.

```
    private void completeBatch(ProducerBatch batch, ProduceResponse.PartitionResponse response, long correlationId,
                               long now, Map<TopicPartition, Metadata.LeaderIdAndEpoch> partitionsWithUpdatedLeaderInfo) {
        Errors error = response.error;

        if (error == Errors.MESSAGE_TOO_LARGE && batch.recordCount > 1 && !batch.isDone() &&
                (batch.magic() >= RecordBatch.MAGIC_VALUE_V2 || batch.isCompressed())) {
            ...
            this.accumulator.splitAndReenqueue(batch);
```

Code after:
        batch.close();
        // Enqueue the batch to the accumulator as if the batch was created by the accumulator.
        accum.reenqueue(batch, now);
        // Re-enqueuing counts as a second attempt, so the delay with jitter is 100 * (1 + 0.2) + 1
        time.sleep(121L);
        // Drain the batch.
        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
        assertFalse(result.readyNodes.isEmpty(), "The batch should be ready");
        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
        assertEquals(1, drained.size(), "Only node1 should be drained");
        assertEquals(1, drained.get(node1.id()).size(), "Only one batch should be drained");
        // Split and reenqueue the batch.
        accum.splitAndReenqueue(drained.get(node1.id()).get(0));
        time.sleep(101L);

        drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
        assertFalse(drained.isEmpty());
        assertFalse(drained.get(node1.id()).isEmpty());
        drained.get(node1.id()).get(0).complete(acked.get(), 100L);
        assertEquals(1, acked.get(), "The first message should have been acked.");
        assertTrue(future1.isDone());
        assertEquals(0, future1.get().offset());

        drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
        assertFalse(drained.isEmpty());
        assertFalse(drained.get(node1.id()).isEmpty());
        drained.get(node1.id()).get(0).complete(acked.get(), 100L);
        assertEquals(2, acked.get(), "Both message should have been acked.");
        assertTrue(future2.isDone());
        assertEquals(1, future2.get().offset());
    }

    // here I am testing the hasRoomFor() behaviour
    // It allows the first record no matter the size
    // but does not allow the second record
    @Test
    public void testHasRoomForAllowsOversizedFirstRecordButRejectsSubsequentRecords() {
        long now = time.milliseconds();
        int smallBatchSize = 1024;

        // Create a large record that exceeds batch size limit
        byte[] largeValue = new byte[4 * 1024]; // 4KB > 1KB

        // Create a small buffer that cannot fit the large record
        ByteBuffer buffer = ByteBuffer.allocate(smallBatchSize);
        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, Compression.NONE, TimestampType.CREATE_TIME, 0L);

        // testing existing code:
        // hasRoomFor() should return true for first record regardless of size
        boolean hasRoomForFirst = builder.hasRoomFor(now, ByteBuffer.wrap(key), ByteBuffer.wrap(largeValue), Record.EMPTY_HEADERS);
        assertTrue(hasRoomForFirst, "hasRoomFor() should return true for first record regardless of size when numRecords == 0");