图解 Kafka 网络层实现机制之Selector 多路复用器

public class Selector implements Selectable, AutoCloseable {
    // 在 Java NIO 中用来监听网络I/O事件
    private final java.nio.channels.Selector nioSelector;
    // channels 管理
    private final Map<String, KafkaChannel> channels; 
    // 发送完成的Send集合
    private final List<Send> completedSends;
    // 已经接收完毕的请求集合
    private final LinkedHashMap<String, NetworkReceive> completedReceives;
    // 立即连接的集合
    private final Set<SelectionKey> immediatelyConnectedKeys;
    // 关闭连接的 channel 集合
    private final Map<String, KafkaChannel> closingChannels;
    // 断开连接的节点集合
    private final Map<String, ChannelState> disconnected;
    // 连接成功的节点集合
    private final List<String> connected;
    // 发送失败的请求集合
    private final List<String> failedSends;
    // 用来构建 KafkaChannel 的工具类
    private final ChannelBuilder channelBuilder;
    // 最大可以接收的数据量大小
    private final int maxReceiveSize;
    // 空闲超时到期连接管理器
    private final IdleExpiryManager idleExpiryManager;
    // 用来管理 ByteBuffer 的内存池
    private final MemoryPool memoryPool;
    // 初始化 Selector
    public Selector(int maxReceiveSize,
            long connectionMaxIdleMs,
            int failedAuthenticationDelayMs,
            Metrics metrics,
            Time time,
            String metricGrpPrefix,
            Map<String, String> metricTags,
            boolean metricsPerConnection,
            boolean recordTimePerConnection,
            ChannelBuilder channelBuilder,
            MemoryPool memoryPool,
            LogContext logContext) {
        try {
            this.nioSelector = java.nio.channels.Selector.open();
        } catch (IOException e) {
            throw new KafkaException(e);
        }
        this.maxReceiveSize = maxReceiveSize;
        this.time = time;
        this.channels = new HashMap<>();
        this.explicitlyMutedChannels = new HashSet<>();
        this.outOfMemory = false;
        this.completedSends = new ArrayList<>();
        this.completedReceives = new LinkedHashMap<>();
        this.immediatelyConnectedKeys = new HashSet<>();
        this.closingChannels = new HashMap<>();
        this.keysWithBufferedRead = new HashSet<>();
        this.connected = new ArrayList<>();
        this.disconnected = new HashMap<>();
        this.failedSends = new ArrayList<>();
        this.log = logContext.logger(Selector.class);
        this.sensors = new SelectorMetrics(metrics, metricGrpPrefix, metricTags, metricsPerConnection);
        this.channelBuilder = channelBuilder;
        this.recordTimePerConnection = recordTimePerConnection;
        this.idleExpiryManager = connectionMaxIdleMs < 0 ? null : new IdleExpiryManager(time, connectionMaxIdleMs);
        this.memoryPool = memoryPool;
        this.lowMemThreshold = (long) (0.1 * this.memoryPool.size());
        this.failedAuthenticationDelayMs = failedAuthenticationDelayMs;
        this.delayedClosingChannels = (failedAuthenticationDelayMs > NO_FAILED_AUTHENTICATION_DELAY) ? new LinkedHashMap<String, DelayedAuthenticationFailureClose>() : null;
    }
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.

@Override
public void connect(String id, InetSocketAddress address, int sendBufferSize, int receiveBufferSize) throws IOException {
    // 1.先确认是否已经被连接过  
    ensureNotRegistered(id);
    // 2.打开一个 SocketChannel
    SocketChannel socketChannel = SocketChannel.open();
    SelectionKey key = null;
    try {
        // 3.设置 socketChannel 信息 
        configureSocketChannel(socketChannel, sendBufferSize, receiveBufferSize);
        // 4.尝试发起连接
        boolean connected = doConnect(socketChannel, address);
        // 5. 将该 socketChannel 注册到 nioSelector 上，并关注 OP_CONNECT 事件
        key = registerChannel(id, socketChannel, SelectionKey.OP_CONNECT);
        // 6.如果立即连接成功了
        if (connected) {
            ...
            // 先将 key 放入 immediatelyConnectedKeys 集合
            immediatelyConnectedKeys.add(key);
            // 并取消对 OP_CONNECT 的监听
            key.interestOps(0);
        }
    } catch (IOException | RuntimeException e) {
        if (key != null)
            immediatelyConnectedKeys.remove(key);
        channels.remove(id);
        socketChannel.close();
        throw e;
    }
}

// 设置 socketChannel 信息 
private void configureSocketChannel(SocketChannel socketChannel, int sendBufferSize, int receiveBufferSize) throws IOException {
    // 1. 设置非阻塞模式
    socketChannel.configureBlocking(false);
    // 2. 创建一个新的 Socket
    Socket socket = socketChannel.socket();
    // 3. 开启长连接 keepalive 探活机制
    socket.setKeepAlive(true);
    // 4. 设置 SocketOptions.SO_SNDBUF，默认12kb
    if (sendBufferSize != Selectable.USE_DEFAULT_BUFFER_SIZE)
        socket.setSendBufferSize(sendBufferSize);
    // 5. 设置 SocketOptions.SO_RCVBUF，默认32kb
    if (receiveBufferSize != Selectable.USE_DEFAULT_BUFFER_SIZE)
        socket.setReceiveBufferSize(receiveBufferSize);
    // 6. 设置 TcpNoDelay 算法，默认为 false 即开启 Nagle 算法，true 为关闭 Nagle 算法
    socket.setTcpNoDelay(true);
}
// 发起连接
protected boolean doConnect(SocketChannel channel, InetSocketAddress address) throws IOException {
    try {
        // 调用socketChannel的connect方法进行发起连接，该方法会向远端发起tcp请求
        // 因为是非阻塞的，返回时，连接不一定已经建立好（即完成3次握手）。连接如果已经建立好则返回true，否则返回false。
        //一般来说server和client在一台机器上，该方法可能返回true。在后面会通过 KSelector.finishConnect() 方法确认连接是否真正建立了。
        return channel.connect(address);
    } catch (UnresolvedAddressException e) {
        throw new IOException("Can't resolve address: " + address, e);
    }
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.

// 将该 socketChannel 注册到 nioSelector 上，并关注感兴趣的事件
protected SelectionKey registerChannel(String id, SocketChannel socketChannel, int interestedOps) throws IOException {
        // 1. 将该 socketChannel 注册到 nioSelector 上，并设置读事件监听
        SelectionKey key = socketChannel.register(nioSelector, interestedOps);
        // 2. 构建 KafkaChannel，将 key 与 KafkaChannel 做注册绑定
        KafkaChannel channel = buildAndAttachKafkaChannel(socketChannel, id, key);
        // 3. 将nodeid，channel 绑定并放入到 channels 集合中
        this.channels.put(id, channel);
        if (idleExpiryManager != null)
            // 4. 更新连接到空闲超时到期连接管理器中，并记录活跃时间
            idleExpiryManager.update(channel.id(), time.nanoseconds());
        return key;
}

// 构建 KafkaChannel 并关联 key 和 Channel，方便查找
private KafkaChannel buildAndAttachKafkaChannel(SocketChannel socketChannel, String id, SelectionKey key) throws IOException {
    try {
        // 1. 构建 KafkaChannel
        KafkaChannel channel = channelBuilder.buildChannel(id, key, maxReceiveSize, memoryPool,
            new SelectorChannelMetadataRegistry());
        // 2. 将 KafkaChannel 注册到 key 上，并做关联，方便查找
        key.attach(channel);
        // 3. 返回建立好的 KafkaChannel
        return channel;
    } catch (Exception e) {
        try {
            socketChannel.close();
        } finally {
            key.cancel();
        }
        throw new IOException("Channel could not be created for socket " + socketChannel, e);
    }
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.

/**
 * 消息预发送
 */
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);
            ...
        }
    }
}

