造轮子记录——内存攒批队列设计

一、背景

"攒批"指的是收到数据之后，不需要立马调用下游业务处理，需要在内存中进行积攒，等到满足条件后（一般是积攒总数目到达最大值，或者已经积攒的元素总大小超过一个设定阈值）再执行数据处理。

想象下，在下面的流水线中，一个瓶子的内容没装满之前显然不能装箱，需要"攒满"一瓶的"数据"之后才能交付下游进行装箱（进行数据消费）。

设计攒批队列最开始的起因是负责的业务中台需要在收到消息之后去调用下游业务的注册接口，但是消息一到就去调用下游接口，就可能导致下游接口的qps压力非常大，多积攒一点数据再去调用下游接口，对于整体调用效率和接口稳定性层面都会有很大提升。

目前开源代码中的攒批队列支持下述功能：

• 使用简单，配置项少，线程安全
• 支持分组攒批
• 支持设定内存上限，可有效避免JVM内存占用问题
• 支持设定时间红线、数量红线、内存占用红线种规则，达到红线后会自动启动消息处理逻辑
• 支持重试
• 回调能力丰富，成功、失败都有对应的回调接口可供实现
• 提供可观测的API，可以了解当前攒批任务的各项指标
• 支持优雅关闭，系统关闭时会将内存中还未被发送的数据全部触发消费逻辑

需要关注的是攒批队列是个单机内存版本的攒批，不适用于分布式场景"攒批"。

攒批队列（waiting-bus-batch 以下简称WBB）项目内部版本已经在多个组件中进行生产使用，功能正常，但是由于是开源项目，相比内部正式版本做了大量代码简化，旨在为读者提供一种设计思路，如需使用到生产环境建议，建议做好功能测试和压力测试工作。

开源代码地址：

github：github.com/AHUCodingBe...

原创不易，如果代码或者文章中的部分设计思路对读者有帮助，欢迎收藏点赞。

二、关键概念

WBB抽象了下述概念：

• 攒批触发条件（Condition） ：接收到的消息追加到攒批队列里，显然不能无限的进行攒批，一般的，当攒批队列的元素总数目大于指定数目，或者队列占据的内存空间超过指定值、以及队列在内存中驻留的时间超过最大值时、都会触发业务指定的消费逻辑。
• 攒批队列（ProducerBatch） ：所谓攒批队列就是指内存中的一个队列，可以往这个队列中持续追加数据，也就是所谓的"攒批"，当攒批队列达到发送条件，攒批队列会转化为一个"攒批消费任务"并交给IO线程池处理，攒批队列的生命周期在消费任务成功执行或者执行失败之后就会结束。
• 攒批消费任务（ProducerBatchTask） ：攒批消费任务指的是将满足攒批触发条件的攒批队列作为任务入参，并调用业务消费逻辑进行消费的过程，需要注意的是在WBB的设计思路中攒批消费任务是在单独的线程池中执行的。
• 攒批分组（batchGroup） ：攒批分组的设计借鉴了Kafka Topic分组的思想，例如在消费上游消息进行内存攒批的场景中，我们接收到的消息由于Tag不同需要加入到不同的分组进行攒批时，就可以使用到"攒批分组"。攒批分组的生命周期比"攒批队列"要长很多，一个分组中一个时刻总会有1个活跃的攒批队列在进行攒批，如下图所示groupA中存在一个攒批ProducerBatch-groupA，当这个攒批攒满以后，会重新创建一个新的攒批队列再次进行攒批，但是这个新的攒批队列所属的分组名称是不变的，仍然是groupA。

• 攒批容器（ProducerBatchContainer） ：攒批容器则是用来管理多个攒批分组的容器，通过这个攒批容器可以知道当前有多少攒批队列在进行攒批。
• 攒批回调（CallBack） ：攒批回调是个可选功能，如果在业务场景中我们需要知悉每次追加的消息内容是否被成功消费，则可以在将数据加入攒批队列时指定回调，这样当这批数据被消费的时候就会触发成功回调函数。

