Kafka 客户端的缓存管理

public final class RecordAccumulator {
    // 消息缓存队列，以batch格式存储
    private final ConcurrentMap<TopicPartition, Deque<ProducerBatch>> batches;    
    // 内存池
    private final BufferPool free;

    public RecordAppendResult append(....) {
        // 计算该消息占用的内存大小    
        int size = Math.max(this.batchSize, AbstractRecords.estimateSizeInBytesUpperBound(maxUsableMagic, compression, key, value, headers));
        // 向内存池申请内存
        buffer = free.allocate(size, maxTimeToBlock);
        .......
    }
}

public final class RecordAccumulator {
    // 内存池
    private final BufferPool free;

    public void deallocate(ProducerBatch batch) {
        incomplete.remove(batch);
        // 检测是否该消息batch数据太大了，如果太大了则需要切割。所以这种情况不需要释放内存
        if (!batch.isSplitBatch())
            // 向内存池释放内存
            free.deallocate(batch.buffer(), batch.initialCapacity());
    }    
}

public class BufferPool {
    private final long totalMemory;  // 整个内存池的总容量
    private long nonPooledAvailableMemory;  // 非缓存池的空闲大小    
    private final int poolableSize;  // 缓存池中，特定内存的大小
    private final ReentrantLock lock;  // 锁用来防治并发
    private final Deque<ByteBuffer> free;  // 缓存池中的空闲内存块
    private final Deque<Condition> waiters;  // 等待申请的用户

    // 参数size表示申请的内存大小，参数maxTimeToBlockMs表示等待的最长时间
    public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedException {
        if (size > this.totalMemory)
            throw new IllegalArgumentException("...")
        ByteBuffer buffer = null;
        this.lock.lock();
        // 如果申请大小等于指定的值，并且缓存池中有空闲的内存块，则直接返回
         if (size == poolableSize && !this.free.isEmpty())
             return this.free.pollFirst();
        // 计算总的空闲内存大小， 等于缓存池的大小 + 非缓存池的空闲大小
        int freeListSize = freeSize() * this.poolableSize;
        if (this.nonPooledAvailableMemory + freeListSize >= size) {
            // 如果空闲内存足够，那么需要保证非缓存池的空闲空间足够
            // 因为所有的内存分配都是从非缓存池开始
            freeUp(size);
            this.nonPooledAvailableMemory -= size;
        } else {
            // 如果空闲内存不够，那么需要等待别的用户释放内存
            int accumulated = 0; // 表示已经成功申请的内存大小
            Condition moreMemory = this.lock.newCondition();

            long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs);
            // 添加到等待集合中
            this.waiters.addLast(moreMemory);
            // 循环等待
            while (accumulated < size) {
                ......
                // 等待通知
                waitingTimeElapsed = !moreMemory.await(remainingTimeToBlockNs, TimeUnit.NANOSECONDS);
                if (waitingTimeElapsed) {
                    // 超过等待时间还没有分配到足够的内存，那么就抛出异常
                    throw new TimeoutException("Failed to allocate memory within the configured max blocking time " + maxTimeToBlockMs + " ms.");
                }
                // 如果申请大小等于指定的值，而这时刚好有其他用户释放了内存，那么就直接从缓存池中获取
                if (accumulated == 0 && size == this.poolableSize && !this.free.isEmpty()) {
                    buffer = this.free.pollFirst();
                    accumulated = size;
                } else {
                    // 继续释放足够的内存
                    freeUp(size - accumulated);
                    // 计算可以分配的内存大小
                    int got = (int) Math.min(size - accumulated, this.nonPooledAvailableMemory);
                    // 更改非缓存池的空闲内存大小
                    this.nonPooledAvailableMemory -= got;
                    accumulated += got;
                }
            }

            // 这里会检查是否还有剩下的空闲内存，如果有则需要通知下一个用户。
            // 因为这时可能有多个用户都释放了内存，而用户释放内存只会通知第一个用户（也就是当前用户），而下个用户还在一直等待中，如果当前用户不主动通知的话，可能造成下个用户等待超时。
            if (!(this.nonPooledAvailableMemory == 0 && this.free.isEmpty()) && !this.waiters.isEmpty())
                    this.waiters.peekFirst().signal();

            // 运行到这里，说明已经成功了申请到了内存          
            if (buffer == null)
                // 表示buffer不是从缓存池中获取的，需要执行内存分配
                return safeAllocateByteBuffer(size);
            else:
                // 已经从缓存池中获取到了，则直接返回
                return buffer;
        }
    }
}

