RocketMQ源码分析之RebalanceService

  前面分析了消息在RocketMQ中是如何存储的，那么消息存储后消费者第一次是如何通过ConsumeQueue找到消息，这其实与RebalanceService是分不开的，本节就来回答这些问题。
   RebalanceService继承ServiceThread，其本质是一个线程，作用是每20秒就会调用一次doRebalance()方法来完成consumer负载均衡、将重新分配的MessageQueue构建PullRequest请求并将其放到PullMessageService服务中的pullRequestQueue队列中。RebalanceService是在consumer启动的过程中启动的。

public class RebalanceService extends ServiceThread {
    private static long waitInterval =
        Long.parseLong(System.getProperty(
            "rocketmq.client.rebalance.waitInterval", "20000"));
    private final InternalLogger log = ClientLogger.getLog();
    private final MQClientInstance mqClientFactory;

    public RebalanceService(MQClientInstance mqClientFactory) {
        this.mqClientFactory = mqClientFactory;
    }

    @Override
    public void run() {
        log.info(this.getServiceName() + " service started");

        while (!this.isStopped()) {
            this.waitForRunning(waitInterval);
            this.mqClientFactory.doRebalance();
        }

        log.info(this.getServiceName() + " service end");
    }

    @Override
    public String getServiceName() {
        return RebalanceService.class.getSimpleName();
    }
}

   由于RebalanceService本质上是个线程，所以当RebalanceService启动后会执行run()方法，该方法的功能是每20秒执行一次doRebalance()，这个间隔时间可以通过-Drocketmq.client.rebalance.waitInterval=XXX来修改。

public void doRebalance() {
    for (Map.Entry<String, MQConsumerInner> entry : this.consumerTable.entrySet()) {
        MQConsumerInner impl = entry.getValue();
        if (impl != null) {
            try {
                impl.doRebalance();
            } catch (Throwable e) {
                log.error("doRebalance exception", e);
            }
        }
    }
}

   doRebalance()方法实现逻辑是：遍历已经注册的消费者并对消费者执行doRebalance()方法。

public void doRebalance(final boolean isOrder) {
    Map<String, SubscriptionData> subTable = this.getSubscriptionInner();
    if (subTable != null) {
        for (final Map.Entry<String, SubscriptionData> entry : subTable.entrySet()) {
            final String topic = entry.getKey();
            try {
                this.rebalanceByTopic(topic, isOrder);
            } catch (Throwable e) {
                if (!topic.startsWith(MixAll.RETRY_GROUP_TOPIC_PREFIX)) {
                    log.warn("rebalanceByTopic Exception", e);
                }
            }
        }
    }

    this.truncateMessageQueueNotMyTopic();
}

   doRebalance(final boolean isOrder)方法实现逻辑是：
   1.获取消费者的订阅信息，这里的订阅信息是应用在调用subscribe方法时写入的，具体如下：

public void subscribe(String topic, String subExpression) throws MQClientException {
    try {
        SubscriptionData subscriptionData = FilterAPI.buildSubscriptionData(this.defaultMQPushConsumer.getConsumerGroup(),
                topic, subExpression);
        this.rebalanceImpl.getSubscriptionInner().put(topic, subscriptionData);
        if (this.mQClientFactory != null) {
            this.mQClientFactory.sendHeartbeatToAllBrokerWithLock();
        }
    } catch (Exception e) {
        throw new MQClientException("subscription exception", e);
    }
}

   2.遍历订阅信息并执行rebalanceByTopic(topic, isOrder)
   rebalanceByTopic(topic, isOrder)实现逻辑按照消费模式分为广播消费和集群消费。

