org.springframework.batch.item.redis.StreamItemReader Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of spring-batch-redis Show documentation
Spring Batch reader and writer implementations for Redis
The newest version!
package org.springframework.batch.item.redis;

import io.lettuce.core.Consumer;
import io.lettuce.core.RedisBusyException;
import io.lettuce.core.RedisClient;
import io.lettuce.core.StreamMessage;
import io.lettuce.core.XGroupCreateArgs;
import io.lettuce.core.XReadArgs;
import io.lettuce.core.XReadArgs.StreamOffset;
import io.lettuce.core.api.StatefulConnection;
import io.lettuce.core.api.sync.BaseRedisCommands;
import io.lettuce.core.api.sync.RedisStreamCommands;
import io.lettuce.core.cluster.RedisClusterClient;
import io.lettuce.core.codec.StringCodec;
import lombok.Setter;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import org.springframework.batch.item.ExecutionContext;
import org.springframework.batch.item.ItemStreamException;
import org.springframework.batch.item.redis.support.CommandBuilder;
import org.springframework.batch.item.redis.support.ConnectionPoolItemStream;
import org.springframework.batch.item.redis.support.PollableItemReader;
import org.springframework.util.Assert;
import org.springframework.util.ClassUtils;

import java.time.Duration;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

@Slf4j
public class StreamItemReader extends ConnectionPoolItemStream implements PollableItemReader> {

    private final Function, BaseRedisCommands> sync;
    private final Long count;
    private final Duration block;
    private final StreamOffset offset;
    private final String consumerGroup;
    private final String consumer;
    private final AckPolicy ackPolicy;
    private Iterator> iterator = Collections.emptyIterator();

    public StreamItemReader(Supplier> connectionSupplier, GenericObjectPoolConfig> poolConfig, Function, BaseRedisCommands> sync, Long count, Duration block, String consumerGroup, String consumer, StreamOffset offset, AckPolicy ackPolicy) {
        super(connectionSupplier, poolConfig);
        Assert.notNull(sync, "A command provider is required");
        this.sync = sync;
        this.count = count;
        this.block = block;
        this.consumerGroup = consumerGroup;
        this.consumer = consumer;
        this.offset = offset;
        this.ackPolicy = ackPolicy;
    }

    @SuppressWarnings("unchecked")
    @Override
    public synchronized void open(ExecutionContext executionContext) {
        super.open(executionContext);
        try (StatefulConnection connection = pool.borrowObject()) {
            RedisStreamCommands commands = (RedisStreamCommands) sync.apply(connection);
            XGroupCreateArgs args = XGroupCreateArgs.Builder.mkstream(true);
            try {
                commands.xgroupCreate(offset, consumerGroup, args);
            } catch (RedisBusyException e) {
                // Consumer Group name already exists, ignore
            }
        } catch (Exception e) {
            throw new ItemStreamException("Failed to initialize the reader", e);
        }
    }

    @Override
    public StreamMessage read() throws Exception {
        throw new IllegalAccessException("read() method is not supposed to be called");
    }

    @Override
    public StreamMessage poll(long timeout, TimeUnit unit) throws Exception {
        if (!iterator.hasNext()) {
            List> messages = readMessages(Duration.ofMillis(unit.toMillis(timeout)));
            if (messages == null || messages.isEmpty()) {
                return null;
            }
            iterator = messages.iterator();
        }
        return iterator.next();
    }

    @SuppressWarnings("unused")
    public List> readMessages() throws Exception {
        return readMessages(block);
    }

    @SuppressWarnings("unchecked")
    private List> readMessages(Duration block) throws Exception {
        XReadArgs args = XReadArgs.Builder.count(count);
        if (block != null) {
            args.block(block);
        }
        try (StatefulConnection connection = pool.borrowObject()) {
            RedisStreamCommands commands = (RedisStreamCommands) sync.apply(connection);
            List> messages = commands.xreadgroup(Consumer.from(consumerGroup, consumer), args, StreamOffset.lastConsumed(offset.getName()));
            if (ackPolicy == AckPolicy.AUTO) {
                ack(messages);
            }
            return messages;
        }
    }

    @SuppressWarnings("unchecked")
    public void ack(List> messages) throws Exception {
        if (messages.isEmpty()) {
            return;
        }
        try (StatefulConnection connection = pool.borrowObject()) {
            RedisStreamCommands commands = (RedisStreamCommands) sync.apply(connection);
            Map>> streams = messages.stream().collect(Collectors.groupingBy(StreamMessage::getStream));
            for (Map.Entry>> entry : streams.entrySet()) {
                String[] messageIds = entry.getValue().stream().map(StreamMessage::getId).toArray(String[]::new);
                log.info("Ack'ing message ids: {}", Arrays.asList(messageIds));
                commands.xack(entry.getKey(), consumerGroup, messageIds);
            }
        }
    }

    public enum AckPolicy {
        AUTO, MANUAL
    }

    public static RedisClientStreamItemReaderBuilder client(RedisClient client) {
        return new RedisClientStreamItemReaderBuilder(client);
    }

    public static RedisClusterClientStreamItemReaderBuilder client(RedisClusterClient client) {
        return new RedisClusterClientStreamItemReaderBuilder(client);
    }

    public static class RedisClientStreamItemReaderBuilder {

        private final RedisClient client;

        public RedisClientStreamItemReaderBuilder(RedisClient client) {
            this.client = client;
        }

        public StreamItemReaderBuilder offset(StreamOffset offset) {
            return new StreamItemReaderBuilder(client, offset);
        }

    }

    public static class RedisClusterClientStreamItemReaderBuilder {

        private final RedisClusterClient client;

        public RedisClusterClientStreamItemReaderBuilder(RedisClusterClient client) {
            this.client = client;
        }

        public StreamItemReaderBuilder offset(StreamOffset offset) {
            return new StreamItemReaderBuilder(client, offset);
        }

    }

    @Setter
    @Accessors(fluent = true)
    public static class StreamItemReaderBuilder extends CommandBuilder {

        public static final Duration DEFAULT_BLOCK = Duration.ofMillis(100);
        public static final long DEFAULT_COUNT = 50;
        public static final String DEFAULT_CONSUMER_GROUP = ClassUtils.getShortName(StreamItemReader.class);
        public static final String DEFAULT_CONSUMER = "consumer1";
        public static final AckPolicy DEFAULT_ACK_POLICY = AckPolicy.AUTO;

        private final StreamOffset offset;
        private Duration block = DEFAULT_BLOCK;
        private Long count = DEFAULT_COUNT;
        private String consumerGroup = DEFAULT_CONSUMER_GROUP;
        private String consumer = DEFAULT_CONSUMER;
        private AckPolicy ackPolicy = DEFAULT_ACK_POLICY;

        public StreamItemReaderBuilder(RedisClient client, StreamOffset offset) {
            super(client, StringCodec.UTF8);
            this.offset = offset;
        }

        public StreamItemReaderBuilder(RedisClusterClient client, StreamOffset offset) {
            super(client, StringCodec.UTF8);
            this.offset = offset;
        }

        public StreamItemReader build() {
            return new StreamItemReader(connectionSupplier, poolConfig, sync, count, block, consumerGroup, consumer, offset, ackPolicy);
        }
    }

}