public class BufferPool {

    public void deallocate(ByteBuffer buffer) {
        deallocate(buffer, buffer.capacity());
    }

    public void deallocate(ByteBuffer buffer, int size) {
        lock.lock();
        try {
            // 如果释放的内存大小等于指定的值，那么就将它添加到缓存池。free列表存储这些内存块
            if (size == this.poolableSize && size == buffer.capacity()) {
                buffer.clear();
                this.free.add(buffer);
            } else {
                // 否则更新非缓存池的空闲大小，这个ByteBuffer实例等待jvm自动gc
                this.nonPooledAvailableMemory += size;
            }
            // 通知队列的第一个用户
            Condition moreMem = this.waiters.peekFirst();
            if (moreMem != null)
                moreMem.signal();
        } finally {
            lock.unlock();
        }
    }
}

public class KafkaConsumer<K, V> implements Consumer<K, V> {

    private final Fetcher<K, V> fetcher; // 消息缓存

    // poll 方法做过简化，省略了Metadata和topic subcribe的步骤 
    private ConsumerRecords<K, V> poll(final long timeoutMs, final boolean includeMetadataInTimeout) {
        // 调用 pollForFetches 获取消息
        final Map<TopicPartition, List<ConsumerRecord<K, V>>> records = pollForFetches(remainingTimeAtLeastZero(timeoutMs, elapsedTime));
        if (!records.isEmpty()) {
            // 注意到这里，成功获取之后，还会调用sendFetches方法，发送请求
            // 因为kafka发送请求都是异步的，这里提前发出请求，可以有效的减少下次缓存为空而需要等待请求的时间
            if (fetcher.sendFetches() > 0 || client.hasPendingRequests()) {
                client.pollNoWakeup();
            }
            return this.interceptors.onConsume(new ConsumerRecords<>(records));
        }
    }

    private Map<TopicPartition, List<ConsumerRecord<K, V>>> pollForFetches(final long timeoutMs) {
        // 首先从缓存中尝试读取
        final Map<TopicPartition, List<ConsumerRecord<K, V>>> records = fetcher.fetchedRecords();
        if (!records.isEmpty()) {
            return records;
        }

        // 如果缓存为空，则向服务端发出请求
        fetcher.sendFetches();
        // 等待服务端的响应
        client.poll(pollTimeout, startMs, () -> {
            // hasCompletedFetches 表示是否有成功的响应
            return !fetcher.hasCompletedFetches();
        });
        // 从缓存中获取消息
        return fetcher.fetchedRecords();
    }
}

public class Fetcher<K, V> implements SubscriptionState.Listener, Closeable {
    // CompletedFetch代表着响应，对应着一个分区的消息
    private final ConcurrentLinkedQueue<CompletedFetch> completedFetches;

    private List<TopicPartition> fetchablePartitions() {
        // 存储着哪些分区不需要请求
        Set<TopicPartition> exclude = new HashSet<>();  
        // 获取分配的分区
        List<TopicPartition> fetchable = subscriptions.fetchablePartitions();
        // nextInLineRecords 表示正在解析的响应
        if (nextInLineRecords != null && !nextInLineRecords.isFetched) {
            exclude.add(nextInLineRecords.partition);
        }
        // 如果在缓存里，该分区已经有了消息，则不需要请求
        for (CompletedFetch completedFetch : completedFetches) {
            exclude.add(completedFetch.partition);
        }
        // 剔除掉那些不需要请求的分区
        fetchable.removeAll(exclude);
        return fetchable;
    }
}