图解 Kafka 源码实现机制之客户端缓存架构

import java.nio.ByteBuffer;
public class MemoryRecords extends AbstractRecords {
  public static MemoryRecordsBuilder builder(..){
        // 重载builder 
        return builder(...);
  }
    
  public static MemoryRecordsBuilder builder(
    ByteBuffer buffer,
    // 消息版本
    byte magic,
    // 消息压缩类型
    CompressionType compressionType,
    // 时间戳
    TimestampType timestampType,
    // 基本位移
    long baseOffset,
    // 日志追加时间
    long logAppendTime,
    // 生产者id
    long producerId,
    // 生产者版本
    short producerEpoch,
    // 批次序列号
    int baseSequence,
    boolean isTransactional,
    // 是否是控制类的批次
    boolean isControlBatch,
    // 分区leader的版本
    int partitionLeaderEpoch){
        // 初始化MemoryRecordsBuilder类
        return new MemoryRecordsBuilder(...);
  }
}

public class MemoryRecordsBuilder implements AutoCloseable {
    // 写操作关闭的输出流
    private static final DataOutputStream CLOSED_STREAM = new DataOutputStream(new OutputStream() {
        // 当向某个ByteBuffer关闭输出流写数据时抛异常
        public void write(int b){
            throw new ...;
        }
    });
    // 日志时间
    private final TimestampType timestampType;
    // 消息压缩类型
    private final CompressionType compressionType;
    // kafka对OutputStream接口的实现类，对ByteBuffer实现了自动扩容功能
    private final ByteBufferOutputStream bufferStream;
    // 消息的版本
    private final byte magic;
    // ByteBuffer的最初始位置
    private final int initialPosition;
    // 基本位移
    private final long baseOffset;
    // 消息追加的时间
    private final long logAppendTime;
    // 是否是控制类的批次
    private final boolean isControlBatch;
    // 分区leader的版本
    private final int partitionLeaderEpoch;
    // 写入上限
    private final int writeLimit;
    // batch头大小字节数
    private final int batchHeaderSizeInBytes;
    // 评估压缩率
    private float estimatedCompressionRatio = 1.0F;
    // 对bufferStream添加压缩功能
    private DataOutputStream appendStream;
    // 是否是事务批次
    private boolean isTransactional;
    // 生产者id
    private long producerId;
    // 生产者版本
    private short producerEpoch;
    // 批次序列号
    private int baseSequence;
    // 压缩前要写入的消息体大小字节数
    private int uncompressedRecordsSizeInBytes = 0; 
    // 压缩前写入的记录数（不包括头）
    private int numRecords = 0;
    // 实际压缩率
    private float actualCompressionRatio = 1;
    // 最大时间戳
    private long maxTimestamp = RecordBatch.NO_TIMESTAMP;
    // 最大时间戳偏移量
    private long offsetOfMaxTimestamp = -1;
    // 最后的偏移量
    private Long lastOffset = null;
    // 第一次追加消息的时间戳
    private Long firstTimestamp = null;
    // 真正保存消息的地方
    private MemoryRecords builtRecords;

public MemoryRecordsBuilder(ByteBuffer buffer,...){
    // 将MemoryRecordsBuilder关联的ByteBuffer封装成ByteBufferOutputStream流
    this(new ByteBufferOutputStream(buffer), ...);
}

// 构造方法
public MemoryRecordsBuilder(
    ByteBufferOutputStream bufferStream,
    ...
    int writeLimit){
        ....
        // 初始位置
        this.initialPosition = bufferStream.position();
        // 1. 根据不同消息版本计算批次Batch头的长度
        this.batchHeaderSizeInBytes = AbstractRecords.recordBatchHeaderSizeInBytes(magic, compressionType);
        // 2. 调整对应的position
        bufferStream.position(initialPosition + batchHeaderSizeInBytes);
        this.bufferStream = bufferStream;
        // 3. 在bufferStream流外层套一层压缩流，再套一层DataOutputStream流
        this.appendStream = new DataOutputStream(compressionType.wrapForOutput(this.bufferStream, magic));
    }
}

