spring-boot 用一个监听器订阅多个 stream

spring-boot 用一个监听器订阅多个 stream。

在一个 监听器里订阅多个 stream，可以减少 java 程序的内存占用。spring-boot 版本：2.3.4.RELEASE, redis 版本: 5.x

import com.google.common.collect.Lists;
import lombok.Builder;
import lombok.Data;
import lombok.Singular;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.dao.QueryTimeoutException;
import org.springframework.data.redis.RedisConnectionFailureException;
import org.springframework.data.redis.connection.stream.*;
import org.springframework.data.redis.core.StreamOperations;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.util.Assert;

import java.util.List;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicBoolean;


@Slf4j
@Builder
@Data
public class MultiStreamContainer {

    private StringRedisTemplate stringRedisTemplate;
    private String group;
    private String custom;
    private Executor executor;

    @Singular
    private List<Stream> streams = Lists.newArrayList();

    private Listener listener;

    private ListenerRunnable runnable;

    private static void prepareStreamAndGroup(StreamOperations<String, ?, ?> ops, String stream, String group) {
        String status = "OK";
        try {
            StreamInfo.XInfoGroups groups = ops.groups(stream);
            if (groups.stream().noneMatch(xInfoGroup -> group.equals(xInfoGroup.groupName()))) {
                status = ops.createGroup(stream, group);
            }
        } catch (Exception exception) {
            RecordId initialRecord = ops.add(ObjectRecord.create(stream, "Initial Record"));
            Assert.notNull(initialRecord, "Cannot initialize stream with key '" + stream + "'");
            status = ops.createGroup(stream, ReadOffset.from(initialRecord), group);
        } finally {
            Assert.isTrue("OK".equals(status), "Cannot create group with name '" + group + "'");
        }
    }

    public void stop() {
        if (runnable != null) {
            runnable.running.set(false);
            runnable = null;
        }
    }

    public void start() {
        stop();
        List<Stream> streamList = streams != null
                ? Lists.newArrayList(streams) : Lists.newArrayList();
        runnable = ListenerRunnable.builder()
                .streamList(streamList)
                .group(group)
                .custom(custom)
                .stringRedisTemplate(stringRedisTemplate)
                .listener(listener)
                .build();
        executor.execute(runnable);
    }

    public interface Listener {
        void onMessage(MapRecord message);
    }

    @Builder
    @Data
    public static class ListenerRunnable implements Runnable {
        private final AtomicBoolean running = new AtomicBoolean(true);
        private final List<Stream> streamList;
        private final String group;
        private final String custom;
        private final StringRedisTemplate stringRedisTemplate;
        private final Listener listener;

        @Override
        public void run() {
            while (running.get()) {
                if (streamList.isEmpty()) {
                    return;
                }

                List<StreamOffset<String>> offsetList = Lists.newArrayList();
                for (Stream readRecord : streamList) {
                    offsetList.add(StreamOffset.create(readRecord.getStream(), ReadOffset.from(readRecord.getId())));
                }

                Consumer consumer = Consumer.from(group, custom);
                StreamReadOptions options = StreamReadOptions.empty()
                        .count(1)
//                            .block(Duration.ofSeconds(2))
//                            .autoAcknowledge()
                        ;
                StreamOffset<String>[] offsets = new StreamOffset[offsetList.size()];
                for (int i = 0; i < offsetList.size(); i++) {
                    offsets[i] = offsetList.get(i);
                    prepareStreamAndGroup(stringRedisTemplate.opsForStream(), offsets[i].getKey(), group);
                }

                List<MapRecord<String, Object, Object>> values = null;
                try {
                    values = stringRedisTemplate.opsForStream().read(consumer, options, offsets);
                } catch (QueryTimeoutException exception) {
                    log.debug("XREADGROUP timeout: {}", exception.getMessage());
                } catch (RedisConnectionFailureException exception) {
                    if (StringUtils.isNotBlank(exception.getMessage()) && exception.getMessage().contains("Pool not open")) {
                        log.debug("XREADGROUP failure : pool not open");
                    } else {
                        log.error("XREADGROUP failure :", exception);
                    }
                }

                if (values != null && !values.isEmpty()) {
                    for (MapRecord val : values) {
                        listener.onMessage(val);
                    }
                }
            }
        }
    }

    @Data
    public static class Stream {
        private final String stream;

        private final String id;

        public static Stream latest(String stream) {
            return new Stream(stream, "$");
        }

        public static Stream lastConsumed(String stream) {
            return new Stream(stream, ">");
        }
    }
}

用法如下：

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.connection.stream.MapRecord;
import org.springframework.data.redis.connection.stream.RecordId;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.util.Map;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;

@Slf4j
@Service
public class NotifyTopicHandler {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    private MultiStreamContainer container = null;

    @PreDestroy
    public void des() {
        if (container != null) {
            container.stop();
        }
    }

    @PostConstruct
    public void d() {
        String stream1 = "NOTIFY:1";
        String stream2 = "NOTIFY:2";

        String group = "my-group";
        String consumer = "consumer-1";
        Executor executor = Executors.newFixedThreadPool(16);

        container = MultiStreamContainer.builder()
                .stringRedisTemplate(stringRedisTemplate)
                .custom(consumer)
                .group(group)
                .stream(MultiStreamContainer.Stream.lastConsumed(stream1))
                .stream(MultiStreamContainer.Stream.lastConsumed(stream2))
                .executor(executor)
                .listener(new MultiStreamContainer.Listener() {
                    @Override
                    public void onMessage(MapRecord message) {
                        RecordId id = message.getId();
                        Map messageValue = (Map) message.getValue();
                        log.info("stream:{} message:{}, id:{}", message.getStream(), messageValue, id);
                        stringRedisTemplate.convertAndSend("USER:xxx", JSON.toJSONString(messageValue));
                    }
                })
                .build();
        container.start();
    }
}