图文 | 探究分析RocketMQ 6种负载均衡策略~

前言

当一个ConsumerGroup中存在多个Consumer时，此时投递一个消息，应该把这个消息分配给哪个Consumer？

有的小伙伴脑海中可能马上就浮现出，随机 、轮询 等等分配策略.....

想法没错，RocektMq中设计的Consumer分配策略（负载均衡）中正包含了这两种，通过Consumer的负载均衡策略，可将主题内的消息分配给指定消费者分组中的多个消费者共同分担，提高消费并发能力和消费者的水平扩展能力。

---

源码

RocketMq中提供了message queue选择的策略接口AllocateMessageQueueStrategy

其中定义了两个方法

1. allocate: 执行策略，获取策略，获取所选择的消息队列集合
2. getName: 获取策略的名称

/**
 * Strategy Algorithm for message allocating between consumers
 */
public interface AllocateMessageQueueStrategy {

  /**
     * Allocating by consumer id
     *
     * @param consumerGroup current consumer group
     * @param currentCID current consumer id
     * @param mqAll message queue set in current topic
     * @param cidAll consumer set in current consumer group
     * @return The allocate result of given strategy
     */
  List<MessageQueue> allocate(
    final String consumerGroup,
    final String currentCID,
    final List<MessageQueue> mqAll,
    final List<String> cidAll
  );

  /**
     * Algorithm name
     *
     * @return The strategy name
     */
  String getName();
}

通过继承关系我们可以看到，RocketMq为我们提供了6种策略

• AllocateMessageQueueAveragely：平均负载策略
• AllocateMessageQueueAveragelyByCircle：环形平均负载策略
• AllocateMessageQueueByConfig：配置负载策略
• AllocateMessageQueueConsistentHash：一致性哈希负载策略
• AllocateMessageQueueByMachineRoom：同机房分配策略
• AllocateMachineRoomNearby ：AllocateMessageQueueByMachineRoom策略的升级版本

下面就具体来了解了解这六种策略~

---

AllocateMessageQueueAveragely

AllocateMessageQueueAveragely: 平均分配策略

在看源码之前，先举个例子，方便能快速认知算法的分配逻辑

举例：现在有10个message queue、3个Consumer，那么分配结果如下:

consumer1分配queueId= 0,1,2,3

consumer2分配queueId= 4,5,6

consumer3分配queueId= 7,8,9

public class AllocateMessageQueueAveragely extends AbstractAllocateMessageQueueStrategy {

  @Override
  public List<MessageQueue> allocate(String consumerGroup, String currentCID, List<MessageQueue> mqAll,
                                     List<String> cidAll) {
    // 参数校验
    List<MessageQueue> result = new ArrayList<MessageQueue>();
    if (!check(consumerGroup, currentCID, mqAll, cidAll)) {
      return result;
    }

    // 找到当前Consumer所处的index
    int index = cidAll.indexOf(currentCID);
    // 取模，如果mod = 0，则说明能够平均分配
    int mod = mqAll.size() % cidAll.size();
    
    // 1. 如果Consumer个数大于等于队列数，那么每个Consumer分配一个队列即可，多余的Consumer则分配不到
    // 2. 如果不能平均分配且index小于mode，则在平均分配以后再加上一个
    int averageSize =
      mqAll.size() <= cidAll.size() ? 1 : (mod > 0 && index < mod ? mqAll.size() / cidAll.size()
                                           + 1 : mqAll.size() / cidAll.size());
    
    // 开始分配的index，正常情况下 index * averageSize即可
    // 但是如果index > mod, 说明当前Consumer不属于多分配的范围之内，需要再 + mod,代表已经多分配的数量
    int startIndex = (mod > 0 && index < mod) ? index * averageSize : index * averageSize + mod;
    
    // 被分配的消息队列的范围上限
    // 因为当Consumer个数大于队列个数时，靠后的Consumer已经没有队列可以分配了
    int range = Math.min(averageSize, mqAll.size() - startIndex);
    for (int i = 0; i < range; i++) {
      result.add(mqAll.get((startIndex + i) % mqAll.size()));
    }
    
    return result;
  }

  @Override
  public String getName() {
    return "AVG";
  }
}

总结一下

首先会用队列数与Consumer总数取模 ，得到mod

如果**mod == 0，那么说明可以 完全平均分配**，每个Consumer可以分配到**mqAll.size() / cidAll.size()个**队列

如果**mod > 0**，那么又分为两种情况，

1. 第一种为队列数小于Consumer个数 ，例如: 5个Consumer，3个队列，那么只有三个Consumer可以每个分配1个队列，剩余两个Consumer无法分配

