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

Code before:

    /**
     * Get the deque for the given topic-partition, creating it if necessary.
     */
    private Deque<ProducerBatch> getOrCreateDeque(TopicPartition tp) {
        TopicInfo topicInfo = topicInfoMap.computeIfAbsent(tp.topic(),
                k -> new TopicInfo(createBuiltInPartitioner(logContext, k, batchSize)));
        return topicInfo.batches.computeIfAbsent(tp.partition(), k -> new ArrayDeque<>());
    }

    BuiltInPartitioner createBuiltInPartitioner(LogContext logContext, String topic, int stickyBatchSize) {
        return new BuiltInPartitioner(logContext, topic, stickyBatchSize);
    }

    /**
     * Deallocate the record batch
     */
    public void deallocate(ProducerBatch batch) {
        incomplete.remove(batch);
        deallocateIfNecessary(batch);
    }

    private void deallocateIfNecessary(ProducerBatch batch) {
        // Only deallocate the batch if it is not a split batch because split batch are allocated outside the
        // buffer pool.
        if (!batch.isSplitBatch() && !batch.isBufferDeallocated() && batch.setBufferDeallocated(true))
            free.deallocate(batch.buffer(), batch.initialCapacity());
    }

    /**
     * Remove from the incomplete list but do not free memory yet
     */
    public void deallocateLater(ProducerBatch batch) {
        incomplete.remove(batch);
    }

    /**
     * Only perform deallocation (and not removal from incomplete set)
     */
    public void deallocateAlreadyRemovedIncomplete(ProducerBatch batch) {
        deallocateIfNecessary(batch);
    }

    /**
     * Package private for unit test. Get the buffer pool remaining size in bytes.
     */
    long bufferPoolAvailableMemory() {
        return free.availableMemory();
    }

    /**

Review comment (performance):
Could we log a warn if this method is called but `batch.isBufferDeallocated()` is true?

Code after:

    /**
     * Get the deque for the given topic-partition, creating it if necessary.
     */
    private Deque<ProducerBatch> getOrCreateDeque(TopicPartition tp) {
        TopicInfo topicInfo = topicInfoMap.computeIfAbsent(tp.topic(),
                k -> new TopicInfo(createBuiltInPartitioner(logContext, k, batchSize)));
        return topicInfo.batches.computeIfAbsent(tp.partition(), k -> new ArrayDeque<>());
    }

    BuiltInPartitioner createBuiltInPartitioner(LogContext logContext, String topic, int stickyBatchSize) {
        return new BuiltInPartitioner(logContext, topic, stickyBatchSize);
    }

    /**
     * Complete and deallocate the record batch
     */
    public void completeAndDeallocateBatch(ProducerBatch batch) {
        completeBatch(batch);
        deallocate(batch);
    }

    /**
     * Only perform deallocation (and not removal from the incomplete set)
     */
    public void deallocate(ProducerBatch batch) {
        // Only deallocate the batch if it is not a split batch because split batch are allocated outside the
        // buffer pool.
        if (!batch.isSplitBatch()) {
            if (batch.isBufferDeallocated()) {
                log.warn("Skipping deallocating a batch that has already been deallocated. Batch is {}, created time is {}", batch, batch.createdMs);
            } else {
                batch.markBufferDeallocated();
                if (batch.isInflight()) {
                    // Create a fresh ByteBuffer to give to BufferPool to reuse since we can't safely call deallocate with the ProduceBatch's buffer
                    free.deallocate(ByteBuffer.allocate(batch.initialCapacity()));
                    throw new IllegalStateException("Attempting to deallocate a batch that is inflight. Batch is " + batch);
                }
                free.deallocate(batch.buffer(), batch.initialCapacity());
            }
        }
    }

    /**
     * Remove from the incomplete list but do not free memory yet
     */
    public void completeBatch(ProducerBatch batch) {
        incomplete.remove(batch);
    }

    /**