三、整体架构

四、快速入门

（1）创建一个生产者配置类，最简单的生产者配置示例如下：

    ProducerConfig producerConfig = new ProducerConfig();
    // 设定攒批达到20条以上的时候再执行业务逻辑
    producerConfig.setBatchCountThreshold(20);
    // 设定攒批已经达到100s的时候再执行业务逻辑
    producerConfig.setLingerMs(100_000);
    // 可以指定攒批重试次数
    producerConfig.setRetries(2); 

此外还可以再去设置攒批最大占据内存大小，攒批容器最大占据内存大小，等参数。详情可以参考 ProducerConfig 类进行设定。

（2）创建 MessageProducer 对象，同时指定攒批条件达成时，需要执行的业务逻辑，示例如下：

    private static MessageProducer getMessageProducer(ProducerConfig producerConfig) {
        return new MessageProducer(producerConfig, (List<Message> arr) -> {
            // do your business here
            return MessageProcessResultEnum.SUCCESS;
        });
    }

在消费代码中，如果主动返回 MessageProcessResultEnum.SUCCESS; WBB会认为当前攒批已经被下游成功消费，返回 SUCCESS 则认为消费成功，返回RETRY 则会进入重试队列，等待重试。需要关注的是如果你的业务代码处理时出现了异常也会加入重试队列。

（3）接收上游的消息开启内存攒批，这里以RocketMQ消息消费为例：

    public static void main(String[] args) {
      
          DefaultMQPushConsumer consumer = new DefaultMQPushConsumer("consumer_group_test");

          ProducerConfig producerConfig = new ProducerConfig();
          // 设定攒批达到20条以上的时候再执行业务逻辑
          producerConfig.setBatchCountThreshold(20);
          // 设定攒批已经达到100s的时候再执行业务逻辑
          producerConfig.setLingerMs(100_000);
          // 可以指定重试次数
          producerConfig.setRetries(2); 
          MessageProducer messageProducer = getMessageProducer(producerConfig);

          try {
              consumer.setNamesrvAddr("localhost:9876");
              consumer.subscribe("test_topic", "*");
              consumer.registerMessageListener(new MessageListenerConcurrently() {
                  @Override
                  public ConsumeConcurrentlyStatus consumeMessage(List<MessageExt> messages, ConsumeConcurrentlyContext context) {
                      List<Message> messageList = new ArrayList<>();
                      for (MessageExt message : messages) {
                          Message msg = convert(message);
                          messageList.add(msg);
                      }
                      // 调用攒批发送器  去发送消息，开启内存攒批
                      messageProducer.send(null, messageList);
                      return ConsumeConcurrentlyStatus.CONSUME_SUCCESS;
                  }
              });
              consumer.start();
          } catch (Exception e) {
              e.printStackTrace();
          }
      }

调用messageProducer的send方法以后，就会将收到的消息加入到攒批过程中，在上面的例子中当攒批队列长度达到20个元素以上或者攒批时间已经超过100s时（或者攒批队列的占据最大内存超过设定值，默认8M）就会触发攒批消费，也就是会执行创建MessageProducer 时指定的业务逻辑。

此外WBB还支持设置攒批回调，也就是为每一次往攒批队列中追加操作设置回调函数，如下面的示例代码所示，当消息被追加到攒批队列中并被成功消费后，会执行回调函数。

    ListenableFuture<Result> listenableFuture = messageProducer.send(null, messageList, (result) -> {
        if (result.isSuccessful()) {
            System.out.println("执行成功 SequenceNo=" + appendSequenceNo);
        } else {
            System.out.println("执行失败 SequenceNo=" + appendSequenceNo);
        }
    });
    // 也可以直接为 listenableFuture 绑定监听 了解每次追加的消息的消费结果

五、关键设计与源码分析

5.1 WBB核心对象建模

面向对象设计思路最重要的就是对事物的抽象，在WBB中主要的建模有两个，一个是"攒批队列"一个是"攒批容器"。

