RocketMQ源码17-consumer 负载均衡（重平衡）

本系列RocketMQ4.8注释github地址，希望对大家有所帮助，要是觉得可以的话麻烦给点一下Star哈

前面我们在分析Consumer消费过程时，有提到一个非常重要的概念，就是重平衡，只有在经过重平衡后，消息的拉取对象PullMessageService才可以去Broker拉取消息，那么这篇文章就单独分析下什么是重平衡?

重平衡分析的前提是：集群消费模式。重平衡要做的也很简单，就是给当前消费者分配属于它的逻辑消费队列。

1.什么是重平衡?

Rebalance(重平衡)机制指的是：将一个Topic下的多个队列，在同一个消费者组(consumer group)下的多个消费者实例(consumer instance)之间进行重新分配。

Rebalance机制的本意是为了提升消息的并行消费能力。例如，⼀个Topic下5个队列，在只有1个消费者的情况下，这个消费者将负责消费这5个队列的消息。如果此时我们增加⼀个消费者，那么就可以给其中⼀个消费者分配2个队列，给另⼀个分配3个队列，从而提升消息的并行消费能力。

由于⼀个队列最多分配给消费者组下的⼀个消费者，因此当某个消费者组下的消费者实例数量大于队列的数量时，多余的消费者实例将分配不到任何队列。

接下来以集群模式下的消息推模式DefaultMQPushConsumerImpl为例，看一下负载均衡的过程。

2. 消费者启动DefaultMQPushConsumerImpl

首先，消费者在启动时会做如下操作：

从NameServer更新当前消费者订阅主题的路由信息；
向Broker发送心跳，注册消费者；
唤醒负载均衡服务，触发一次负载均衡；

public class DefaultMQPushConsumerImpl implements MQConsumerInner {
    public synchronized void start() throws MQClientException {
        // ...
        // 更新当前消费者订阅主题的路由信息
        this.updateTopicSubscribeInfoWhenSubscriptionChanged();
        this.mQClientFactory.checkClientInBroker();
        // 向Broker发送心跳
        this.mQClientFactory.sendHeartbeatToAllBrokerWithLock();
        // 唤醒负载均衡服务
        this.mQClientFactory.rebalanceImmediately();
    }
}
复制代码

2.1 更新主题路由信息

为了保证消费者拿到的主题路由信息是最新的（topic下有几个消息队列、消息队列的分布信息等），在进行负载均衡之前首先要更新主题的路由信息，在updateTopicSubscribeInfoWhenSubscriptionChanged方法中可以看到，首先获取了当前消费者订阅的所有主题信息（一个消费者可以订阅多个主题），然后进行遍历，向NameServer发送请求，更新每一个主题的路由信息，保证路由信息是最新的：

public class DefaultMQPushConsumerImpl implements MQConsumerInner {
    private void updateTopicSubscribeInfoWhenSubscriptionChanged() {
        // 获取当前消费者订阅的主题信息
        Map<String, SubscriptionData> subTable = this.getSubscriptionInner();
        if (subTable != null) {
            // 遍历订阅的主题信息
            for (final Map.Entry<String, SubscriptionData> entry : subTable.entrySet()) {
                final String topic = entry.getKey();
                // 从NameServer更新主题的路由信息
                this.mQClientFactory.updateTopicRouteInfoFromNameServer(topic);
            }
        }
    }
}
复制代码

2.2 注册消费者

2.2.1 发送心跳

由于Broker需要感知消费者数量的增减，所以每个消费者在启动的时候，会调用sendHeartbeatToAllBrokerWithLock向Broker发送心跳包，进行消费者注册：

public class MQClientInstance {
    public void sendHeartbeatToAllBrokerWithLock() {
        if (this.lockHeartbeat.tryLock()) {
            try {
                // todo 调用sendHeartbeatToAllBroker向Broker发送心跳
                this.sendHeartbeatToAllBroker();
                this.uploadFilterClassSource();
            } catch (final Exception e) {
                log.error("sendHeartbeatToAllBroker exception", e);
            } finally {
                this.lockHeartbeat.unlock();
            }
        } else {
            log.warn("lock heartBeat, but failed. [{}]", this.clientId);
        }
    }
}
复制代码

在sendHeartbeatToAllBroker方法中，可以看到从brokerAddrTable中获取了所有的Broker进行遍历（主从模式下也会向从节点发送请求注册），调用MQClientAPIImpl的sendHearbeat方法向每一个Broker发送心跳请求进行注册：

public class MQClientInstance {
    // Broker路由表
    private final ConcurrentMap<String/* Broker Name */, HashMap<Long/* brokerId */, String/* address */>> brokerAddrTable =
        new ConcurrentHashMap<String, HashMap<Long, String>>();
    // 发送心跳
    private void sendHeartbeatToAllBroker() {
        final HeartbeatData heartbeatData = this.prepareHeartbeatData();
        // ...
        if (!this.brokerAddrTable.isEmpty()) {
            long times = this.sendHeartbeatTimesTotal.getAndIncrement();
            // 获取所有的Broker进行遍历， key为 Broker Name， value为同一个name下的所有Broker实例（主从模式下Broker的name一致）
            Iterator<Entry<String, HashMap<Long, String>>> it = this.brokerAddrTable.entrySet().iterator();
            while (it.hasNext()) {
                Entry<String, HashMap<Long, String>> entry = it.next();
                String brokerName = entry.getKey(); // broker name
                // 获取同一个Broker Name下的所有Broker实例
                HashMap<Long, String> oneTable = entry.getValue();
                if (oneTable != null) {
                    // 遍历所有的实例
                    for (Map.Entry<Long, String> entry1 : oneTable.entrySet()) {
                        Long id = entry1.getKey();
                        String addr = entry1.getValue();
                        if (addr != null) { // 如果地址不为空
                            // ...
                            try {
                                // 发送心跳
                                int version = this.mQClientAPIImpl.sendHearbeat(addr, heartbeatData, clientConfig.getMqClientApiTimeout());
                                // ...
                            } catch (Exception e) {
                                // ...
                            }
                        }
                    }
                }
            }
        }
    }
}
复制代码

