diff --git a/src/main/java/org/springframework/data/redis/core/ConvertingCursor.java b/src/main/java/org/springframework/data/redis/core/ConvertingCursor.java index e37562687..ababeba1e 100644 --- a/src/main/java/org/springframework/data/redis/core/ConvertingCursor.java +++ b/src/main/java/org/springframework/data/redis/core/ConvertingCursor.java @@ -122,4 +122,12 @@ public class ConvertingCursor implements Cursor { return delegate.getPosition(); } + /* + * (non-Javadoc) + * @see org.springframework.data.redis.core.Cursor#limit(long) + */ + @Override + public Cursor limit(long count) { + return new ConvertingCursor<>(delegate.limit(count), converter); + } } diff --git a/src/main/java/org/springframework/data/redis/core/Cursor.java b/src/main/java/org/springframework/data/redis/core/Cursor.java index bd0ac13d6..dc36a86cf 100644 --- a/src/main/java/org/springframework/data/redis/core/Cursor.java +++ b/src/main/java/org/springframework/data/redis/core/Cursor.java @@ -16,14 +16,19 @@ package org.springframework.data.redis.core; import java.io.Closeable; -import java.util.Iterator; + +import org.springframework.data.redis.util.BoundedIterator; /** + * Cursor abstraction to scan over the keyspace or elements within a data structure using a variant of a {@code SCAN} + * command. + * * @author Christoph Strobl + * @author Mark Paluch * @param * @since 1.4 */ -public interface Cursor extends Iterator, Closeable { +public interface Cursor extends BoundedIterator, Closeable { /** * Get the reference cursor.
@@ -50,4 +55,14 @@ public interface Cursor extends Iterator, Closeable { */ long getPosition(); + /** + * Limit the maximum number of elements to be returned from this cursor. The returned cursor object can be used to + * iterate over the remaining items and to {@link #close() release} associated resources. The returned cursor is not + * attached to the state of {@code this} cursor and this object should be no longer used. + * + * @return a new {@link Cursor} with detached iteration state. + * @since 2.5 + */ + @Override + Cursor limit(long count); } diff --git a/src/main/java/org/springframework/data/redis/core/ScanCursor.java b/src/main/java/org/springframework/data/redis/core/ScanCursor.java index 5e49de95a..4693bfe00 100644 --- a/src/main/java/org/springframework/data/redis/core/ScanCursor.java +++ b/src/main/java/org/springframework/data/redis/core/ScanCursor.java @@ -22,6 +22,7 @@ import java.util.NoSuchElementException; import org.springframework.dao.InvalidDataAccessApiUsageException; import org.springframework.lang.Nullable; +import org.springframework.util.Assert; import org.springframework.util.CollectionUtils; /** @@ -43,6 +44,7 @@ public abstract class ScanCursor implements Cursor { private @Nullable Iterator delegate; private @Nullable final ScanOptions scanOptions; private long position; + private final long limit; /** * Crates new {@link ScanCursor} with {@code id=0} and {@link ScanOptions#NONE} @@ -81,6 +83,23 @@ public abstract class ScanCursor implements Cursor { this.cursorId = cursorId; this.state = CursorState.READY; this.delegate = Collections. emptyList().iterator(); + this.limit = -1; + } + + /** + * Crates a new {@link ScanCursor}. + * + * @param source + * @param limit + * @since 2.5 + */ + private ScanCursor(ScanCursor source, long limit) { + + this.scanOptions = source.scanOptions; + this.cursorId = source.cursorId; + this.state = source.state; + this.delegate = source.delegate; + this.limit = limit; } private void scan(long cursorId) { @@ -178,6 +197,10 @@ public abstract class ScanCursor implements Cursor { assertCursorIsOpen(); + if (limit != -1 && getPosition() > limit - 1) { + return false; + } + while (!delegate.hasNext() && !CursorState.FINISHED.equals(state)) { scan(cursorId); } @@ -283,10 +306,43 @@ public abstract class ScanCursor implements Cursor { return position; } + /* + * (non-Javadoc) + * @see org.springframework.data.redis.core.Cursor#limit(long) + */ + @Override + public ScanCursor limit(long count) { + + Assert.isTrue(count >= 0, "Count must be greater or equal to zero"); + + return new ScanCursorWrapper<>(this, count); + } + /** * @author Thomas Darimont */ enum CursorState { READY, OPEN, FINISHED, CLOSED; } + + /** + * Wrapper for a concrete {@link ScanCursor} forwarding {@link #doScan(long, ScanOptions)}. + * + * @param + * @since 2.5 + */ + private static class ScanCursorWrapper extends ScanCursor { + + private final ScanCursor delegate; + + public ScanCursorWrapper(ScanCursor delegate, long limit) { + super(delegate, limit); + this.delegate = delegate; + } + + @Override + protected ScanIteration doScan(long cursorId, ScanOptions options) { + return delegate.doScan(cursorId, options); + } + } } diff --git a/src/main/java/org/springframework/data/redis/util/BoundedIterator.java b/src/main/java/org/springframework/data/redis/util/BoundedIterator.java new file mode 100644 index 000000000..0540befba --- /dev/null +++ b/src/main/java/org/springframework/data/redis/util/BoundedIterator.java @@ -0,0 +1,36 @@ +/* + * Copyright 2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.data.redis.util; + +import java.util.Iterator; + +/** + * Extension to {@link Iterator} that can be {@link #limit(long) limited} to a maximum number of items. + * + * @author Mark Paluch + * @since 2.5 + */ +public interface BoundedIterator extends Iterator { + + /** + * Limit the maximum number of elements to return. The limit is only applied to the returned instance and not applied + * to {@code this} iterator. + * + * @param count the maximum number of elements of iterator to return. Must be greater or equal to zero. + * @return a new instance of {@link BoundedIterator} with {@code count} applied. + */ + BoundedIterator limit(long count); +} diff --git a/src/test/java/org/springframework/data/redis/core/ScanCursorUnitTests.java b/src/test/java/org/springframework/data/redis/core/ScanCursorUnitTests.java index d722d3539..058797f32 100644 --- a/src/test/java/org/springframework/data/redis/core/ScanCursorUnitTests.java +++ b/src/test/java/org/springframework/data/redis/core/ScanCursorUnitTests.java @@ -17,7 +17,6 @@ package org.springframework.data.redis.core; import static org.assertj.core.api.Assertions.*; -import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -32,7 +31,10 @@ import org.junit.jupiter.api.Test; import org.springframework.dao.InvalidDataAccessApiUsageException; /** + * Unit tests for {@link ScanCursor}. + * * @author Christoph Strobl + * @author Mark Paluch */ class ScanCursorUnitTests { @@ -198,6 +200,43 @@ class ScanCursorUnitTests { assertThat(cursor.getCursorId()).isEqualTo(0L); } + @Test // GH-1575 + void limitShouldApplyLimitation() { + + LinkedList> values = new LinkedList<>(); + values.add(createIteration(1, "spring")); + values.add(createIteration(2, "data")); + values.add(createIteration(3, "redis")); + values.add(createIteration(0)); + Cursor cursor = initCursor(values).limit(2); + + List result = new ArrayList<>(); + while (cursor.hasNext()) { + result.add(cursor.next()); + } + + assertThat(result).hasSize(2).contains("spring", "data"); + } + + @Test // GH-1575 + void limitShouldNotLimitOriginalCursor() { + + LinkedList> values = new LinkedList<>(); + values.add(createIteration(1, "spring")); + values.add(createIteration(2, "data")); + values.add(createIteration(3, "redis")); + values.add(createIteration(0)); + Cursor cursor = initCursor(values); + cursor.limit(1); + + List result = new ArrayList<>(); + while (cursor.hasNext()) { + result.add(cursor.next()); + } + + assertThat(result).hasSize(3); + } + private CapturingCursorDummy initCursor(Queue> values) { CapturingCursorDummy cursor = new CapturingCursorDummy(values); cursor.open();