// 判断 channelid 是否存在
private KafkaChannel openOrClosingChannelOrFail(String id) {
    // 通过 channelid 先从 channels 集合中获取
    KafkaChannel channel = this.channels.get(id);
    // 如果为空那么再从 closingChannels 集合中获取
    if (channel == null)
        channel = this.closingChannels.get(id);
    // 如果还为空则抛异常    
    if (channel == null)
        throw new IllegalStateException("Attempt to retrieve channel for which there is no connection. Connection id " + id + " existing connections " + channels.keySet());
    return channel;
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.

@Override
public void poll(long timeout) throws IOException {
    ...
    // 1. 先将上次的结果清理掉
    clear();
    boolean dataInBuffers = !keysWithBufferedRead.isEmpty();
    ...
    /* check ready keys */
    long startSelect = time.nanoseconds();
    // 2. 调用nioSelector.select线程阻塞等待I/O事件并设置阻塞时间，等待I/O事件就绪发生，然后返回已经监控到了多少准备就绪的事件
    int numReadyKeys = select(timeout);
    long endSelect = time.nanoseconds();
    // 记录耗时
    this.sensors.selectTime.record(endSelect - startSelect, time.milliseconds());
    // 3. 监听到事件发生或立即连接集合不为空或存在缓存数据
    if (numReadyKeys > 0 || !immediatelyConnectedKeys.isEmpty() || dataInBuffers) {
        // 4. 获取监听到的准备就绪事件集合 
        Set<SelectionKey> readyKeys = this.nioSelector.selectedKeys();
        // 在SSL连接才可能会存在缓存数据
        if (dataInBuffers) {
            // 清除所有的就绪事件
            keysWithBufferedRead.removeAll(readyKeys); //so no channel gets polled twice
            Set<SelectionKey> toPoll = keysWithBufferedRead;
            // 重新初始化
            keysWithBufferedRead = new HashSet<>(); //poll() calls will repopulate if needed
            // 处理事件
            pollSelectionKeys(toPoll, false, endSelect);
        }
        // 5. 处理监听到的准备就绪事件
        pollSelectionKeys(readyKeys, false, endSelect);
        // 6. 就绪事件集合清理
        // Clear all selected keys so that they are included in the ready count for the next select
        readyKeys.clear();
        // 7. 处理立即连接集合
        pollSelectionKeys(immediatelyConnectedKeys, true, endSelect);
        // 8. 立即连接集合清理
        immediatelyConnectedKeys.clear();
        ...
        maybeCloseOldestConnection(endSelect);
    } else {
        ...
    }
    ...
}
// 调用nioselector.select进行阻塞监听就绪事件
private int select(long timeoutMs) throws IOException {
  if (timeoutMs < 0L)
      throw new IllegalArgumentException("timeout should be >= 0");
  if (timeoutMs == 0L)
      return this.nioSelector.selectNow();
  else
      return this.nioSelector.select(timeoutMs);
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.

private void attemptWrite(SelectionKey key, KafkaChannel channel, long nowNanos) throws IOException {
    // 
    if (channel.hasSend()
            && channel.ready()
            && key.isWritable()
            && !channel.maybeBeginClientReauthentication(() -> nowNanos)) {
        // 进行写操作
        write(channel);
    }
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.

// 执行写操作 
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);
        }
    }
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.

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;
}
/**
 * adds a receive to completed receives
 */