攒批队列的基本属性和行为如下，一个攒批队列实际最重要的就是一个List结构，此外还需要记录重试信息、当前队列的空间占用等信息，并暴露tryAppend() 方法由攒批容器来调用，从而进行数据追加。

攒批容器则是主要是对多个攒批队列的管理，承担管理者的角色，对于需要攒批的数据需要找到匹配的攒批队列将数据追加进去，此外还需要对外暴露部分接口，让外部领域了解攒批的具体运行情况。

5.2 攒批功能实现

在WBB中一个攒批操作，实际上就是一个将消息追加到消息队列中的动作。核心入口位于com.waiting.bus.core.containers.ProducerBatchContainer#doAppend方法中。在这个方法中通过 Semaphore 来保证所有攒批队列的总容量不会超过设定值（避免内存爆炸问题），每次往队列中追加元素，都需要使用semaphore来确保能够申请到一批存储空间。此外在攒批时会考虑是否满足"转化条件"也就是说是不是要停止攒批开始消费过程，关键代码如下：

    private void doAppend(List<Message> messageList, String groupName) throws InterruptedException, ProducerException {
        if (closed) {
            throw new IllegalStateException("cannot append after the waitingBus container was closed");
        }

        int sizeInBytes = DataSizeCalculator.calculateMessageByteSize(messageList);
        ensureSize(sizeInBytes);

        long maxBlockMs = producerConfig.getMaxBlockMs();
        if (maxBlockMs >= 0) {
             // 申请内存资源
            boolean acquired = semaphore.tryAcquire(sizeInBytes, maxBlockMs, TimeUnit.MILLISECONDS);
            if (!acquired) {
                throw new TimeoutException("failed to acquire memory 
                                           within the configured max blocking time " +
                producerConfig.getMaxBlockMs() + " ms");
            }
        } else {
            // 申请内存资源
            semaphore.acquire(sizeInBytes);
        }

        try {
            ProducerBatchHolder holder = getOrCreateProducerBatchHolder(groupName);
            synchronized (holder) {
                appendToHolder(groupName, messageList, sizeInBytes, holder);
            }
        } catch (Exception e) {
            semaphore.release(sizeInBytes);
            throw new ProducerException(e);
        }
    }

    private synchronized void appendToHolder(String groupName, List<Message> messages, 
                                             int sizeInBytes, ProducerBatchHolder holder) {
        if (holder.producerBatch != null) {
            boolean appendResult = holder.producerBatch.tryAppend(messages, sizeInBytes);
            if (appendResult) {
                if (holder.producerBatch.isMeetSendCondition()) {
                    holder.transferProducerBatch(ioThreadPool, producerConfig, messageProcessFunction, retryQueue, successQueue, failureQueue);
                }
                return;
            } else {
                holder.transferProducerBatch(ioThreadPool, producerConfig, messageProcessFunction, retryQueue, successQueue, failureQueue);
            }
        }
        holder.producerBatch = new ProducerBatch(producerConfig.getBatchSizeThresholdInBytes(), producerConfig.getBatchCountThreshold(), producerConfig.getMaxReservedAttempts(), groupName);
        LOGGER.info("groupName={} has created a new batch  batchId={}", groupName, holder.producerBatch.getBatchId());
        holder.producerBatch.tryAppend(messages, sizeInBytes);
        batchCount.incrementAndGet();
        // 检查攒批是否达到预设条件
        if (holder.producerBatch.isMeetSendCondition()) {
            holder.transferProducerBatch(ioThreadPool, producerConfig, messageProcessFunction, retryQueue, successQueue, failureQueue);
        }
    }
 