private void rebalanceByTopic(final String topic, final boolean isOrder) {
    switch (messageModel) {
        case BROADCASTING: {
            Set<MessageQueue> mqSet = this.topicSubscribeInfoTable.get(topic);
            if (mqSet != null) {
                boolean changed = this.updateProcessQueueTableInRebalance(topic, mqSet, isOrder);
                if (changed) {
                    this.messageQueueChanged(topic, mqSet, mqSet);
                    log.info("messageQueueChanged {} {} {} {}",
                        consumerGroup,
                        topic,
                        mqSet,
                        mqSet);
                }
            } else {
                log.warn("doRebalance, {}, but the topic[{}] not exist.", consumerGroup, topic);
            }
            break;
        }
        case CLUSTERING: {
            Set<MessageQueue> mqSet = this.topicSubscribeInfoTable.get(topic);
            List<String> cidAll = this.mQClientFactory.findConsumerIdList(topic, consumerGroup);
            if (null == mqSet) {
                if (!topic.startsWith(MixAll.RETRY_GROUP_TOPIC_PREFIX)) {
                    log.warn("doRebalance, {}, but the topic[{}] not exist.", consumerGroup, topic);
                }
            }

            if (null == cidAll) {
                log.warn("doRebalance, {} {}, get consumer id list failed", consumerGroup, topic);
            }

            if (mqSet != null && cidAll != null) {
                List<MessageQueue> mqAll = new ArrayList<MessageQueue>();
                mqAll.addAll(mqSet);

                Collections.sort(mqAll);
                Collections.sort(cidAll);

                AllocateMessageQueueStrategy strategy = this.allocateMessageQueueStrategy;

                List<MessageQueue> allocateResult = null;
                try {
                    allocateResult = strategy.allocate(
                        this.consumerGroup,
                        this.mQClientFactory.getClientId(),
                        mqAll,
                        cidAll);
                } catch (Throwable e) {
                    log.error("AllocateMessageQueueStrategy.allocate Exception. allocateMessageQueueStrategyName={}", strategy.getName(),
                        e);
                    return;
                }

                Set<MessageQueue> allocateResultSet = new HashSet<MessageQueue>();
                if (allocateResult != null) {
                    allocateResultSet.addAll(allocateResult);
                }

                boolean changed = this.updateProcessQueueTableInRebalance(topic, allocateResultSet, isOrder);
                if (changed) {
                    log.info(
                        "rebalanced result changed. allocateMessageQueueStrategyName={}, group={}, topic={}, clientId={}, mqAllSize={}, cidAllSize={}, rebalanceResultSize={}, rebalanceResultSet={}",
                        strategy.getName(), consumerGroup, topic, this.mQClientFactory.getClientId(), mqSet.size(), cidAll.size(),
                        allocateResultSet.size(), allocateResultSet);
                    this.messageQueueChanged(topic, mqSet, allocateResultSet);
                }
            }
            break;
        }
        default:
            break;
    }
}

  集群模式：
  （1）从topicSubscribeInfoTable中获取topic对应的MessageQueue集合mqSet
  （2）根据consumerGroup和topic从broker中获取其所有的consumer客户端ID列表cidAll（这里是随机选择一个broker来获取消费组内所有的消费者信息，因为在消费者启动的过程中会向所有的broker发送心跳信息，心跳信息中就包含注册的消费者）
  （3）对mqAll和cidAll进行排序（注意：排序可以保证一个consumerGroup中的所有consumer看到的视图保持一致，确保一个MessageQueue不会被多个消费者分配）
  （4）调用allocate方法为consumer分配MessageQueue，其结果存储在allocateResult并将其放在allocateResultSet中
  （5）调用updateProcessQueueTableInRebalance方法实现以下功能：对比新分配的MessageQueue及consumer当前的负载集合，看看是否有MessageQueue的变化，并依据变化做出不同的处理，具体如下：

