Kafka2.0生产者客户端源码分析

1 KafkaProducer 构造器

1. 初始化参数配置。
2. 初始化记录累加器 RecordAccumulator。
3. 初始化 Kafka 连接 KafkaClient，发现集群的所有节点加入缓存。
4. 初始化实现了 Runnable 接口的 Sender 对象，并在 ioThread 中启动线程。

2 发送消息

1. 执行消息拦截器
2. 查询 Kafka 集群元数据
3. 序列化 key、value
4. 获取分区
5. 把消息添加到记录累加器中
6. 当 batch 满了，或者创建了新的 batch 后，唤醒 Sender 线程

  核心源码如下

public Future<RecordMetadata> send(ProducerRecord<K, V> record, Callback callback) {
	// 执行拦截器
	ProducerRecord<K, V> interceptedRecord = this.interceptors.onSend(record);
	return doSend(interceptedRecord, callback);
}
private Future<RecordMetadata> doSend(ProducerRecord<K, V> record, Callback callback) {
	TopicPartition tp = null;
	// 获取元数据
	ClusterAndWaitTime clusterAndWaitTime = waitOnMetadata(record.topic(), record.partition(), maxBlockTimeMs);
	Cluster cluster = clusterAndWaitTime.cluster;
	// 序列化 key、value
	byte[] serializedKey = keySerializer.serialize(record.topic(), record.headers(), record.key());
	byte[] serializedValue = valueSerializer.serialize(record.topic(), record.headers(), record.value());
	// 获取分区。
	// 如果为空，会计算 key 的 hash 值，再和该主题的分区总数取余得到分区号；
	// 如果 key 也为空，客户端会生成递增的随机整数，再和该主题的分区总数区域得到分区号。
	int partition = partition(record, serializedKey, serializedValue, cluster);
	tp = new TopicPartition(record.topic(), partition);
	// 校验序列化后的记录是否超过限制
	int serializedSize = AbstractRecords.estimateSizeInBytesUpperBound(apiVersions.maxUsableProduceMagic(),
			compressionType, serializedKey, serializedValue, headers);
	ensureValidRecordSize(serializedSize);
	// 时间戳，默认是 KafkaProducer 初始化时间
	long timestamp = record.timestamp() == null ? time.milliseconds() : record.timestamp();
	// 初始化回调和响应的拦截器对象
	Callback interceptCallback = new InterceptorCallback<>(callback, this.interceptors, tp);
	// 把消息添加到记录累加器中
	RecordAccumulator.RecordAppendResult result = accumulator.append(tp, timestamp, serializedKey,
			serializedValue, headers, interceptCallback, remainingWaitMs);
	if (result.batchIsFull || result.newBatchCreated) {
		// 当 batch 满了，或者创建了新的 batch 后，唤醒 Sender 线程
		this.sender.wakeup();
	}
	return result.future;
}

 2.1 查询元数据

1. 如果根据指定的主题和分区能在缓存中查找到，则直接返回元数据，结束流程。
2. 否则，设置需要更新元数据的标记 needUpdate=true，并获取当前的 version。
3. 唤醒 Sender 线程，当 Sender 线程判断 needUpdate=true 时，发送获取元数据的请求到 broker，获取到后更新 needUpdate=true，version+1。
4. 当前线程判断，如果 version 变大，说明元数据已更新，则跳出循环，拉取新的元数据，判断是否匹配到主题和分区，如果没有匹配到，返回第2步。
5. 如果 version 没变大，说明元数据还没更新，则调用 wait(long timeout) 方法，等待 timeout 时间后，返回第4步。
6. 当第4步获取到匹配的元数据后，返回给 doSend 方法。

  核心源码如下