在MQClientAPIImpl的sendHearbeat方法中，可以看到构建了HEART_BEAT请求，然后向Broker发送：

public class MQClientAPIImpl {
   public int sendHearbeat(final String addr, final HeartbeatData heartbeatData, final long timeoutMillis
    ) throws RemotingException, MQBrokerException, InterruptedException {
        // 创建HEART_BEAT请求
        RemotingCommand request = RemotingCommand.createRequestCommand(RequestCode.HEART_BEAT, null);
        request.setLanguage(clientConfig.getLanguage());
        request.setBody(heartbeatData.encode());
        // 发送请求
        RemotingCommand response = this.remotingClient.invokeSync(addr, request, timeoutMillis);
        // ...
    }
}
复制代码

2.2.2 broker心跳请求处理

Broker在启动时注册了HEART_BEAT请求的处理器，可以看到请求处理器是ClientManageProcessor：

public class BrokerController {
    public void registerProcessor() {
        ClientManageProcessor clientProcessor = new ClientManageProcessor(this);
        // 注册HEART_BEAT请求的处理器ClientManageProcessor
        this.remotingServer.registerProcessor(RequestCode.HEART_BEAT, clientProcessor, this.heartbeatExecutor);
    }
}
复制代码

进入到ClientManageProcessor的processRequest方法，如果请求是HEART_BEAT类型会调用heartBeat方法进行处理，这里也能看还有UNREGISTER_CLIENT类型的请求，从名字上可以看出是与取消注册有关的（这个稍后再说）：

public class ClientManageProcessor extends AsyncNettyRequestProcessor implements NettyRequestProcessor {
    @Override
    public RemotingCommand processRequest(ChannelHandlerContext ctx, RemotingCommand request)
        throws RemotingCommandException {
        switch (request.getCode()) {
            case RequestCode.HEART_BEAT: // 处理心跳请求
                return this.heartBeat(ctx, request);
            case RequestCode.UNREGISTER_CLIENT: // 取消注册请求
                return this.unregisterClient(ctx, request);
            case RequestCode.CHECK_CLIENT_CONFIG:
                return this.checkClientConfig(ctx, request);
            default:
                break;
        }
        return null;
    }
}
复制代码

进入到heartBeat方法，可以看到，调用了ConsumerManager的registerConsumer注册消费者：

public class ClientManageProcessor extends AsyncNettyRequestProcessor implements NettyRequestProcessor {
    public RemotingCommand heartBeat(ChannelHandlerContext ctx, RemotingCommand request) {
        // ...
        for (ConsumerData data : heartbeatData.getConsumerDataSet()) {
            // ...
            // 注册Consumer
            boolean changed = this.brokerController.getConsumerManager().registerConsumer(
                data.getGroupName(), clientChannelInfo, data.getConsumeType(), data.getMessageModel(), data.getConsumeFromWhere(),
                data.getSubscriptionDataSet(), isNotifyConsumerIdsChangedEnable);
            // ...
        }
        // ...
        return response;
    }
}
复制代码

进行注册

ConsumerManager的registerConsumer方法的理逻辑如下：

根据组名称获取该消费者组的信息ConsumerGroupInfo对象。如果获取为空，会创建一个ConsumerGroupInfo，记录了消费者组的相关信息；
判断消费者是否发生了变更，如果如果发生了变化，会触发CHANGE变更事件（这个稍后再看）；
触发REGISTER注册事件；

public class ConsumerManager {
    public boolean registerConsumer(final String group, final ClientChannelInfo clientChannelInfo,
        ConsumeType consumeType, MessageModel messageModel, ConsumeFromWhere consumeFromWhere,
        final Set<SubscriptionData> subList, boolean isNotifyConsumerIdsChangedEnable) {
        // 根据组名称获取消费者组信息
        ConsumerGroupInfo consumerGroupInfo = this.consumerTable.get(group);
        if (null == consumerGroupInfo) { // 如果为空新增ConsumerGroupInfo对象
            ConsumerGroupInfo tmp = new ConsumerGroupInfo(group, consumeType, messageModel, consumeFromWhere);
            ConsumerGroupInfo prev = this.consumerTable.putIfAbsent(group, tmp);
            consumerGroupInfo = prev != null ? prev : tmp;
        }
        boolean r1 =
            consumerGroupInfo.updateChannel(clientChannelInfo, consumeType, messageModel,
                consumeFromWhere);
        boolean r2 = consumerGroupInfo.updateSubscription(subList);
        // 如果有变更
        if (r1 || r2) {
            if (isNotifyConsumerIdsChangedEnable) {
                // 通知变更
                this.consumerIdsChangeListener.handle(ConsumerGroupEvent.CHANGE, group, consumerGroupInfo.getAllChannel());
            }
        }
        // 注册Consumer
        this.consumerIdsChangeListener.handle(ConsumerGroupEvent.REGISTER, group, subList);
        return r1 || r2;
    }
}
复制代码

进入到DefaultConsumerIdsChangeListener的handle方法中，可以看到如果是REGISTER事件，会通过ConsumerFilterManager的register方法进行注册，注册的详细过程这里先不展开讲解：

public class DefaultConsumerIdsChangeListener implements ConsumerIdsChangeListener {
    @Override
    public void handle(ConsumerGroupEvent event, String group, Object... args) {
        if (event == null) {
            return;