private void addToCompletedReceives(KafkaChannel channel, NetworkReceive networkReceive, long currentTimeMs) {
    if (hasCompletedReceive(channel))
        throw new IllegalStateException("Attempting to add second completed receive to channel " + channel.id());
    // 将 networkReceive 添加到已经接收完毕网络请求集合里 
    this.completedReceives.put(channel.id(), networkReceive);
    ...
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.

@Override
public List<Send> completedSends() {
    return this.completedSends;
}
1.
2.
3.
4.

@Override
public Collection<NetworkReceive> completedReceives() {
    return this.completedReceives.values();
}
1.
2.
3.
4.

@Override
public Map<String, ChannelState> disconnected() {
    return this.disconnected;
}
1.
2.
3.
4.

@Override
public List<String> connected() {
    return this.connected;
}
1.
2.
3.
4.

/**
 * check if channel is ready
 */
@Override
public boolean isChannelReady(String id) {
    // 从 Channels 集合中获取该id对应的 channel 
    KafkaChannel channel = this.channels.get(id);
    // 然后 channel 不为空 则判断是否准备好
    return channel != null && channel.ready();
}
// KafkaChannel 类方法
public boolean ready() {
    // 判断传输层是否准备好，默认是 PlaintextTransportLayer
    return transportLayer.ready() && authenticator.complete();
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.

/**
 * adds a receive to completed receives
 */
private void addToCompletedReceives(KafkaChannel channel, NetworkReceive networkReceive, long currentTimeMs) {
    if (hasCompletedReceive(channel))
        throw new IllegalStateException("Attempting to add second completed receive to channel " + channel.id());
    // 将 channel id 添加到已经接收完毕的网络请求集合中
    this.completedReceives.put(channel.id(), networkReceive);
    sensors.recordCompletedReceive(channel.id(), networkReceive.size(), currentTimeMs);
}
/**
 * Check if given channel has a completed receive
 */
private boolean hasCompletedReceive(KafkaChannel channel) {
    // 判断已经接收完毕的网络集合中是否存在该 channel id
    return completedReceives.containsKey(channel.id());
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.

private static class IdleExpiryManager {
    // lru 连接集合
    private final Map<String, Long> lruConnections;
    // 连接最大的空闲时间 默认9分钟
    private final long connectionsMaxIdleNanos;
    private long nextIdleCloseCheckTime;
    // 初始化管理器
    public IdleExpiryManager(Time time, long connectionsMaxIdleMs) {
        this.connectionsMaxIdleNanos = connectionsMaxIdleMs * 1000 * 1000;
        // initial capacity and load factor are default, we set them explicitly because we want to set accessOrder = true
        // 初始化lru连接集合，设置初始容量，扩容因子，是否排序
        this.lruConnections = new LinkedHashMap<>(16, .75F, true);
        this.nextIdleCloseCheckTime = time.nanoseconds() + this.connectionsMaxIdleNanos;
    } 
    // 更新活跃时间
    public void update(String connectionId, long currentTimeNanos) {
        lruConnections.put(connectionId, currentTimeNanos);
    }
    ...
    // 删除连接
    public void remove(String connectionId) {
        lruConnections.remove(connectionId);
    }
}
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.