private ClusterAndWaitTime waitOnMetadata(String topic, Integer partition, long maxWaitMs) {
	// 获取缓存的集群信息
	Cluster cluster = metadata.fetch();
	Integer partitionsCount = cluster.partitionCountForTopic(topic);
	// Return cached metadata if we have it, and if the record's partition is either undefined
	// or within the known partition range
	// 如果缓存中的数据满足条件，直接返回缓存中的元数据。
	if (partitionsCount != null && (partition == null || partition < partitionsCount))
		return new ClusterAndWaitTime(cluster, 0);

	long begin = time.milliseconds();
	long remainingWaitMs = maxWaitMs;
	long elapsed;
	do {
		// 更新元数据的标记 needUpdate=true，并获取当前的 version。
		int version = metadata.requestUpdate();
		sender.wakeup(); // 唤醒 Sender 线程
		try {
			metadata.awaitUpdate(version, remainingWaitMs); // 等待更新
		} catch (TimeoutException ex) {
		}
		cluster = metadata.fetch(); // 重新获取元数据
		elapsed = time.milliseconds() - begin;
		if (elapsed >= maxWaitMs) // 超出最大等待时间，抛出异常
			throw new TimeoutException("Failed to update metadata after " + maxWaitMs + " ms.");
		remainingWaitMs = maxWaitMs - elapsed;
		partitionsCount = cluster.partitionCountForTopic(topic);
	} while (partitionsCount == null); // 分区数量是 0，继续上述循环
	if (partition != null && partition >= partitionsCount) { // 当指定的分区号大于等于分数总数时，异常
		throw new KafkaException(String.format("Invalid partition given with record: %d is not in the range [0...%d).", partition, partitionsCount));
	}
	return new ClusterAndWaitTime(cluster, elapsed);
}
// 等待更新
public synchronized void awaitUpdate(final int lastVersion, final long maxWaitMs) throws InterruptedException {
     long begin = System.currentTimeMillis();
     long remainingWaitMs = maxWaitMs;
     // 版本号<=当前版本号，说明未更新，需要继续循环等待更新
     while ((this.version <= lastVersion) && !isClosed()) {
         if (remainingWaitMs != 0)
             wait(remainingWaitMs); // 等待一会再判断
         long elapsed = System.currentTimeMillis() - begin;
         if (elapsed >= maxWaitMs) // 超过了最大等待时间
             throw new TimeoutException("Failed to update metadata after " + maxWaitMs + " ms.");
         remainingWaitMs = maxWaitMs - elapsed;
     }
 }

 2.2 消息添加到累加器 RecordAccumulator

  2.2.1 缓冲池 BufferPool

  Kafka 使用缓冲池技术给消息分配堆字节缓存 HeapByteBuffer，缓冲池的空闲队列 free 存放了空闲的缓存队列，优先直接从中取出第一个进行分配缓存，如果缓冲池不够了，利用 ReentrantLock + Condition 构造等待队列，等待缓冲池足够分配。
  Kafka 在处理消息响应时，释放分配的内存，并把加入空闲队列 free。