// 追加新记录
public Long append(long timestamp, ByteBuffer key, ByteBuffer value, Header[] headers){
   return appendWithOffset(nextSequentialOffset(), timestamp, key, value, headers);
}

// 计算下一个连续偏移量
private long nextSequentialOffset(){
  // lastOffset用来记录当前写入Record的offset，每次当有新Record写入时，都会递增它。
  return lastOffset == null ? baseOffset : lastOffset + 1;
}

// 根据偏移量追加消息
private Long appendWithOffset(
  long offset,
  boolean isControlRecord, 
  long timestamp, 
  ByteBuffer key,
  ByteBuffer value, 
  Header[] headers){
    try {
        // 检查isControl标志是否一致
        if (isControlRecord != isControlBatch)
            throw new ...;
        // 保证offset是递增的
        if (lastOffset != null && offset <= lastOffset)
            throw new ...;
        // 检查时间戳      
        if (timestamp < 0 && timestamp != RecordBatch.NO_TIMESTAMP)
           throw new ...;
        // 只有V2版本才有header
        if (magic < RecordBatch.MAGIC_VALUE_V2 && headers != null && headers.length > 0)
           throw new ...;
        // 更新firstTimestamp
        if (firstTimestamp == null)
           firstTimestamp = timestamp;
        // V2版本消息写入
        if (magic > RecordBatch.MAGIC_VALUE_V1)         {
            appendDefaultRecord(offset, timestamp, key, value, headers);
            return null;
        } else {
            //V0、V1 版本消息写入(此处不进行剖析)
            return appendLegacyRecord(offset, timestamp, key, value, magic);
        }
    } catch (IOException e) {
        // 抛异常
    }
}

private void appendDefaultRecord(
  long offset, 
  long timestamp, 
  ByteBuffer key, 
  ByteBuffer value,
  Header[] headers) throws IOException {
    // 1. 检查appendStream状态是否可以写
    ensureOpenForRecordAppend();
    // 2. 计算写入多少偏移量
    int offsetDelta = (int) (offset - baseOffset);
    // 3.计算本次写与第一次写之间时间差
    long timestampDelta = timestamp - firstTimestamp;
    // 4.使用DefaultRecord.writeTo()方法会按照V2 版本格式写入appendStream流中，并返回压缩前的消息大小
    int sizeInBytes = DefaultRecord.writeTo(appendStream, offsetDelta, timestampDelta, key, value, headers);
    // 5. 消息写入成功后更新RecordBatch的元信息
    recordWritten(offset, timestamp, sizeInBytes);
}

// 判断appendStream状态是否为CLOSED_STREAM 
private void ensureOpenForRecordAppend(){
    if (appendStream == CLOSED_STREAM)
        throw new ...;
}

// 消息写入成功后更新RecordBatch的元信息
private void recordWritten(long offset, long timestamp, int size){
  ....
  // 压缩前写入的记录数 + 1
  numRecords += 1;
  // 压缩前要写入的消息体大小字节数 + size
  uncompressedRecordsSizeInBytes += size;
  // 最后的偏移量 + offset
  lastOffset = offset;
  if (magic > RecordBatch.MAGIC_VALUE_V0 && timestamp > maxTimestamp) {
      // 赋值最大时间戳
      maxTimestamp = timestamp;
      // 赋值最大时间戳偏移量
      offsetOfMaxTimestamp = offset;
  }
}