细心的读者可能发现了 追加消息的时候检查只能检查攒批队列的长度或者占据的内存大小是否触发攒批发送条件，如果设定的条件最多内存驻留5s，时间到时无论攒批了多少元素，都需要进行下游消费，这种攒批条件应该如何实现。事实上，WBB的攒批时间检测则是交给ExpireBatchHandler 来实现的，核心代码如下图所示：

    private void loopHandleExpireBatch() {
        while (!closed) {
            try {
                // 从攒批容器中获取过期"攒批"
                List<ProducerBatch> expiredBatches1 = producerBatchContainer.getExpiredBatches();
                // 从重试队列中获取过期的"攒批"
                List<ProducerBatch> expiredBatches2 = retryQueue.getExpiredBatches(1000);
                expiredBatches1.addAll(expiredBatches2);
                for (ProducerBatch b : expiredBatches1) {
                    ioThreadPool.submit(new SendProducerBatchTask(b, producerConfig, messageProcessFunction, retryQueue, successQueue, failureQueue));
                }
            } catch (Exception e) {
                LOGGER.error("loopHandleExpireBatch Exception, e=", e);
            }
        }
    }

ExpireBatchHandler本质是一个线程对象，这个线程在启动之后会不断检测 攒批容器和重试队列是否存在过期的攒批任务，如果存在就会将这些到达存活时间的攒批队列转为"攒批消费任务"并执行指定的业务逻辑。

5.3 攒批重试

WBB支持重试逻辑，当业务对攒批的处理出现异常或者明确返回失败的时候，如果还有重试机会，WBB会将执行失败的攒批队列加入重试队列，整个重试的逻辑使用的是JAVA并发包的神器 ------ DelayQueue

    DelayQueue<ProducerBatch> retryDelayBatchesQueue = new DelayQueue<ProducerBatch>();

业务执行失败的攒批如果还有重试机会就会加入到DelayQueue 进行排队。这里额外解释下DelayQueue的功能，DelayQueue 是一个无界的 BlockingQueue，用于放置实现了 Delayed 接口的对象，其中的对象只能在其到期时才能从队列中取走。这种队列是线程安全的，适合用来实现延时功能。

在WBB中当一个 ProducerBatch 攒批失败以后，会给这个攒批计算一个延时时间值，保证这个攒批只有在延时时间到了以后才能被再次取到。如下面的代码所示 在加入重试队列之前会对这个队列计算一个下次能够被取出的时间点，这样在指定时间点未到的时候无法从队列中取出这个元素。

    long retryBackoffMs = calculateRetryBackoffMs();
    batch.setNextRetryMs(System.currentTimeMillis() + retryBackoffMs);
    retryQueue.put(batch);

有兴趣的同学可以去了解下 DelayQueue的poll方法实现，代码很短，其实就是一个经典的"生产者消费者"模型。这里也考察下读者，如果现在让你写出生产者消费者模型的代码你能有多少种实现方案？这也是个面试常考的问题......

在WBB设计中使用到了 BlockingQueue、EvictingQueue（GUAVA） 、DelayQueue三种队列，后期可能会总结一篇文章专门介绍下这些队列的使用的原理。

此外有过业务开发经验的同学一般都会清楚，重试一般情况下会有个延时梯度，也就是说第一次失败以后会间隔较短的一段时间重试，后面每次重试都会加大重试时间，在WBB的重试的时间模型设计的比较简单，就是在用户设定的基础重试时间的基础上每次乘以2，比如用户规定基础重试时间是1s那么后面依次的重试时间为2s、4s、8s.....直到达到用户规定的最大重试时间, 代码核心实现如下。

    private long calculateRetryBackoffMs() {
        int retry = batch.getRetries();
        long retryBackoffMs = producerConfig.getBaseRetryBackoffMs() * LongMath.pow(2, retry);
        if (retryBackoffMs <= 0) {
            retryBackoffMs = producerConfig.getMaxRetryBackoffMs();
        }
        return Math.min(retryBackoffMs, producerConfig.getMaxRetryBackoffMs());
    }

5.3 优雅关闭实现

优雅关闭是一个很重要的需求，考虑下面的几种情况。

情况1： 需要关停组件，但是此时还有数据在内存攒批没有达到攒批完成条件触发消费，这时候如果不做任何处理可能这批数据就没能正确的执行业务逻辑。