2. 第二种就是队列数大于Consumer个数 ，例如: 4个Consumer，10个队列，那么mod = 2，那么说明在平均每个Consumer分配10 / 4 = 2个队列的情况下，还多出来2个可以被分配的队列，那么这两个队列的分配我理解为优先分配 ，即优先给前面的队列分配

   根据源码，当index < mod时，平均分配的结果是mqAll.size() / cidAll.size() + 1，即在平均分配的基础上+ 1，所以当index = 0、1的时候，即最前面两个Consumer可以实现多分配

---

AllocateMessageQueueAveragelyByCircle

AllocateMessageQueueAveragelyByCircle: 环形平均分配策略

举例：现在有10个message queue、3个Consumer，那么分配结果如下:

consumer1分配queueId= 0,1,2,9

consumer2分配queueId= 3,4,5

consumer3分配queueId= 6,7,8

public class AllocateMessageQueueAveragelyByCircle extends AbstractAllocateMessageQueueStrategy {

  @Override
  public List<MessageQueue> allocate(String consumerGroup, String currentCID, List<MessageQueue> mqAll,
                                     List<String> cidAll) {

    List<MessageQueue> result = new ArrayList<MessageQueue>();
    if (!check(consumerGroup, currentCID, mqAll, cidAll)) {
      return result;
    }

    // 当前Consumer的index
    int index = cidAll.indexOf(currentCID);
    
    // 从index开始遍历mqAll
    for (int i = index; i < mqAll.size(); i++) {
      // 与Consumer总个数取模，如果等于当前Consumer的index，则代表归属于当前Consumer
      if (i % cidAll.size() == index) {
        result.add(mqAll.get(i));
      }
    }
    return result;
  }

  @Override
  public String getName() {
    return "AVG_BY_CIRCLE";
  }
}

这个环形平均分配策略 简单来说就是一人一个

如果队列个数有10个，Consumer个数有3个

那么就C1、C2、C3、C1、C2、C3、C1、C2、C3、C1，一人一个分配

最终C1分配4个，C2、C3分配3个

---

AllocateMessageQueueConsistentHash

AllocateMessageQueueConsistentHash: 一致性hash分配策略

public class AllocateMessageQueueConsistentHash extends AbstractAllocateMessageQueueStrategy {

  // 虚拟节点个数（防止数据倾斜）
  private final int virtualNodeCnt;
  
  // 自定义key hash计算function
  private final HashFunction customHashFunction;

  public AllocateMessageQueueConsistentHash() {
    this(10);
  }

  public AllocateMessageQueueConsistentHash(int virtualNodeCnt) {
    this(virtualNodeCnt, null);
  }

  public AllocateMessageQueueConsistentHash(int virtualNodeCnt, HashFunction customHashFunction) {
    if (virtualNodeCnt < 0) {
      throw new IllegalArgumentException("illegal virtualNodeCnt :" + virtualNodeCnt);
    }
    this.virtualNodeCnt = virtualNodeCnt;
    this.customHashFunction = customHashFunction;
  }

  @Override
  public List<MessageQueue> allocate(String consumerGroup, String currentCID, List<MessageQueue> mqAll,
                                     List<String> cidAll) {

    List<MessageQueue> result = new ArrayList<MessageQueue>();
    if (!check(consumerGroup, currentCID, mqAll, cidAll)) {
      return result;
    }

    // 将consumerId封装为节点
    Collection<ClientNode> cidNodes = new ArrayList<ClientNode>();
    for (String cid : cidAll) {
      cidNodes.add(new ClientNode(cid));
    }

    // 生成hash环
    final ConsistentHashRouter<ClientNode> router;
    if (customHashFunction != null) {
      // 存在自定义key hash计算function
      router = new ConsistentHashRouter<ClientNode>(cidNodes, virtualNodeCnt, customHashFunction);
    } else {
      router = new ConsistentHashRouter<ClientNode>(cidNodes, virtualNodeCnt);
    }

    List<MessageQueue> results = new ArrayList<MessageQueue>();
    // 遍历所有队列
    for (MessageQueue mq : mqAll) {
      // 进行hash路由，得到路由到的节点
      ClientNode clientNode = router.routeNode(mq.toString());
      // 如果路由节点key(节点封装的consumerId)与当前consumerId一致，则归属于当前节点
      if (clientNode != null && currentCID.equals(clientNode.getKey())) {
        results.add(mq);
      }
    }

    return results;

  }

  @Override
  public String getName() {
    return "CONSISTENT_HASH";
  }

  private static class ClientNode implements Node {
    private final String clientID;

    public ClientNode(String clientID) {
      this.clientID = clientID;
    }

    @Override
    public String getKey() {
      return clientID;
    }
  }

}

一致性hash策略，将所有consumerId封装为节点，再构建hash环