public boolean hasRoomFor(long timestamp, ByteBuffer key, ByteBuffer value, Header[] headers){
    // 检查两个状态
    // (1)appendStream流状态
    // (2)当前已经写入的预估字节数是否超过了writeLimit写入上限
    if (isFull())
        return false;
    // 每个RecordBatch至少可以写入一个Record，此时如果一个Record都没有，则可以继续写入
    if (numRecords == 0)
        return true;
    final int recordSize;
    if (magic < RecordBatch.MAGIC_VALUE_V2) {
        // 预估V0、V1旧版本的Record大小
        recordSize = Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, key, value);
    } else {
        // 预估V2版本写入的Record大小
        int nextOffsetDelta = lastOffset == null ? 0 : (int) (lastOffset - baseOffset + 1);
        ...
        recordSize = DefaultRecord.sizeInBytes(nextOffsetDelta, timestampDelta, key, value, headers);
    }

    // 已写入字节数 + 本次写入Record的预估字节数不能超过writeLimit写入上限
    return this.writeLimit >= estimatedBytesWritten() + recordSize;
}

public boolean isFull(){
      return appendStream == CLOSED_STREAM || 
      (this.numRecords > 0 && this.writeLimit <= estimatedBytesWritten());
}

public class ByteBufferOutputStream extends OutputStream {
   // 扩容因子1.1倍
   private static final float REALLOCATION_FACTOR = 1.1f;
   // 初始容量
   private final int initialCapacity;
   // 初始位置
   private final int initialPosition;
   // 计算是否需要扩容
   public void ensureRemaining(int remainingBytesRequired){
     // 当写入字节数大于buffer当前剩余字节数就开启扩容
     if (remainingBytesRequired > buffer.remaining())
     expandBuffer(remainingBytesRequired);
  }
  
  // 扩容
  private void expandBuffer(int remainingRequired){
    // 1. 评估需要多少空间
    int expandSize = Math.max((int) (buffer.limit() * REALLOCATION_FACTOR), buffer.position() + remainingRequired);
    // 2. 申请新的ByteBuffer
    ByteBuffer temp = ByteBuffer.allocate(expandSize);
    // 3. 获取写入上限
    int limit = limit();
    // 4. 写状态转换为读状态
    buffer.flip();
    // 5. 将buffer读到新申请的temp里
    temp.put(buffer);
    // 6. 修改写模式的limit上限
    buffer.limit(limit);
    // 7. 更新原来的buffer的position，防止被重复消费
    buffer.position(initialPosition);
    // 8. 将引用指向新申请的ByteBuffer
    buffer = temp;
  }
}

public final class ProducerBatch {
    // 批次最终状态
    private enum FinalState { ABORTED, FAILED, SUCCEEDED }
    // 批次创建时间  
    final long createdMs;
    // 批次对应的主题分区
    final TopicPartition topicPartition;
    // 请求结果的future
    final ProduceRequestResult produceFuture;
    // 用来存储消息的callback和响应数据
    private final List<Thunk> thunks = new ArrayList<>();
    // 封装MemoryRecords对象，用来存储消息的ByteBuffer
    private final MemoryRecordsBuilder recordsBuilder;
    // batch的失败重试次数
    private final AtomicInteger attempts = new AtomicInteger(0);
    // 是否是被分裂的批次
    private final boolean isSplitBatch;
    // ProducerBatch的最终状态
    private final AtomicReference<FinalState> finalState = new AtomicReference<>(null);
    // Record个数
    int recordCount;
    // 最大Record字节数
    int maxRecordSize;
    // 最后一次失败重试发送的时间戳
    private long lastAttemptMs;
    // 最后一次向该ProducerBatch追加Record的时间戳
    private long lastAppendTime;
    // Sender子线程拉取批次的时间
    private long drainedMs;
    // 是否正在重试过，如果ProducerBatch中的数据发送失败，则会重新尝试发送
    private boolean retry;
}

// 构造函数
public ProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long createdMs, boolean isSplitBatch){
    ...
    // 请求结果的future
    this.produceFuture = new ProduceRequestResult(topicPartition);
    ...
}

public class ProduceRequestResult {
    // 通过一个count为1的CountDownLatch对象间接地实现了Future的功能。
    private final CountDownLatch latch = new CountDownLatch(1);
    private final TopicPartition topicPartition;
    // 用来记录broker端关联ProducerBatch中第一条Record分配的offset值
    // 这样每个Record的真实offset就可以根据自身在ProducerBatch的位置计算出来了(baseOffset + relativeOffset)
    private volatile Long baseOffset = null;
    
