图解 Kafka 网络层源码实现机制之收发消息全过程

/**
 * 消息预发送
 */
public void send(Send send) {
    // 1. 从服务端获取 connectionId
    String connectionId = send.destination();
    // 2. 从数据包中获取对应连接
    KafkaChannel channel = openOrClosingChannelOrFail(connectionId);
    // 3. 如果关闭连接集合中存在该连接
    if (closingChannels.containsKey(connectionId)) {
        // 把 connectionId 放入 failedSends 集合里
        this.failedSends.add(connectionId);
    } else {
        try {
            // 4. 暂存数据预发送，并没有真正的发送,一次只能发送一个
            channel.setSend(send);
        } catch (Exception e) {
            // 5. 更新 KafkaChannel 的状态为发送失败  
            channel.state(ChannelState.FAILED_SEND);
            // 6. 把 connectionId 放入 failedSends 集合里
            this.failedSends.add(connectionId);
            // 7. 关闭连接
            close(channel, CloseMode.DISCARD_NO_NOTIFY);
            ...
        }
    }
}

public void setSend(Send send) {
  if (this.send != null)
      throw new IllegalStateException("Attempt to begin a send operation with prior send operation still in progress, connection id is " + id);
  // 设置要发送消息的字段
  this.send = send;
  // 调用传输层增加写事件
  this.transportLayer.addInterestOps(SelectionKey.OP_WRITE);
}

// PlaintextTransportLayer 类方法
@Override
public void addInterestOps(int ops) {
    //通过 key.interestOps() | ops 来添加事件
    key.interestOps(key.interestOps() | ops);
}

@Override
public void poll(long timeout) throws IOException {
    ...
    // 调用nioSelector.select线程阻塞等待I/O事件并设置阻塞时间，等待I/O事件就绪发生，然后返回已经监控到了多少准备就绪的事件
    int numReadyKeys = select(timeout);
    // 监听到事件发生或立即连接集合不为空或存在缓存数据
    if (numReadyKeys > 0 || !immediatelyConnectedKeys.isEmpty() || dataInBuffers) {
        // 在SSL连接才可能会存在缓存数据
        if (dataInBuffers) {
            // 处理事件
            pollSelectionKeys(toPoll, false, endSelect);
        }
        // 处理监听到的准备就绪事件
        pollSelectionKeys(readyKeys, false, endSelect);
        // 处理立即连接集合
        pollSelectionKeys(immediatelyConnectedKeys, true, endSelect);
    } else {
        ...
    }
    ...
}

void pollSelectionKeys(Set<SelectionKey> selectionKeys,boolean isImmediatelyConnected,long currentTimeNanos) {
    //1. 循环调用当前监听到的事件(原顺序或者洗牌后顺序)
    for (SelectionKey key : determineHandlingOrder(selectionKeys)) {
        // 2. 之前创建连接，把kafkachanel注册到key上，这里就是获取对应的 channel
        KafkaChannel channel = channel(key);
        ...
        // 3. 获取节点id
        String nodeId = channel.id();
        ...
        try {
            ...
            // 4. 读事件是否准备就绪了
            if (channel.ready() && (key.isReadable() || channel.hasBytesBuffered()) && !hasCompletedReceive(channel) && !explicitlyMutedChannels.contains(channel)) {
                // 尝试处理读事件
                attemptRead(channel);
        }
        ...
        try {
           // 5. 尝试处理写事件
           attemptWrite(key, channel, nowNanos);
        } catch (Exception e) {
            sendFailed = true;
            throw e;
        }
    } catch (Exception e) {
        ...
    } finally {
       ....
    }
  }
}

private void attemptWrite(SelectionKey key, KafkaChannel channel, long nowNanos) throws IOException {
    // 此处需要满足4个条件才可以进行写操作      
    if (channel.hasSend()
            && channel.ready()
            && key.isWritable()
            && !channel.maybeBeginClientReauthentication(() -> nowNanos)) {
        // 进行写操作
        write(channel);
    }
}
// channel 连接就绪
public boolean ready() {
    return transportLayer.ready() && authenticator.complete();
}
// java nio SelectionKey
public final boolean isWritable() {
    return (readyOps() & OP_WRITE) != 0;
}

// 执行写操作 
void write(KafkaChannel channel) throws IOException {
    // 1.获取 channel 对应的节点id    
    String nodeId = channel.id();
    // 2. 将保存在 send 上的数据真正发送出去，但是一次不一定能发送完，会返回已经发出的字节数
    long bytesSent = channel.write();
    // 3. 判断是否发送完成，未完成返回null，等待下次poll继续发送
    Send send = channel.maybeCompleteSend();
    // 4. 说明已经发出或者发送完成
    if (bytesSent > 0 || send != null) {
        long currentTimeMs = time.milliseconds();
        if (bytesSent > 0)
            // 记录发送字节 Metrics 信息
            this.sensors.recordBytesSent(nodeId, bytesSent, currentTimeMs);
        // 发送完成
        if (send != null) {
            // 将 send 添加到 completedSends
            this.completedSends.add(send);
            //  记录发送完成 Metrics 信息
            this.sensors.recordCompletedSend(nodeId, send.size(), currentTimeMs);
        }
    }
}

public long write() throws IOException {
    // 判断 send 是否为空，如果为空表示已经发送完毕了
    if (send == null)
        return 0;

    midWrite = true;
    // 调用ByteBufferSend.writeTo把数据真正发送出去
    return send.writeTo(transportLayer);
}

