Java中如何用Akka Stream Graph DSL实现指定批量流处理？

我之前刚好在Java Akka Streams项目里处理过类似的批量+缓冲区溢出触发的场景，给你分享一个可行的实现方案，完全基于Graph DSL和自定义流阶段来满足你的需求。

整体思路 #

你的需求核心是Worker2需要双条件触发处理：

1. 正常情况下，每攒够3个字符串就拼接映射为数字输出
2. 当Worker2的输入缓冲区达到最大容量时，立即处理当前所有缓存的消息，清空缓冲区后继续接收新消息

我们可以通过自定义GraphStage来实现这个双逻辑，再结合Graph DSL搭建完整的流拓扑。

完整代码实现 #

1. 整体Graph DSL拓扑搭建 #

首先搭建包含Broadcast、Worker1、Worker2的完整流结构：

import akka.NotUsed;
import akka.stream.ClosedShape;
import akka.stream.Graph;
import akka.stream.javadsl.Broadcast;
import akka.stream.javadsl.GraphDSL;
import akka.stream.javadsl.Sink;
import akka.stream.javadsl.Source;
import java.util.List;

public class AkkaStreamBatchDemo {

    public static void main(String[] args) {
        // 示例原始字符串流（你可以替换成自己的业务流）
        Source<String, NotUsed> source = Source.from(List.of("a", "b", "c", "a", "d", "e", "f", "g", "h"));

        // Worker1：实时替换'a'为'A'并输出
        Sink<String, NotUsed> sink1 = Sink.foreach(str -> System.out.println("Worker1 Output: " + str));

        // Worker2：接收处理后的批量结果并输出（模拟慢Sink）
        Sink<String, NotUsed> sink2 = Sink.foreach(numStr -> {
            try {
                Thread.sleep(1000); // 模拟Sink2的慢处理速度
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
            System.out.println("Worker2 Output: " + numStr);
        });

        // 构建Graph DSL
        Graph<ClosedShape, NotUsed> streamGraph = GraphDSL.create(builder -> {
            // 创建Broadcast节点，将流广播到两个Worker
            Broadcast<String> broadcast = builder.add(Broadcast.create(2));

            // Worker1分支：简单的字符串替换
            builder.from(broadcast.out(0))
                    .map(str -> str.replace('a', 'A'))
                    .to(sink1);

            // Worker2分支：自定义批量处理阶段 + 拼接映射
            builder.from(broadcast.out(1))
                    // 自定义阶段：每3个触发，或缓冲区满（这里设为5）时触发
                    .via(new BatchOnCountOrBufferFull<>(3, 5))
                    // 将字符串列表拼接并映射为数字（替换成你的业务逻辑）
                    .map(strList -> mapStringsToNumber(strList))
                    // 显式设置输入缓冲区大小（和自定义阶段的maxBufferSize一致）
                    .withAttributes(akka.stream.Attributes.inputBuffer(5, 5))
                    .to(sink2);

            return ClosedShape.getInstance();
        });

        // 启动流
        source.runWith(streamGraph, akka.actor.ActorSystem.create("BatchStreamDemo"));
    }

    // 示例：将字符串列表拼接后映射为数字（根据你的业务需求修改）
    private static String mapStringsToNumber(List<String> strList) {
        String joined = String.join("", strList);
        // 这里用哈希值模拟数字映射，你可以替换成实际的转数字逻辑
        return String.valueOf(joined.hashCode());
    }
}

2. 自定义批量处理GraphStage #

这个自定义阶段是核心，用来实现双条件触发的逻辑：

import akka.stream.Attributes;
import akka.stream.FanOutShape1;
import akka.stream.Inlet;
import akka.stream.Outlet;
import akka.stream.stage.GraphStage;
import akka.stream.stage.GraphStageLogic;
import akka.stream.stage.InHandler;
import akka.stream.stage.OutHandler;

import java.util.ArrayList;
import java.util.List;

public class BatchOnCountOrBufferFull<T> extends GraphStage<FanOutShape1<T, List<T>>> {
    private final int targetBatchSize; // 正常触发的批量大小（3）
    private final int maxBufferSize; // 缓冲区最大容量（X）

    private final Inlet<T> inlet = Inlet.create("BatchOnCountOrBufferFull.in");
    private final Outlet<List<T>> outlet = Outlet.create("BatchOnCountOrBufferFull.out", this::pull);
    private final FanOutShape1<T, List<T>> shape = new FanOutShape1<>(inlet, outlet);

    public BatchOnCountOrBufferFull(int targetBatchSize, int maxBufferSize) {
        this.targetBatchSize = targetBatchSize;
        this.maxBufferSize = maxBufferSize;
    }

    @Override
    public FanOutShape1<T, List<T>> shape() {
        return shape;
    }

    @Override
    public GraphStageLogic createLogic(Attributes inheritedAttributes) {
        return new GraphStageLogic(shape) {
            private final List<T> buffer = new ArrayList<>();

            {
                // 处理上游输入
                setHandler(inlet, new InHandler() {
                    @Override
                    public void onPush() throws Exception {
                        T element = grab(inlet);
                        buffer.add(element);

                        // 触发条件：要么攒够目标数量，要么输入缓冲区已满
                        boolean shouldEmit = buffer.size() >= targetBatchSize || isInputBufferFull();

                        if (shouldEmit) {
                            // 输出当前缓冲区的副本，避免并发问题
                            emit(outlet, new ArrayList<>(buffer));
                            buffer.clear();
                        } else {
                            // 继续拉取上游数据
                            pull(inlet);
                        }
                    }

                    @Override
                    public void onUpstreamFinish() throws Exception {
                        // 上游结束时，输出剩余的所有缓存数据
                        if (!buffer.isEmpty()) {
                            emit(outlet, new ArrayList<>(buffer));
                        }
                        completeStage();
                    }
                });

                // 处理下游拉取
                setHandler(outlet, new OutHandler() {
                    @Override
                    public void onPull() throws Exception {
                        if (!hasBeenPulled(inlet)) {
                            pull(inlet);
                        }
                    }
                });
            }

            // 判断输入缓冲区是否已满
            private boolean isInputBufferFull() {
                return inlet.hasBeenPulled() && !inlet.isClosed() && inlet.available() == 0;
            }
        };
    }
}

关键细节说明 #

1. 双触发逻辑：自定义阶段在onPush中同时检查两个条件，满足任一条件就输出缓存的消息并清空缓冲区
2. 缓冲区控制：通过Attributes.inputBuffer显式设置Worker2分支的输入缓冲区大小，确保和自定义阶段的maxBufferSize一致
3. 上游结束处理：在onUpstreamFinish中输出剩余缓存，避免丢失未处理的消息
4. 线程安全：输出时创建缓冲区的副本，避免并发修改问题

内容的提问来源于stack exchange，提问作者Flame_Phoenix