在构造时会创建虚拟节点，防止出现数据倾斜 ，即virtualNodeCnt字段，代表创建的虚拟节点的个数

同时也支持自定义key的hash计算function

hash环构造完毕，此时遍历所有messageQueue，计算当前队列所属的节点，如果节点封装的consumerId跟当前的consumerId一致，就代表该队列归属于当前Consumer

如果没有自定义key的hash function，则默认使用的是MD5，源码如下

private static class MD5Hash implements HashFunction {
  MessageDigest instance;

  public MD5Hash() {
    try {
      // md5
      instance = MessageDigest.getInstance("MD5");
    } catch (NoSuchAlgorithmException e) {
    }
  }

  @Override
  public long hash(String key) {
    instance.reset();
    instance.update(key.getBytes());
    byte[] digest = instance.digest();

    long h = 0;
    for (int i = 0; i < 4; i++) {
      h <<= 8;
      h |= ((int) digest[i]) & 0xFF;
    }
    return h;
  }
}

---

AllocateMessageQueueByConfig

AllocateMessageQueueByConfig: 配置分配策略

这个没啥好说的，在执行前调用setMessageQueueList设置分配的MessageQueueList

public class AllocateMessageQueueByConfig extends AbstractAllocateMessageQueueStrategy {
  private List<MessageQueue> messageQueueList;

  @Override
  public List<MessageQueue> allocate(String consumerGroup, String currentCID, List<MessageQueue> mqAll,
                                     List<String> cidAll) {
    // 返回手动设置的messageQueueList
    return this.messageQueueList;
  }

  @Override
  public String getName() {
    return "CONFIG";
  }

  public List<MessageQueue> getMessageQueueList() {
    return messageQueueList;
  }

  // 手动设置messageQueueList
  public void setMessageQueueList(List<MessageQueue> messageQueueList) {
    this.messageQueueList = messageQueueList;
  }
}

---

AllocateMessageQueueByMachineRoom

AllocateMessageQueueByMachineRoom: 同机房负载策略

public class AllocateMessageQueueByMachineRoom extends AbstractAllocateMessageQueueStrategy {
  
  // 机房名称
  private Set<String> consumeridcs;

  @Override
  public List<MessageQueue> allocate(String consumerGroup, String currentCID, List<MessageQueue> mqAll,
                                     List<String> cidAll) {

    List<MessageQueue> result = new ArrayList<MessageQueue>();
    if (!check(consumerGroup, currentCID, mqAll, cidAll)) {
      return result;
    }
    
    // 计算当前consumer的index
    int currentIndex = cidAll.indexOf(currentCID);
    if (currentIndex < 0) {
      return result;
    }
    
    // 遍历所有MessageQueue，过滤出符合指定机房的msgQueue
    List<MessageQueue> premqAll = new ArrayList<MessageQueue>();
    for (MessageQueue mq : mqAll) {
      String[] temp = mq.getBrokerName().split("@");
      if (temp.length == 2 && consumeridcs.contains(temp[0])) {
        premqAll.add(mq);
      }
    }

    // 此处跟平均分配一样，就不过多赘述了
    int mod = premqAll.size() / cidAll.size();
    int rem = premqAll.size() % cidAll.size();
    int startIndex = mod * currentIndex;
    int endIndex = startIndex + mod;
    for (int i = startIndex; i < endIndex; i++) {
      result.add(premqAll.get(i));
    }
    if (rem > currentIndex) {
      result.add(premqAll.get(currentIndex + mod * cidAll.size()));
    }
    return result;
  }

  @Override
  public String getName() {
    return "MACHINE_ROOM";
  }

  public Set<String> getConsumeridcs() {
    return consumeridcs;
  }

  public void setConsumeridcs(Set<String> consumeridcs) {
    this.consumeridcs = consumeridcs;
  }
}

同机房负载策略，其根本还是用的平均负载策略，只不过在执行平均负载前，会过滤出同一机房的messageQueue ，在此基础上再进行的平均负载

所以在执行该策略前，需要设置consumeridcs，即目标机房名称

---

AllocateMachineRoomNearby

AllocateMachineRoomNearby: 就近机房负载策略

public class AllocateMachineRoomNearby extends AbstractAllocateMessageQueueStrategy {

  // 分配策略
  private final AllocateMessageQueueStrategy allocateMessageQueueStrategy;//actual allocate strategy
  
  // 机房匹配
  private final MachineRoomResolver machineRoomResolver;

  public AllocateMachineRoomNearby(AllocateMessageQueueStrategy allocateMessageQueueStrategy,
                                   MachineRoomResolver machineRoomResolver) throws NullPointerException {
    if (allocateMessageQueueStrategy == null) {
      throw new NullPointerException("allocateMessageQueueStrategy is null");
    }

    if (machineRoomResolver == null) {
      throw new NullPointerException("machineRoomResolver is null");
    }

    this.allocateMessageQueueStrategy = allocateMessageQueueStrategy;
    this.machineRoomResolver = machineRoomResolver;
  }