情况2： 需要关停组件，但是此时还有消息正在加入攒批队列，还没有成功追加到攒批队列中，此时如果直接关停，可能导致这些消息没有被正确处理。

情况3： 需要关停组件，但是用户的消费业务逻辑还没有执行完毕，期望的是用户指定的消费业务逻辑执行完毕后再退出。

WBB在设计时充分考虑了上面的这三种问题。

对于情况1，ExpireBatchHandler中设计了一个close标记位，当close标记位为false的时候所有的while循环都会退出执行，例如ExpireBatchHandler的部分实现如下：

    private volatile boolean closed;

    // ...

    public void run() {
        // 循环检查有无过期攒批
        loopHandleExpireBatch();
        // 提交未完成的批次
        submitIncompleteBatches();
    }

    // ...

    private void loopHandleExpireBatch() {
        while (!closed) {
            try {
                List<ProducerBatch> expiredBatches1 = producerBatchContainer.getExpiredBatches();
                List<ProducerBatch> expiredBatches2 = retryQueue.getExpiredBatches(1000);
                expiredBatches1.addAll(expiredBatches2);
                for (ProducerBatch b : expiredBatches1) {
                    ioThreadPool.submit(new SendProducerBatchTask(b, producerConfig, messageProcessFunction, retryQueue, successQueue, failureQueue));
                }
            } catch (Exception e) {
                LOGGER.error("loopHandleExpireBatch Exception, e=", e);
            }
        }
    }
    // ...

    public void close() {
        this.closed = true;
        interrupt();
    }

在上面的代码中当close方法调用之后，while循环会及时退出，然后会检测攒批容器中是否存在未完成的攒批，如果有未完成的攒批会直接提交攒批队列。

对于情况2 ，则在攒批容器中设计了一个AtomicInteger类型的appendsInProgress作为追加计数器，如下面的代码所示，这样在关闭的时候只需要不停检查计数器是否归零就可以判断出是否存在元素还正处于追加过程中。

     public void append(List<Message> items, String groupName) throws InterruptedException, ProducerException {
        appendsInProgress.incrementAndGet();
        try {
            if (closed) {
                throw new IllegalStateException("cannot append after the waitingBus container was closed");
            }
            doAppend(items, groupName);
        } finally {
            appendsInProgress.decrementAndGet();
        }
    }

对于情况3则比较简单，因为业务的消费逻辑都在单独的线程池中执行的，只需要执行线程池的优雅关闭策略即可。

    private long closeIOThreadPool(long timeoutMs) throws InterruptedException, ProducerException {
        long startMs = System.currentTimeMillis();
        ioThreadPool.shutdown();
        if (ioThreadPool.awaitTermination(timeoutMs, TimeUnit.MILLISECONDS)) {
            LOGGER.debug("The biz-logic ioThreadPool is terminated");
        } else {
            LOGGER.warn("The biz-logic ioThreadPool is not fully terminated");
            throw new ProducerException("the biz-logic ioThreadPool is not fully terminated");
        }
        long nowMs = System.currentTimeMillis();
        return Math.max(0, timeoutMs - nowMs + startMs);
    }

当然优雅关闭的这些策略能够实现的前提是组件不能被强制Kill，因为WBB本质不是分布式组件，仅仅是基于内存设计的， 所以WBB不能完全避免消息丢失问题，这个也是这个组件目前比较不完善的地方。

六、总结

本文重点介绍了一种攒批队列的设计与实现，在业务开发过程中如果遇到需要进行"延迟处理"的场景时可以考虑本文的部分设计思路和源码进行实现。

在实现WBB的过程中一个很深的感触就是一个很简单的需求往往需要注意的细节越多，攒批的功能很简单但是需要考虑内存占用、并发、线程安全、组件启停等方方面面，想要写好一个工具其实需要花费很多心思。

最后将最近一段时间比较喜欢的Charles Eames的一句名言送给各位读者"The details are not the details. They make the design"我翻译为 "细节成就设计"，在软件开发中要永远尽可能多的去考虑细节。

本文内容为原创内容，创作不易，如要转载请注明出处。