private boolean updateProcessQueueTableInRebalance(final String topic, final Set<MessageQueue> mqSet,
    final boolean isOrder) {
    boolean changed = false;
    //遍历当前负载的MessageQueue集合，如果MessageQueue不在新分配的MessageQueue集合mqSet中
    //说明本次消息队列负载后该mq被分配给其他消费者了，这种情况需要暂停消息队列的消费，将其对应ProcessQueue的dropped
    //属性设置为true并调用removeUnnecessaryMessageQueue方法保存消费进度然后根据removeUnnecessaryMessageQueue方法返回结果
    //判断是否从processQueueTable中将其删除
    Iterator<Entry<MessageQueue, ProcessQueue>> it = this.processQueueTable.entrySet().iterator();
    while (it.hasNext()) {
        Entry<MessageQueue, ProcessQueue> next = it.next();
        MessageQueue mq = next.getKey();
        ProcessQueue pq = next.getValue();

        if (mq.getTopic().equals(topic)) {
            if (!mqSet.contains(mq)) {
                pq.setDropped(true);
                if (this.removeUnnecessaryMessageQueue(mq, pq)) {
                    it.remove();
                    changed = true;
                    log.info("doRebalance, {}, remove unnecessary mq, {}", consumerGroup, mq);
                }
            } else if (pq.isPullExpired()) {
                switch (this.consumeType()) {
                    case CONSUME_ACTIVELY:
                        break;
                    case CONSUME_PASSIVELY:
                        pq.setDropped(true);
                        if (this.removeUnnecessaryMessageQueue(mq, pq)) {
                            it.remove();
                            changed = true;
                            log.error("[BUG]doRebalance, {}, remove unnecessary mq, {}, because pull is pause, so try to fixed it",
                                consumerGroup, mq);
                        }
                        break;
                    default:
                        break;
                }
            }
        }
    }
    //遍历新分配的MessageQueue集合mqSet，如果processQueueTable中不包含mq则说明该mq是本次新增加的
    //这种情况下首先从内存中删除该mq的消费进度，然后从磁盘中读取该mq的消费进度并构建PullRequest
    List<PullRequest> pullRequestList = new ArrayList<PullRequest>();
    for (MessageQueue mq : mqSet) {
        if (!this.processQueueTable.containsKey(mq)) {
            if (isOrder && !this.lock(mq)) {
                log.warn("doRebalance, {}, add a new mq failed, {}, because lock failed", consumerGroup, mq);
                continue;
            }

            this.removeDirtyOffset(mq);
            ProcessQueue pq = new ProcessQueue();
            long nextOffset = this.computePullFromWhere(mq);
            if (nextOffset >= 0) {
                ProcessQueue pre = this.processQueueTable.putIfAbsent(mq, pq);
                if (pre != null) {
                    log.info("doRebalance, {}, mq already exists, {}", consumerGroup, mq);
                } else {
                    log.info("doRebalance, {}, add a new mq, {}", consumerGroup, mq);
                    PullRequest pullRequest = new PullRequest();
                    pullRequest.setConsumerGroup(consumerGroup);
                    pullRequest.setNextOffset(nextOffset);
                    pullRequest.setMessageQueue(mq);
                    pullRequest.setProcessQueue(pq);
                    pullRequestList.add(pullRequest);
                    changed = true;
                }
            } else {
                log.warn("doRebalance, {}, add new mq failed, {}", consumerGroup, mq);
            }
        }
    }
    //将PullRequest放到PullMessageService中的pullRequestQueue并唤醒PullMessageService线程
    this.dispatchPullRequest(pullRequestList);

    return changed;
}

  （6）判断rebalance结果是否发生变化，如果发生变化则执行messageQueueChanged(topic, mqSet, allocateResultSet)更新订阅信息的版本为当前时间、更新pullThresholdForQueue和pullThresholdSizeForQueue、通知所有的broker
  广播模式：
  (1)从topicSubscribeInfoTable中获取topic对应的MessageQueue集合mqSet
  (2)执行updateProcessQueueTableInRebalance方法，该方法在集群模式中已经详细解释
   3.执行truncateMessageQueueNotMyTopic()方法完成以下功能：获取消费者的订阅信息，遍历processQueueTable(记录的是消费者负载的MessageQueue及ProcessQueue)判断processQueueTable中的MessageQueue是否被包含在消费者的订阅关系中，如果没有包含，则从processQueueTable中删除MessageQueue并将其对应ProcessQueue的dropped属性设置为true(表示该ProcessQueue中的消息将不再被消费)

private void truncateMessageQueueNotMyTopic() {
    Map<String, SubscriptionData> subTable = this.getSubscriptionInner();

    for (MessageQueue mq : this.processQueueTable.keySet()) {
        if (!subTable.containsKey(mq.getTopic())) {

            ProcessQueue pq = this.processQueueTable.remove(mq);
            if (pq != null) {
                pq.setDropped(true);
                log.info("doRebalance, {}, truncateMessageQueueNotMyTopic remove unnecessary mq, {}", consumerGroup, mq);
            }
        }
    }
}

  在updateProcessQueueTableInRebalance方法中构建PullRequest(消息拉取任务)前会先计算拉取消息的nextOffset，计算nextOffset的方法是computePullFromWhere(MessageQueue mq)。在RocketMQ中提供了三种消费点（注意：如果是全新的消费者则从指定的消费点开始消费数据），分别是CONSUME_FROM_LAST_OFFSET、CONSUME_FROM_FIRST_OFFSET和CONSUME_FROM_TIMESTAMP，设置消费点是在应用中通过setConsumeFromWhere(…)方法来设置。
  （1）CONSUME_FROM_LAST_OFFSET（从队列最新的偏移量开始消费）
  （2）CONSUME_FROM_FIRST_OFFSET（从头开始消费）
  （3）CONSUME_FROM_TIMESTAMP（从指定的时间开始消费）
  上面三种消费点计算nextOffset的逻辑基本一致，都会从磁盘中获取消息队列的消费进度，根据消费进度可以分为三种情况：
  （1）消费进度大于等0，这种情况下将消费进度直接赋给nextOffset并返回；
  （2）消费进度等于-1，这种情况表示消费者第一次启动需要根据应用设置的消费点来分情况讨论：如果是CONSUME_FROM_LAST_OFFSET，则获取该消息队列当前最大的偏移量；如果是CONSUME_FROM_FIRST_OFFSET，则直接返回0从头开始消费；如果是CONSUME_FROM_TIMESTAMP，则根据应用设置的时间调用searchOffset来查找偏移量；
  （3）消费进度小于-1，表示该消息进度文件中存储了错误的偏移量，返回-1;