    // 构造函数
    public ProduceRequestResult(TopicPartition topicPartition){
        this.topicPartition = topicPartition;
    }
    // 当等到响应会会调该函数唤醒阻塞的主线程
    public void done(){
        if (baseOffset == null)
            throw new ...;
        this.latch.countDown();
    }
    // 调用await()方法的线程会被挂起，它会等待直到count值为0才继续执行
    public void await() throws InterruptedException {
        latch.await();
    }
    // 和await()类似，只不过等待一定的时间后count值还没变为0的话就会继续执行
    public boolean await(long timeout, TimeUnit unit) throws InterruptedException {
        return latch.await(timeout, unit);
    }
}

public FutureRecordMetadata tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, Callback callback, long now){
    // 1.检查MemoryRecordsBuilder是否还有空间写入
    if (!recordsBuilder.hasRoomFor(timestamp, key, value, headers)) {
        return null;
    } else {
        // 2.调用append()方法写入Record
        Long checksum = this.recordsBuilder.append(timestamp, key, value, headers);
        // 3. 更新最大Record字节数
        this.maxRecordSize = Math.max(this.maxRecordSize, AbstractRecords.estimateSizeInBytesUpperBound(magic(),recordsBuilder.compressionType(), key, value, headers));
        ...
        // 4.构建FutureRecordMetadata对象
        FutureRecordMetadata future = new FutureRecordMetadata(this.produceFuture, this.recordCount,timestamp, checksum,key == null ? -1 : key.length,value == null ? -1 : value.length, Time.SYSTEM);
        // 5. 将Callback和FutureRecordMetadata记录到thunks集合中
        thunks.add(new Thunk(callback, future));
        // 6. 更新Record记录数
        this.recordCount++;
        // 7. 返回FutureRecordMetadata
        return future;
    }
}

public RecordMetadata get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
    ...
    // 依赖ProduceRequestResult的CountDown来实现阻塞等待
    boolean occurred = this.result.await(timeout, unit);
    ...
    // 调用value()方法返回RecordMetadata对象
    return valueOrError();
}
    
RecordMetadata valueOrError() throws ExecutionException {
    ...
    return value();
}

RecordMetadata value(){
    ...
     // 将 partition、baseOffset、relativeOffset、时间戳（LogAppendTime | CreateTimeStamp）等信息封装成 RecordMetadata 对象返回
    return new RecordMetadata(
      result.topicPartition(), 
      ...);
}

private long timestamp(){
    return result.hasLogAppendTime() ? result.logAppendTime() : createTimestamp;
}

public boolean done(long baseOffset, long logAppendTime, RuntimeException exception){
    // 1.根据exception决定本次ProducerBatch发送的最终状态
    final FinalState tryFinalState = (exception == null) ? FinalState.SUCCEEDED : FinalState.FAILED;
    ....
    // 2.通过CAS操作更新finalState状态，只有第一次更新的时候，才会触发completeFutureAndFireCallbacks()方法
    if (this.finalState.compareAndSet(null, tryFinalState)) {
        // 3.执行回调
        completeFutureAndFireCallbacks(baseOffset, logAppendTime, exception);
        return true;
    }
    ....
    return false;
}

private void completeFutureAndFireCallbacks(long baseOffset, long logAppendTime, RuntimeException exception){
  // 1.更新ProduceRequestResult中的相关字段
  produceFuture.set(baseOffset, logAppendTime, exception);

  // 2.遍历thunks集合，触发每个Record的Callback回调
  for (Thunk thunk : thunks) {
      try {
          if (exception == null) {
           // 3.获取消息元数据
           RecordMetadata metadata = thunk.future.value();
           if (thunk.callback != null)
             //4.调用回调方法
             thunk.callback.onCompletion(metadata, null);
          } else {
              if (thunk.callback != null)
                  // 4.调用回调方法
                  thunk.callback.onCompletion(null, exception);
          }
      } 
      ....
  }
  // 4.调用底层 CountDownLatch.countDown()方法，阻塞在其上的主线程。
  produceFuture.done();
}