@Override
// 将字节流数据写入到channel中
public long writeTo(GatheringByteChannel channel) throws IOException {
    // 1.调用nio底层write方法把buffers写入传输层返回写入的字节数
    long written = channel.write(buffers);
    if (written < 0)
        throw new EOFException("Wrote negative bytes to channel. This shouldn't happen.");
    // 2.计算还剩多少字节没有写入传输层
    remaining -= written;
    // 每次发送 都检查是否
    pending = TransportLayers.hasPendingWrites(channel);
    return written;
}

// 可能完成发送
public Send maybeCompleteSend() {
    // send 不为空且已经发送完毕
    if (send != null && send.completed()) {
        midWrite = false;
        // 当写数据完毕后，取消传输层对 OP_WRITE 事件的监听，完成一次写操作
        transportLayer.removeInterestOps(SelectionKey.OP_WRITE);
        // 将 send 赋值给结果集 result
        Send result = send;
        // 此时读完后将 send 清空，以便下次写
        send = null;
        // 最后返回结果集 result，完成一次写操作
        return result;
    }
    return null;
}
// PlaintextTransportLayer 类方法
@Override
public void removeInterestOps(int ops) {
    // 通过 key.interestOps() & ~ops 来删除事件
    key.interestOps(key.interestOps() & ~ops);
}
// ByteBufferSend
@Override
public boolean completed() {
    return remaining <= 0 && !pending;
}

private void handleCompletedSends(List<ClientResponse> responses, long now) {
    // if no response is expected then when the send is completed, return it
    // 上面发送完成将 send 添加到 completedSends 集合，然后遍历这个集合
    for (Send send : this.selector.completedSends())     {
        // 获取 inFlightRequests 集合发往对应 Broker 的最后一个请求元素
        InFlightRequest request = this.inFlightRequests.lastSent(send.destination());
        // 判断是否期望进行响应
        if (!request.expectResponse) {
            // 如果不期望进行响应就删除inFlightRequests集合发往对应 Broker 请求队列的第一个元素
            this.inFlightRequests.completeLastSent(send.destination());
            // 把请求添加到 responses 集合里
            responses.add(request.completed(null, now));
        }
    }
}

private void attemptRead(KafkaChannel channel) throws IOException {
    // 获取 channel 对应的节点 id
    String nodeId = channel.id();
    // 将从传输层中读取数据到NetworkReceive对象中
    long bytesReceived = channel.read();
    if (bytesReceived != 0) {
        ...
        // 判断 NetworkReceive 对象是否已经读完了
        NetworkReceive receive = channel.maybeCompleteReceive();
        // 当读完后把这个 NetworkReceive 对象添加到已经接收完毕网络请求集合里
        if (receive != null) {
            addToCompletedReceives(channel, receive, currentTimeMs);
        }
    }
    ...
}
// KafkaChannel 方法
public long read() throws IOException {
    if (receive == null) {
        // 初始化 NetworkReceive 对象
        receive = new NetworkReceive(maxReceiveSize, id, memoryPool);
    }
    // 尝试把 channel 的数据读到 NetworkReceive 对象中
    long bytesReceived = receive(this.receive);
    ...
    return bytesReceived;
}

// 判断 NetworkReceive 对象是否已经读完了
// 如果此时并没有读完一个完整的NetworkReceive对象,则下次触发读事件会继续填充整个NetworkReceive对象，
// 如果读完一个完整的NetworkReceive对象则将其置空，下次触发读事件时会创建一个全新的NetworkReceive对象。
public NetworkReceive maybeCompleteReceive() {
  if (receive != null && receive.complete()) {
      receive.payload().rewind();
      NetworkReceive result = receive;
      receive = null;
      return result;
  }
  return null;
}
// NetworkReceive
public boolean complete() {
    return !size.hasRemaining() && buffer != null && !buffer.hasRemaining();
}

private void handleCompletedReceives(List<ClientResponse> responses, long now) {
  // 当读完后把这个 NetworkReceive 对象添加到已经接收完毕网络请求集合里，然后遍历这个集合
  for (NetworkReceive receive : this.selector.completedReceives()) {
      // 获取发送请求的node id 
      String source = receive.source();
      // 从 InFlightRequest 集合取出对应的元素并删除
      InFlightRequest req = inFlightRequests.completeNext(source);
      // 解析该响应
      Struct responseStruct = parseStructMaybeUpdateThrottleTimeMetrics(receive.payload(), req.header,
          throttleTimeSensor, now);
      ....
      // 添加响应到响应结果集合中
      responses.add(req.completed(response, now));
  }
}

private static Struct parseStructMaybeUpdateThrottleTimeMetrics(ByteBuffer responseBuffer, RequestHeader requestHeader,Sensor throttleTimeSensor, long now) {
    // 获取响应头
    ResponseHeader responseHeader = ResponseHeader.parse(responseBuffer,requestHeader.apiKey().responseHeaderVersion(requestHeader.apiVersion()));
    // 获取响应体
    Struct responseBody = requestHeader.apiKey().parseResponse(requestHeader.apiVersion(), responseBuffer);
    // 对比响应头 correlationId 和响应体的 correlationId 是否一致，否则抛异常
    correlate(requestHeader, responseHeader);
    ...
    return responseBody;
}

private void completeResponses(List<ClientResponse> responses) {
    // 遍历响应结果集合
    for (ClientResponse response : responses) {
        try {
            response.onComplete();
        } catch (Exception e) {
            log.error("Uncaught error in request completion:", e);
        }
    }
}

//ClientResponse 类
public void onComplete() {
    if (callback != null)
        callback.onComplete(this);
}