public long computePullFromWhere(MessageQueue mq) {
    long result = -1;
    final ConsumeFromWhere consumeFromWhere = this.defaultMQPushConsumerImpl.getDefaultMQPushConsumer().getConsumeFromWhere();
    final OffsetStore offsetStore = this.defaultMQPushConsumerImpl.getOffsetStore();
    switch (consumeFromWhere) {
        case CONSUME_FROM_LAST_OFFSET_AND_FROM_MIN_WHEN_BOOT_FIRST:
        case CONSUME_FROM_MIN_OFFSET:
        case CONSUME_FROM_MAX_OFFSET:
        case CONSUME_FROM_LAST_OFFSET: {
            long lastOffset = offsetStore.readOffset(mq, ReadOffsetType.READ_FROM_STORE);
            if (lastOffset >= 0) {
                result = lastOffset;
            }
            // First start,no offset
            else if (-1 == lastOffset) {
                if (mq.getTopic().startsWith(MixAll.RETRY_GROUP_TOPIC_PREFIX)) {
                    result = 0L;
                } else {
                    try {
                        result = this.mQClientFactory.getMQAdminImpl().maxOffset(mq);
                    } catch (MQClientException e) {
                        result = -1;
                    }
                }
            } else {
                result = -1;
            }
            break;
        }
        case CONSUME_FROM_FIRST_OFFSET: {
            long lastOffset = offsetStore.readOffset(mq, ReadOffsetType.READ_FROM_STORE);
            if (lastOffset >= 0) {
                result = lastOffset;
            } else if (-1 == lastOffset) {
                result = 0L;
            } else {
                result = -1;
            }
            break;
        }
        case CONSUME_FROM_TIMESTAMP: {
            long lastOffset = offsetStore.readOffset(mq, ReadOffsetType.READ_FROM_STORE);
            if (lastOffset >= 0) {
                result = lastOffset;
            } else if (-1 == lastOffset) {
                if (mq.getTopic().startsWith(MixAll.RETRY_GROUP_TOPIC_PREFIX)) {
                    try {
                        result = this.mQClientFactory.getMQAdminImpl().maxOffset(mq);
                    } catch (MQClientException e) {
                        result = -1;
                    }
                } else {
                    try {
                        long timestamp = UtilAll.parseDate(this.defaultMQPushConsumerImpl.getDefaultMQPushConsumer().getConsumeTimestamp(),UtilAll.YYYYMMDDHHMMSS).getTime();
                        result = this.mQClientFactory.getMQAdminImpl().searchOffset(mq, timestamp);
                    } catch (MQClientException e) {
                        result = -1;
                    }
                }
            } else {
                result = -1;
            }
            break;
        }

        default:
            break;
    }

    return result;
}

  现在回答文章开始提出的问题，消费者在第一次启动时是如何通过ConsumeQueue来查找消息？在consumer启动的过程中会启动RebalanceService服务，在RebalanceService启动的过程中会构造PullRequest(消息拉取任务)，在PullRequest中会设置获取消息的nextOffset（ConsumeQueue中的逻辑偏移量），而计算nextOffset的方法参照computePullFromWhere。PullRequest会被放到PullMessageService的pullRequestQueue队列中并唤醒PullMessageService，PullMessageService负责将PullRequest发送到broker，broker会根据请求封装好数据并将其发送给消费者消费。