public class BufferPool {
  // 整个BufferPool总内存大小 默认32M
  private final long totalMemory;
  // 当前BufferPool管理的单个ByteBuffer大小，16k 
  private final int poolableSize;
  // 因为有多线程并发分配和回收ByteBuffer，用锁控制并发，保证线程安全。
  private final ReentrantLock lock;
  // 对应一个ArrayDeque<ByteBuffer> 队列，其中缓存了固定大小的 ByteBuffer 对象
  private final Deque<ByteBuffer> free;
  // 此队列记录因申请不到足够空间而阻塞的线程对应的Condition 对象
  private final Deque<Condition> waiters;
  // 非池化可用的内存即totalMemory减去free列表中的全部ByteBuffer的大小
  private long nonPooledAvailableMemory;
  // 构造函数
  public BufferPool(long memory, int poolableSize, Metrics metrics, Time time, String metricGrpName){
    ...
    // 总的内存
    this.totalMemory = memory;
    // 默认的池外内存，就是总的内存
    this.nonPooledAvailableMemory = memory;
  }
}

// 分配指定空间的缓存，如果缓冲区中没有足够的空闲空间，那么会阻塞线程，直到超时或得到足够空间
public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedException {
  // 1.判断申请的内存是否大于总内存
  if (size > this.totalMemory)
      throw new IllegalArgumentException("Attempt to allocate " + size + " bytes, but there is a hard limit of "+ this.totalMemory + " on memory allocations.");
  // 初始化buffer
  ByteBuffer buffer = null;
  // 2.加锁，保证线程安全。
  this.lock.lock();
  // 如果当前BufferPool处于关闭状态，则直接抛出异常
  if (this.closed) {
      this.lock.unlock();
      throw new KafkaException("Producer closed while allocating memory");
  }
  ....
  try {
      // 3.申请内存大小恰好为16k 且free缓存池不为空
      if (size == poolableSize && !this.free.isEmpty())
      // 从free队列取出一个ByteBuffer
      return this.free.pollFirst();
      
      // 对于申请内存大小非16k情况
      // 先计算free缓存池总空间大小，判断是否足够
      int freeListSize = freeSize() * this.poolableSize;
      // 4.当前BufferPool能够释放出申请内存大小的空间
      if (this.nonPooledAvailableMemory + freeListSize >= size) {
          // 5.如果size大于非池化可用内存大小，就循环从free缓存池里释放出来空闲Bytebuffer补充到nonPooledAvailableMemory中，直到满足size大小为止。
          freeUp(size);
          // 释放非池化可用内存大小
          this.nonPooledAvailableMemory -= size;
      } else {
          // 如果当前BufferPool不够提供申请内存大小，则需要阻塞当前线程
          // 累计已经释放的内存
          int accumulated = 0;
          // 创建对应的Condition，阻塞自己等待别的线程释放内存
          Condition moreMemory = this.lock.newCondition();
          try {
              // 计算当前线程最大阻塞时长
              long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs);
              // 把自己添加到等待队列中末尾，保持公平性，先来的先获取内存，防止饥饿
              this.waiters.addLast(moreMemory);
              // 循环等待直到分配成功或超时
              while (accumulated < size) {
                  ....
                  try {
                    // 当前线程阻塞等待，返回结果为false则表示阻塞超时
                   waitingTimeElapsed = !moreMemory.await(remainingTimeToBlockNs, TimeUnit.NANOSECONDS);
                  } finally {
                      ....
                  }
                  ....   
                  // 申请内存大小是16k，且free缓存池有了空闲的ByteBuffer
                  if (accumulated == 0 && size == this.poolableSize && !this.free.isEmpty()) {
                    // 从free队列取出一个ByteBuffer
                    buffer = this.free.pollFirst();
                    // 计算累加器
                    accumulated = size;
                  } else {
                      // 释放空间给非池化可用内存，并继续等待空闲空间，如果分配多了只取够size的空间
                      freeUp(size - accumulated);
                      int got = (int) Math.min(size - accumulated, this.nonPooledAvailableMemory);
                      // 释放非池化可用内存大小
                      this.nonPooledAvailableMemory -= got;
                      // 累计分配了多少空间
                      accumulated += got;
                  }
              }
              accumulated = 0;
          } finally {
              // 如果循环有异常，将已释放的空间归还给非池化可用内存
              this.nonPooledAvailableMemory += accumulated;
              //把自己从等待队列中移除并结束
              this.waiters.remove(moreMemory);
          }
      }
  } finally {
     // 当非池化可用内存有内存或free缓存池有空闲ByteBufer且等待队列里有线程正在等待
      try {
          if (!(this.nonPooledAvailableMemory == 0 && this.free.isEmpty()) && !this.waiters.isEmpty())
           // 唤醒队列里正在等待的线程
           this.waiters.peekFirst().signal();
      } finally {
          // 解锁
          lock.unlock();
      }
  }
  // 说明空间足够，并且有足够空闲的了。可以执行真正的分配空间了。
  if (buffer == null)
      // 没有正好的buffer，从缓冲区外(JVM Heap)中直接分配内存
      return safeAllocateByteBuffer(size);
  else
      // 直接复用free缓存池的ByteBuffer
      return buffer;
}