// 缓冲池
public class BufferPool {
	// 可用总内存 buffer.memory
    private final long totalMemory;
	// 一批消息的大小 batch.size
    private final int poolableSize;
    private final ReentrantLock lock;
	// 空闲缓存队列
    private final Deque<ByteBuffer> free;
	// 等待队列
    private final Deque<Condition> waiters;
    // 可用未分配的内存总量是nonPooledAvailableMemory和free * poolableSize中字节缓冲区的总和。
    private long nonPooledAvailableMemory;
}
// 字节缓冲分配
public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedException {
	if (size > this.totalMemory)
		throw new IllegalArgumentException("消息大小超过总内存");

	ByteBuffer buffer = null;
	this.lock.lock();
	try {
		// 直接在空闲队列分配
		if (size == poolableSize && !this.free.isEmpty())
			return this.free.pollFirst();

		// 计算空闲队列总大小
		int freeListSize = this.free.size() * this.poolableSize;
		if (this.nonPooledAvailableMemory + freeListSize >= size) { // 可用的总内存(未分配的+空闲队列)>消息大小
			// we have enough unallocated or pooled memory to immediately
			// satisfy the request, but need to allocate the buffer
			freeUp(size);
			this.nonPooledAvailableMemory -= size; // 未分配内存总数-消息大小
		} else { // 内存不够分配
			int accumulated = 0;
			Condition moreMemory = this.lock.newCondition();
			try {
				long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs);
				this.waiters.addLast(moreMemory); //　加入等待队列
				// loop over and over until we have a buffer or have reserved
				// enough memory to allocate one
				while (accumulated < size) {　//  轮询，直到足够分配内存
					long startWaitNs = time.nanoseconds();
					long timeNs;
					boolean waitingTimeElapsed;
					try { // 等待一段时间
						waitingTimeElapsed = !moreMemory.await(remainingTimeToBlockNs, TimeUnit.NANOSECONDS);
					}

					remainingTimeToBlockNs -= timeNs;

					// 直接在空闲队列分配
					if (accumulated == 0 && size == this.poolableSize && !this.free.isEmpty()) {
						buffer = this.free.pollFirst();
						accumulated = size;
					} else { // 内存不够，accumulated累加计数
						freeUp(size - accumulated);
						int got = (int) Math.min(size - accumulated, this.nonPooledAvailableMemory);
						this.nonPooledAvailableMemory -= got;
						accumulated += got;
					}
				}
				accumulated = 0; // 清空
			}
		}
	}

	if (buffer == null) // 没有在空闲队列分配到内存，需要在堆上分配内存
		return new HeapByteBuffer(size, size);
	else
		return buffer;
}
private void freeUp(int size) {
	while (!this.free.isEmpty() && this.nonPooledAvailableMemory < size)
		this.nonPooledAvailableMemory += this.free.pollLast().capacity(); // 释放空闲队列的内存
}
// 处理生产者响应消息时，释放分配的内存
public void deallocate(ByteBuffer buffer, int size) {
	lock.lock();
	try {
		if (size == this.poolableSize && size == buffer.capacity()) {
			buffer.clear();
			this.free.add(buffer); // 加到空闲队列
		} else {
			this.nonPooledAvailableMemory += size; // 增加未分配内存数量
		}
		Condition moreMem = this.waiters.peekFirst();
		if (moreMem != null)
			moreMem.signal();
	} finally {
		lock.unlock();
	}
}

  2.2.2 消息缓存 CopyOnWriteMap

  累加器使用 CopyOnWriteMap 来缓存消息，key 是主题分区信息，value 是个双端队列，队列中的对象是压缩后的批量消息。

// 累加器缓存
ConcurrentMap<TopicPartition, Deque<ProducerBatch>> batches = new CopyOnWriteMap<>();

  CopyOnWriteMap 是线程安全的，是由 Kafka 实现的写时复制 Map，内部定义了 volatile 的 Map，读时不用加锁，直接读取，写时需要加锁，然后拷贝一个 Map 副本进行实际的写入，写入完成后再把原来的 Map 指向修改后的 Map。
  双端队列 Deque 实际上就是 ArrayDeque，非线程安全的，需要手动同步。使用双端队列可以在消息发送失败时，把消息直接放回队列头部进行重试。

// 累加消息到缓存
public RecordAppendResult append(TopicPartition tp,
								 long timestamp,
								 byte[] key,
								 byte[] value,
								 Header[] headers,
								 Callback callback,
								 long maxTimeToBlock) throws InterruptedException {
	ByteBuffer buffer = null;
	try {
		Deque<ProducerBatch> dq = getOrCreateDeque(tp); // 检查 batches 是否有该分区的映射，如果没有，则创建一个
		synchronized (dq) { // 加锁后分配
			RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callback, dq);
			if (appendResult != null)
				return appendResult;
		}

		byte maxUsableMagic = apiVersions.maxUsableProduceMagic();
		// 计算消息大小
		int size = Math.max(this.batchSize, AbstractRecords.estimateSizeInBytesUpperBound(maxUsableMagic, compression, key, value, headers));
		buffer = free.allocate(size, maxTimeToBlock); // 利用 BufferPool 分配字节缓存
		synchronized (dq) { // 加锁后分配
			RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callback, dq);
			// 构造出压缩后的批量消息对象 ProducerBatch
			MemoryRecordsBuilder recordsBuilder = recordsBuilder(buffer, maxUsableMagic);
			ProducerBatch batch = new ProducerBatch(tp, recordsBuilder, time.milliseconds());
			FutureRecordMetadata future = Utils.notNull(batch.tryAppend(timestamp, key, value, headers, callback, time.milliseconds()));

			dq.addLast(batch); // 加入双端队列

			return new RecordAppendResult(future, dq.size() > 1 || batch.isFull(), true);
		}
	}
}