  @Override
  public List<MessageQueue> allocate(String consumerGroup, String currentCID, List<MessageQueue> mqAll,
                                     List<String> cidAll) {

    List<MessageQueue> result = new ArrayList<MessageQueue>();
    if (!check(consumerGroup, currentCID, mqAll, cidAll)) {
      return result;
    }

    // 将MessageQueue进行分组
    Map<String/*machine room */, List<MessageQueue>> mr2Mq = new TreeMap<String, List<MessageQueue>>();
    for (MessageQueue mq : mqAll) {
      // 匹配msgQueue对应的broker分区
      String brokerMachineRoom = machineRoomResolver.brokerDeployIn(mq);
      if (StringUtils.isNoneEmpty(brokerMachineRoom)) {
        if (mr2Mq.get(brokerMachineRoom) == null) {
          mr2Mq.put(brokerMachineRoom, new ArrayList<MessageQueue>());
        }
        mr2Mq.get(brokerMachineRoom).add(mq);
      } else {
        throw new IllegalArgumentException("Machine room is null for mq " + mq);
      }
    }

    // 将consumer分组
    Map<String/*machine room */, List<String/*clientId*/>> mr2c = new TreeMap<String, List<String>>();
    for (String cid : cidAll) {
      // 匹配consumer对应的分区
      String consumerMachineRoom = machineRoomResolver.consumerDeployIn(cid);
      if (StringUtils.isNoneEmpty(consumerMachineRoom)) {
        if (mr2c.get(consumerMachineRoom) == null) {
          mr2c.put(consumerMachineRoom, new ArrayList<String>());
        }
        mr2c.get(consumerMachineRoom).add(cid);
      } else {
        throw new IllegalArgumentException("Machine room is null for consumer id " + cid);
      }
    }

    List<MessageQueue> allocateResults = new ArrayList<MessageQueue>();
    // 找到当前consumer对应的分区，并拿到分区对应的msgQueue以及所有consumer
    String currentMachineRoom = machineRoomResolver.consumerDeployIn(currentCID);
    List<MessageQueue> mqInThisMachineRoom = mr2Mq.remove(currentMachineRoom);
    List<String> consumerInThisMachineRoom = mr2c.get(currentMachineRoom);
    
    // 如果都存在，则在mqInThisMachineRoom、consumerInThisMachineRoom前提下，根据传入的分配策略给当前consumer进行分配
    if (mqInThisMachineRoom != null && !mqInThisMachineRoom.isEmpty()) {
      allocateResults.addAll(allocateMessageQueueStrategy.allocate(consumerGroup, currentCID, mqInThisMachineRoom, consumerInThisMachineRoom));
    }

    // 如果有的队列所属的分区，没有同分区consumer，那么将这些队列按照自定义的分配策略分配给所有consumer
    for (Entry<String, List<MessageQueue>> machineRoomEntry : mr2Mq.entrySet()) {
      if (!mr2c.containsKey(machineRoomEntry.getKey())) { // no alive consumer in the corresponding machine room, so all consumers share these queues
        allocateResults.addAll(allocateMessageQueueStrategy.allocate(consumerGroup, currentCID, machineRoomEntry.getValue(), cidAll));
      }
    }

    return allocateResults;
  }

  @Override
  public String getName() {
    return "MACHINE_ROOM_NEARBY" + "-" + allocateMessageQueueStrategy.getName();
  }

  /**
     * A resolver object to determine which machine room do the message queues or clients are deployed in.
     *
     * AllocateMachineRoomNearby will use the results to group the message queues and clients by machine room.
     *
     * The result returned from the implemented method CANNOT be null.
     */
  public interface MachineRoomResolver {
    
    // 解析messageQueue所属的broker属于哪个机房
    String brokerDeployIn(MessageQueue messageQueue);

    // 解析consumerId属于哪个机房
    String consumerDeployIn(String clientID);
    
  }
}

就近机房分配策略，我们需要自定义MachineRoomResolver来定义consumer以及messageQueue的分区分配策略，以及自定义的 allocateMessageQueueStrategy

在执行分配时，会将所有consuemr、messageQueue按照自定义的MachineRoomResolver来进行分组

随后找到当前consumer所属的分区，拿到同一分区的所有consumer、messageQueue，如果都不为空则根据自定义的allocateMessageQueueStrategy进行分配

此外，如果有的队列所属的分区，没有同分区consumer，那么将这些队列按照自定义的分配策略分配给所有consumer

---

总结

本文总结了RocketMq关于Consumer的6中分配策略，如有误，请指正~