private ByteBuffer safeAllocateByteBuffer(int size){
  boolean error = true;
  try {
      //分配空间
      ByteBuffer buffer = allocateByteBuffer(size);
      error = false;
      //返回buffer
      return buffer;
  } finally {
    if (error) {
        //分配失败了, 加锁，操作内存pool
        this.lock.lock();
        try {
            //归还空间给非池化可用内存
            this.nonPooledAvailableMemory += size;
            if (!this.waiters.isEmpty())
                //有其他在等待的线程的话，唤醒其他线程
                this.waiters.peekFirst().signal();
        } finally {
            // 加锁不忘解锁
            this.lock.unlock();
        }
    }
  }
}

protected ByteBuffer allocateByteBuffer(int size){
    // 从JVM Heap中分配空间
    return ByteBuffer.allocate(size);
}

// 不断从free队列中释放空闲的ByteBuffer来补充非池化可用内存
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){
    // 1.加锁，保证线程安全。
    lock.lock();
    try {
    // 2.如果待释放的size大小为16k，则直接放入free队列中
        if (size == this.poolableSize && size == buffer.capacity()) {
            // 清空buffer
            buffer.clear();
            // 释放buffer到free队列里
            this.free.add(buffer);
        } else {
            //如果非16k，则由JVM GC来回收ByteBuffer并增加非池化可用内存
            this.nonPooledAvailableMemory += size;
        }
        // 3.唤醒waiters中的第一个阻塞线程
        Condition moreMem = this.waiters.peekFirst();
        if (moreMem != null)
            moreMem.signal();
    } finally {
        // 释放锁
        lock.unlock();
    }
}  

  public class CopyOnWriteMap<K, V> implements ConcurrentMap<K, V> {
    // volatile Map
    private volatile Map<K, V> map;
    // 构造函数
    public CopyOnWriteMap(){
        this.map = Collections.emptyMap();
    } 

// 获取集合中队列
public V get(Object k){
    return map.get(k);
}

public synchronized V putIfAbsent(K k, V v){
    if (!containsKey(k))
        return put(k, v);
    else
        return get(k);
}

// 判断队列是否存在
public boolean containsKey(Object k){
    return map.containsKey(k);
}   

public synchronized V put(K k, V v){
    Map<K, V> copy = new HashMap<K, V>(this.map);
    V prev = copy.put(k, v);
    this.map = Collections.unmodifiableMap(copy);
    return prev;
}