From 5548b8e21619cf5be1509e93edcb46445c0c5bfc Mon Sep 17 00:00:00 2001 From: sanghun Date: Tue, 14 Jul 2026 01:18:40 +0900 Subject: [PATCH 1/3] Add node-source association to ScanCursor #3287 SCAN cursors are node-local state: a continuation request must be routed to the node that issued the cursor. ScanCursor now carries an optional source RedisURI so drivers can associate a cursor with its issuing node. The attribute is package-private and exposed across packages through ScanCursorAccessor (following the OrderingReadFromAccessor pattern), leaving the public ScanCursor API unchanged. The immutable INITIAL/FINISHED cursors reject setSource. --- src/main/java/io/lettuce/core/ScanCursor.java | 29 +++++++++++++ .../io/lettuce/core/ScanCursorAccessor.java | 41 +++++++++++++++++++ .../io/lettuce/core/ScanCursorUnitTests.java | 7 ++++ 3 files changed, 77 insertions(+) create mode 100644 src/main/java/io/lettuce/core/ScanCursorAccessor.java diff --git a/src/main/java/io/lettuce/core/ScanCursor.java b/src/main/java/io/lettuce/core/ScanCursor.java index dfb71fb87b..146dbb7aa7 100644 --- a/src/main/java/io/lettuce/core/ScanCursor.java +++ b/src/main/java/io/lettuce/core/ScanCursor.java @@ -6,6 +6,7 @@ * Generic Cursor data structure. * * @author Mark Paluch + * @author Sanghun Lee * @since 3.0 */ public class ScanCursor { @@ -24,6 +25,8 @@ public class ScanCursor { private boolean finished; + private RedisURI source; + /** * Creates a new {@link ScanCursor}. */ @@ -72,6 +75,27 @@ public void setFinished(boolean finished) { this.finished = finished; } + /** + * Returns the {@link RedisURI} of the node that issued this cursor, if known, or {@code null}. Scan cursors are node-local + * state: continuation requests must be routed to the node that served the initial request. This hint is maintained by the + * driver (e.g. Master/Replica connections) and is not sent to Redis. Internal API, accessed across packages via + * {@link ScanCursorAccessor}. + * + * @return the node that issued this cursor or {@code null} if unknown. + */ + RedisURI getSource() { + return source; + } + + /** + * Associate this cursor with the node that issued it. Internal API, set by the driver via {@link ScanCursorAccessor}. + * + * @param source the node that issued this cursor, may be {@code null}. + */ + void setSource(RedisURI source) { + this.source = source; + } + /** * Creates a Scan-Cursor reference. * @@ -100,6 +124,11 @@ public void setFinished(boolean finished) { throw new UnsupportedOperationException("setFinished not supported on " + getClass().getSimpleName()); } + @Override + void setSource(RedisURI source) { + throw new UnsupportedOperationException("setSource not supported on " + getClass().getSimpleName()); + } + } } diff --git a/src/main/java/io/lettuce/core/ScanCursorAccessor.java b/src/main/java/io/lettuce/core/ScanCursorAccessor.java new file mode 100644 index 0000000000..27ec702493 --- /dev/null +++ b/src/main/java/io/lettuce/core/ScanCursorAccessor.java @@ -0,0 +1,41 @@ +package io.lettuce.core; + +/** + * Accessor for the source-node association of a {@link ScanCursor}. Internal utility class that lets the driver (e.g. + * Master/Replica connection routing) associate a scan cursor with the node that issued it across package boundaries, without + * widening the public {@link ScanCursor} API. + * + * @author Sanghun Lee + * @since 7.7 + */ +public abstract class ScanCursorAccessor { + + /** + * Utility constructor. + */ + private ScanCursorAccessor() { + } + + /** + * Returns the {@link RedisURI} of the node that issued {@code cursor}, if known. + * + * @param cursor the scan cursor. + * @return the node that issued the cursor or {@code null} if unknown. + * @since 7.7 + */ + public static RedisURI getSource(ScanCursor cursor) { + return cursor.getSource(); + } + + /** + * Associates {@code cursor} with the node that issued it. + * + * @param cursor the scan cursor. + * @param source the node that issued the cursor, may be {@code null}. + * @since 7.7 + */ + public static void setSource(ScanCursor cursor, RedisURI source) { + cursor.setSource(source); + } + +} diff --git a/src/test/java/io/lettuce/core/ScanCursorUnitTests.java b/src/test/java/io/lettuce/core/ScanCursorUnitTests.java index 0f6345abe8..809cc98347 100644 --- a/src/test/java/io/lettuce/core/ScanCursorUnitTests.java +++ b/src/test/java/io/lettuce/core/ScanCursorUnitTests.java @@ -9,6 +9,7 @@ /** * @author Mark Paluch + * @author Sanghun Lee */ @Tag(UNIT_TEST) class ScanCursorUnitTests { @@ -30,4 +31,10 @@ void setFinishedOnImmutableInstance() { assertThatThrownBy(() -> ScanCursor.INITIAL.setFinished(false)).isInstanceOf(UnsupportedOperationException.class); } + @Test + void setSourceOnImmutableInstance() { + assertThatThrownBy(() -> ScanCursor.INITIAL.setSource(RedisURI.create("localhost", 6379))) + .isInstanceOf(UnsupportedOperationException.class); + } + } From 2aed4720461b227148220aa93e809371d3215b19 Mon Sep 17 00:00:00 2001 From: sanghun Date: Tue, 14 Jul 2026 01:18:40 +0900 Subject: [PATCH 2/3] Propagate scan cursor source hint in RedisCommandBuilder #3287 Scan command builders now copy the source-node hint from the input cursor to the cursor that carries the response, so reusing the returned cursor object for continuation calls retains the association (as ScanIterator and ScanStream already do). Applied uniformly across all SCAN/HSCAN/SSCAN/ZSCAN builders including streaming variants. --- .../io/lettuce/core/RedisCommandBuilder.java | 23 +++++++++++++ .../core/RedisCommandBuilderUnitTests.java | 33 +++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/src/main/java/io/lettuce/core/RedisCommandBuilder.java b/src/main/java/io/lettuce/core/RedisCommandBuilder.java index b6ef02d68d..4bab8383a5 100644 --- a/src/main/java/io/lettuce/core/RedisCommandBuilder.java +++ b/src/main/java/io/lettuce/core/RedisCommandBuilder.java @@ -63,6 +63,7 @@ * @author Ali Takavci * @author Seonghwan Lee * @author dae won + * @author Sanghun Lee */ @SuppressWarnings({ "unchecked", "varargs" }) class RedisCommandBuilder extends BaseRedisCommandBuilder { @@ -1739,6 +1740,7 @@ Command> hscan(K key, ScanCursor scanCursor, ScanArgs scanArgs(scanCursor, scanArgs, args); MapScanOutput output = new MapScanOutput<>(codec); + associateSourceHint(output, scanCursor); return createCommand(HSCAN, output, args); } @@ -1753,6 +1755,7 @@ Command> hscanNovalues(K key, ScanCursor scanCursor, Scan args.add(NOVALUES); KeyScanOutput output = new KeyScanOutput<>(codec); + associateSourceHint(output, scanCursor); return createCommand(HSCAN, output, args); } @@ -1809,6 +1812,7 @@ Command hscanStreaming(KeyValueStreamingChannel ch scanArgs(scanCursor, scanArgs, args); KeyValueScanStreamingOutput output = new KeyValueScanStreamingOutput<>(codec, channel); + associateSourceHint(output, scanCursor); return createCommand(HSCAN, output, args); } @@ -1825,6 +1829,7 @@ Command hscanNoValuesStreaming(KeyStreamingChannel ch args.add(NOVALUES); KeyScanStreamingOutput output = new KeyScanStreamingOutput<>(codec, channel); + associateSourceHint(output, scanCursor); return createCommand(HSCAN, output, args); } @@ -2644,6 +2649,7 @@ Command> scan(ScanCursor scanCursor, ScanArgs scanArgs) { scanArgs(scanCursor, scanArgs, args); KeyScanOutput output = new KeyScanOutput<>(codec); + associateSourceHint(output, scanCursor); return createCommand(SCAN, output, args); } @@ -2658,6 +2664,18 @@ protected void scanArgs(ScanCursor scanCursor, ScanArgs scanArgs, CommandArgs output, ScanCursor scanCursor) { + + if (scanCursor.getSource() != null) { + output.get().setSource(scanCursor.getSource()); + } + } + Command scanStreaming(KeyStreamingChannel channel) { notNull(channel); LettuceAssert.notNull(channel, "KeyStreamingChannel " + MUST_NOT_BE_NULL); @@ -2687,6 +2705,7 @@ Command scanStreaming(KeyStreamingChannel channel, Sc scanArgs(scanCursor, scanArgs, args); KeyScanStreamingOutput output = new KeyScanStreamingOutput<>(codec, channel); + associateSourceHint(output, scanCursor); return createCommand(SCAN, output, args); } @@ -3088,6 +3107,7 @@ Command> sscan(K key, ScanCursor scanCursor, ScanArgs s scanArgs(scanCursor, scanArgs, args); ValueScanOutput output = new ValueScanOutput<>(codec); + associateSourceHint(output, scanCursor); return createCommand(SSCAN, output, args); } @@ -3123,6 +3143,7 @@ Command sscanStreaming(ValueStreamingChannel channel, scanArgs(scanCursor, scanArgs, args); ValueScanStreamingOutput output = new ValueScanStreamingOutput<>(codec, channel); + associateSourceHint(output, scanCursor); return createCommand(SSCAN, output, args); } @@ -4658,6 +4679,7 @@ Command> zscan(K key, ScanCursor scanCursor, Scan scanArgs(scanCursor, scanArgs, args); ScoredValueScanOutput output = new ScoredValueScanOutput<>(codec); + associateSourceHint(output, scanCursor); return createCommand(ZSCAN, output, args); } @@ -4693,6 +4715,7 @@ Command zscanStreaming(ScoredValueStreamingChannel ch scanArgs(scanCursor, scanArgs, args); ScoredValueScanStreamingOutput output = new ScoredValueScanStreamingOutput<>(codec, channel); + associateSourceHint(output, scanCursor); return createCommand(ZSCAN, output, args); } diff --git a/src/test/java/io/lettuce/core/RedisCommandBuilderUnitTests.java b/src/test/java/io/lettuce/core/RedisCommandBuilderUnitTests.java index 8c0d2b9b28..d0f56e9944 100644 --- a/src/test/java/io/lettuce/core/RedisCommandBuilderUnitTests.java +++ b/src/test/java/io/lettuce/core/RedisCommandBuilderUnitTests.java @@ -26,6 +26,7 @@ * * @author Mark Paluch * @author dae won + * @author Sanghun Lee */ @Tag(UNIT_TEST) class RedisCommandBuilderUnitTests { @@ -837,4 +838,36 @@ void shouldCorrectlyConstructClientNoTouchOff() { .isEqualTo("*3\r\n" + "$6\r\n" + "CLIENT\r\n" + "$8\r\n" + "NO-TOUCH\r\n" + "$3\r\n" + "OFF\r\n"); } + @Test + void scanShouldPropagateCursorSourceToResultCursor() { + + RedisURI source = RedisURI.create("localhost", 6482); + + ScanCursor scanCursor = ScanCursor.of("42"); + scanCursor.setSource(source); + + assertThat(sut.scan(scanCursor, null).getOutput().get().getSource()).isSameAs(source); + assertThat(sut.hscan(MY_KEY, scanCursor, null).getOutput().get().getSource()).isSameAs(source); + assertThat(sut.hscanNovalues(MY_KEY, scanCursor, null).getOutput().get().getSource()).isSameAs(source); + assertThat(sut.sscan(MY_KEY, scanCursor, null).getOutput().get().getSource()).isSameAs(source); + assertThat(sut.zscan(MY_KEY, scanCursor, null).getOutput().get().getSource()).isSameAs(source); + assertThat(sut.scanStreaming(k -> { + }, scanCursor, null).getOutput().get().getSource()).isSameAs(source); + assertThat(sut.hscanStreaming((k, v) -> { + }, MY_KEY, scanCursor, null).getOutput().get().getSource()).isSameAs(source); + assertThat(sut.hscanNoValuesStreaming(k -> { + }, MY_KEY, scanCursor, null).getOutput().get().getSource()).isSameAs(source); + assertThat(sut.sscanStreaming(v -> { + }, MY_KEY, scanCursor, null).getOutput().get().getSource()).isSameAs(source); + assertThat(sut.zscanStreaming(sv -> { + }, MY_KEY, scanCursor, null).getOutput().get().getSource()).isSameAs(source); + } + + @Test + void scanWithoutCursorSourceShouldLeaveResultCursorSourceEmpty() { + + assertThat(sut.scan(ScanCursor.of("42"), null).getOutput().get().getSource()).isNull(); + assertThat(sut.scan(ScanCursor.INITIAL, null).getOutput().get().getSource()).isNull(); + } + } From 7df3592d7b973839c945e88858ce0978d9389953 Mon Sep 17 00:00:00 2001 From: sanghun Date: Tue, 14 Jul 2026 01:18:40 +0900 Subject: [PATCH 3/3] Pin Master/Replica scan continuations to the issuing node #3287 MasterReplicaChannelWriter now recognizes scan commands: continuation requests (a cursor carrying a source) are pinned to the issuing node, while the node selected for an initial request is stamped onto the returned cursor. MasterReplicaConnectionProvider gains a node-reporting overload and getPinnedConnectionAsync, which fails fast with a descriptive RedisException when the issuing node is no longer part of the topology instead of silently routing elsewhere. Read load-balancing for non-scan reads and for initial scan selection is unchanged. Adds unit coverage plus a {SCAN,HSCAN,SSCAN,ZSCAN} x {sync,async,reactive} integration matrix under ReadFrom.ANY with a ReadFrom.UPSTREAM control. --- .../MasterReplicaChannelWriter.java | 106 ++++- .../MasterReplicaConnectionProvider.java | 93 ++++- .../MasterReplicaChannelWriterUnitTests.java | 148 +++++++ ...terReplicaConnectionProviderUnitTests.java | 67 ++++ ...sterReplicaStickyScanIntegrationTests.java | 364 ++++++++++++++++++ 5 files changed, 765 insertions(+), 13 deletions(-) create mode 100644 src/test/java/io/lettuce/core/masterreplica/MasterReplicaStickyScanIntegrationTests.java diff --git a/src/main/java/io/lettuce/core/masterreplica/MasterReplicaChannelWriter.java b/src/main/java/io/lettuce/core/masterreplica/MasterReplicaChannelWriter.java index 7bf36b4958..32824b313b 100644 --- a/src/main/java/io/lettuce/core/masterreplica/MasterReplicaChannelWriter.java +++ b/src/main/java/io/lettuce/core/masterreplica/MasterReplicaChannelWriter.java @@ -19,15 +19,23 @@ */ package io.lettuce.core.masterreplica; +import java.util.ArrayList; import java.util.Collection; +import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.function.Consumer; import io.lettuce.core.ClientOptions; import io.lettuce.core.ReadFrom; import io.lettuce.core.RedisChannelWriter; import io.lettuce.core.RedisException; +import io.lettuce.core.ScanCursor; +import io.lettuce.core.ScanCursorAccessor; import io.lettuce.core.api.StatefulRedisConnection; import io.lettuce.core.internal.LettuceAssert; +import io.lettuce.core.models.role.RedisNodeDescription; +import io.lettuce.core.output.CommandOutput; +import io.lettuce.core.protocol.CommandType; import io.lettuce.core.protocol.ConnectionFacade; import io.lettuce.core.protocol.ConnectionIntent; import io.lettuce.core.protocol.ProtocolKeyword; @@ -39,6 +47,7 @@ * * @author Mark Paluch * @author Jim Brunner + * @author Sanghun Lee */ class MasterReplicaChannelWriter implements RedisChannelWriter { @@ -78,8 +87,8 @@ public RedisCommand write(RedisCommand command) { ConnectionIntent connectionIntent = inTransaction ? ConnectionIntent.WRITE : (readOnlyCommands.isReadOnly(command) ? ConnectionIntent.READ : ConnectionIntent.WRITE); - CompletableFuture> future = (CompletableFuture) masterReplicaConnectionProvider - .getConnectionAsync(connectionIntent); + CompletableFuture> future = (CompletableFuture) getConnectionAsync(command, + connectionIntent); if (isEndTransaction(command.getType())) { inTransaction = false; @@ -94,6 +103,90 @@ public RedisCommand write(RedisCommand command) { return command; } + /** + * Obtain a connection to run {@code command} on. Scan commands ({@code SCAN}, {@code HSCAN}, {@code SSCAN}, {@code ZSCAN}) + * receive node-affine routing: scan cursors are node-local state, so continuation requests are pinned to the node that + * issued the cursor, and the node selected for the initial request is associated with the cursor that is returned to the + * caller, see {@link ScanCursor#getSource()}. + */ + private CompletableFuture> getConnectionAsync(RedisCommand command, + ConnectionIntent connectionIntent) { + + if (connectionIntent == ConnectionIntent.READ) { + + ScanCursor cursor = getScanCursor(command); + + if (cursor != null) { + + if (ScanCursorAccessor.getSource(cursor) != null) { + return masterReplicaConnectionProvider.getPinnedConnectionAsync(ScanCursorAccessor.getSource(cursor)); + } + + // Node selection happens-before the command is dispatched, so the association is visible to the caller + // by the time the response (carrying this cursor) completes. + return masterReplicaConnectionProvider.getConnectionAsync(ConnectionIntent.READ, + node -> ScanCursorAccessor.setSource(cursor, node.getUri())); + } + } + + return masterReplicaConnectionProvider.getConnectionAsync(connectionIntent); + } + + /** + * Batched commands are dispatched over a single connection: associate the selected node with each scan cursor in the batch + * that does not carry a source yet, so individual continuations of these scans are pinned correctly. Pinning of batched + * scan continuations is not supported (the standard sync/async/reactive APIs dispatch scans as single commands); a + * batched continuation is routed like any other read command. + * + * @return a stamper for the batch, or {@code null} if the batch contains no scan command to associate. + */ + private static Consumer scanCursorStamper( + Collection> commands) { + + List cursors = null; + + for (RedisCommand command : commands) { + + ScanCursor cursor = getScanCursor(command); + + if (cursor != null && ScanCursorAccessor.getSource(cursor) == null) { + + if (cursors == null) { + cursors = new ArrayList<>(2); + } + + cursors.add(cursor); + } + } + + if (cursors == null) { + return null; + } + + List cursorsToStamp = cursors; + return node -> cursorsToStamp.forEach(cursor -> ScanCursorAccessor.setSource(cursor, node.getUri())); + } + + /** + * @return the {@link ScanCursor} that will carry the response of a scan command, or {@code null} if {@code command} is not + * a scan command. + */ + private static ScanCursor getScanCursor(RedisCommand command) { + + ProtocolKeyword type = command.getType(); + + if (type != CommandType.SCAN && type != CommandType.HSCAN && type != CommandType.SSCAN && type != CommandType.ZSCAN) { + return null; + } + + CommandOutput output = command.getOutput(); + // The cursor is created eagerly in the ScanOutput constructor, so get() returns it before the response arrives; + // it is the same object that later receives the response cursor and is handed back to the caller. + Object value = output != null ? output.get() : null; + + return value instanceof ScanCursor ? (ScanCursor) value : null; + } + @SuppressWarnings("unchecked") private static void writeCommand(RedisCommand command, StatefulRedisConnection connection, Throwable throwable) { @@ -131,8 +224,13 @@ private static void writeCommand(RedisCommand command, StatefulR // Currently: Retain order ConnectionIntent connectionIntent = inTransaction ? ConnectionIntent.WRITE : getIntent(commands); - CompletableFuture> future = (CompletableFuture) masterReplicaConnectionProvider - .getConnectionAsync(connectionIntent); + Consumer scanCursorStamper = connectionIntent == ConnectionIntent.READ + ? scanCursorStamper(commands) + : null; + + CompletableFuture> future = (CompletableFuture) (scanCursorStamper != null + ? masterReplicaConnectionProvider.getConnectionAsync(connectionIntent, scanCursorStamper) + : masterReplicaConnectionProvider.getConnectionAsync(connectionIntent)); for (RedisCommand command : commands) { if (isEndTransaction(command.getType())) { diff --git a/src/main/java/io/lettuce/core/masterreplica/MasterReplicaConnectionProvider.java b/src/main/java/io/lettuce/core/masterreplica/MasterReplicaConnectionProvider.java index 4fb8002638..9a5ad6727b 100644 --- a/src/main/java/io/lettuce/core/masterreplica/MasterReplicaConnectionProvider.java +++ b/src/main/java/io/lettuce/core/masterreplica/MasterReplicaConnectionProvider.java @@ -15,10 +15,13 @@ import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; import java.util.function.Function; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.util.function.Tuple2; +import reactor.util.function.Tuples; import io.lettuce.core.ConnectionFuture; import io.lettuce.core.OrderingReadFromAccessor; import io.lettuce.core.ReadFrom; @@ -40,6 +43,7 @@ * Connection provider for master/replica setups. The connection provider * * @author Mark Paluch + * @author Sanghun Lee * @since 4.1 */ class MasterReplicaConnectionProvider { @@ -106,6 +110,20 @@ public StatefulRedisConnection getConnection(ConnectionIntent intent) { * @throws RedisException if the host is not part of the cluster */ public CompletableFuture> getConnectionAsync(ConnectionIntent intent) { + return getConnectionAsync(intent, null); + } + + /** + * Variant of {@link #getConnectionAsync(ConnectionIntent)} that reports the selected node to {@code selectedNodeListener}, + * so callers can associate node-local state (such as scan cursors) with the node serving the command. + * + * @param intent command intent + * @param selectedNodeListener callback notified with the selected node, may be {@code null}. + * @return the connection. + * @throws RedisException if the host is not part of the cluster + */ + CompletableFuture> getConnectionAsync(ConnectionIntent intent, + Consumer selectedNodeListener) { if (debugEnabled) { logger.debug("getConnectionAsync(" + intent + ")"); @@ -132,31 +150,88 @@ public Iterator iterator() { } if (selection.size() == 1) { - return getConnection(selection.get(0)); + RedisNodeDescription node = selection.get(0); + notifySelectedNode(selectedNodeListener, node); + return getConnection(node); } try { - Flux> connections = Flux.empty(); + Flux>> connections = Flux.empty(); for (RedisNodeDescription node : selection) { - connections = connections.concatWith(Mono.fromFuture(getConnection(node))); + connections = connections + .concatWith(Mono.fromFuture(getConnection(node)).map(connection -> Tuples.of(node, connection))); } + Mono>> selected; + if (OrderingReadFromAccessor.isOrderSensitive(readFrom)) { - return connections.filter(StatefulConnection::isOpen).next().switchIfEmpty(connections.next()).toFuture(); + selected = connections.filter(it -> it.getT2().isOpen()).next().switchIfEmpty(connections.next()); + } else { + selected = connections.filter(it -> it.getT2().isOpen()).collectList().filter(it -> !it.isEmpty()) + .map(it -> it.get(ThreadLocalRandom.current().nextInt(it.size()))) + .switchIfEmpty(connections.next()); } - return connections.filter(StatefulConnection::isOpen).collectList().filter(it -> !it.isEmpty()).map(it -> { - int index = ThreadLocalRandom.current().nextInt(it.size()); - return it.get(index); - }).switchIfEmpty(connections.next()).toFuture(); + // doOnNext runs before the returned future completes, so the selected node (possibly the switchIfEmpty + // fallback, which is still part of the known topology) is reported before the command is dispatched. + return selected.doOnNext(it -> notifySelectedNode(selectedNodeListener, it.getT1())).map(Tuple2::getT2) + .toFuture(); } catch (RuntimeException e) { throw Exceptions.bubble(e); } } - return getConnection(getMaster()); + RedisNodeDescription master = getMaster(); + notifySelectedNode(selectedNodeListener, master); + return getConnection(master); + } + + /** + * Retrieve a {@link StatefulRedisConnection} to the node identified by {@code source}, used to pin scan cursor continuation + * requests to the node that issued the cursor. Scan cursors are node-local state: applying them to another node yields + * undefined results, so if the source node is no longer part of the known topology this method fails with a descriptive + * {@link RedisException} instead of silently routing elsewhere; callers are expected to restart their scan. + * + * @param source the node that issued the scan cursor. + * @return the connection. + */ + CompletableFuture> getPinnedConnectionAsync(RedisURI source) { + + if (debugEnabled) { + logger.debug("getPinnedConnectionAsync(" + source + ")"); + } + + // Snapshot under the lock: a topology refresh mutates knownNodes in place and a torn read here would surface + // as a spurious scan abort. + List nodes; + + stateLock.lock(); + try { + nodes = new ArrayList<>(knownNodes); + } finally { + stateLock.unlock(); + } + + RedisNodeDescription node = findNodeByUri(nodes, source); + + if (node == null) { + CompletableFuture> failed = new CompletableFuture<>(); + failed.completeExceptionally(new RedisException(String.format( + "Cannot route scan continuation to %s: the node that issued the scan cursor is no longer available (Known nodes: %s). Restart the scan iteration", + source, nodes))); + return failed; + } + + return getConnection(node); + } + + private static void notifySelectedNode(Consumer listener, RedisNodeDescription node) { + + if (listener != null) { + listener.accept(node); + } } protected CompletableFuture> getConnection(RedisNodeDescription redisNodeDescription) { diff --git a/src/test/java/io/lettuce/core/masterreplica/MasterReplicaChannelWriterUnitTests.java b/src/test/java/io/lettuce/core/masterreplica/MasterReplicaChannelWriterUnitTests.java index 7acce94210..e2df8c2604 100644 --- a/src/test/java/io/lettuce/core/masterreplica/MasterReplicaChannelWriterUnitTests.java +++ b/src/test/java/io/lettuce/core/masterreplica/MasterReplicaChannelWriterUnitTests.java @@ -28,6 +28,7 @@ import java.util.Collections; import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.function.Consumer; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; @@ -38,9 +39,16 @@ import org.mockito.quality.Strictness; import io.lettuce.core.ClientOptions; +import io.lettuce.core.KeyScanCursor; +import io.lettuce.core.RedisException; +import io.lettuce.core.RedisURI; +import io.lettuce.core.ScanCursorAccessor; import io.lettuce.core.api.StatefulRedisConnection; import io.lettuce.core.codec.StringCodec; +import io.lettuce.core.models.role.RedisNodeDescription; +import io.lettuce.core.output.KeyScanOutput; import io.lettuce.core.output.StatusOutput; +import io.lettuce.core.protocol.AsyncCommand; import io.lettuce.core.protocol.Command; import io.lettuce.core.protocol.CommandType; import io.lettuce.core.protocol.ConnectionIntent; @@ -50,6 +58,7 @@ /** * @author Mark Paluch * @author Jim Brunner + * @author Sanghun Lee */ @Tag(UNIT_TEST) @ExtendWith(MockitoExtension.class) @@ -195,6 +204,145 @@ void shouldDeriveIntentFromCommandBatchTypeAfterDiscardedTransaction() { verify(connectionProvider).getConnectionAsync(ConnectionIntent.READ); } + @Test + void shouldAssociateSourceWithInitialScanCommand() { + + MasterReplicaChannelWriter writer = new MasterReplicaChannelWriter(connectionProvider, clientResources, clientOptions); + + RedisURI uri = RedisURI.create("localhost", 6482); + RedisNodeDescription node = mock(RedisNodeDescription.class); + when(node.getUri()).thenReturn(uri); + + when(connectionProvider.getConnectionAsync(eq(ConnectionIntent.READ), any())).thenAnswer(invocation -> { + Consumer listener = invocation.getArgument(1); + listener.accept(node); + return CompletableFuture.completedFuture(connection); + }); + + Command> scan = scanCommand(); + + writer.write(scan); + + assertThat(ScanCursorAccessor.getSource(scan.getOutput().get())).isEqualTo(uri); + verify(connectionProvider).getConnectionAsync(eq(ConnectionIntent.READ), any()); + verify(connectionProvider, never()).getPinnedConnectionAsync(any()); + } + + @Test + void shouldPinScanContinuationToSourceNode() { + + MasterReplicaChannelWriter writer = new MasterReplicaChannelWriter(connectionProvider, clientResources, clientOptions); + + RedisURI uri = RedisURI.create("localhost", 6482); + + when(connectionProvider.getPinnedConnectionAsync(uri)).thenReturn(CompletableFuture.completedFuture(connection)); + + Command> scan = scanCommand(); + ScanCursorAccessor.setSource(scan.getOutput().get(), uri); + + writer.write(scan); + + verify(connectionProvider).getPinnedConnectionAsync(uri); + verify(connectionProvider, never()).getConnectionAsync(any(ConnectionIntent.class)); + verify(connectionProvider, never()).getConnectionAsync(any(ConnectionIntent.class), any()); + } + + @Test + void shouldFailScanContinuationWhenSourceNodeIsGone() { + + MasterReplicaChannelWriter writer = new MasterReplicaChannelWriter(connectionProvider, clientResources, clientOptions); + + RedisURI uri = RedisURI.create("localhost", 6482); + + CompletableFuture> failed = new CompletableFuture<>(); + failed.completeExceptionally(new RedisException("Cannot route scan continuation")); + when(connectionProvider.getPinnedConnectionAsync(uri)).thenReturn(failed); + + Command> scan = scanCommand(); + ScanCursorAccessor.setSource(scan.getOutput().get(), uri); + AsyncCommand> asyncCommand = new AsyncCommand<>(scan); + + writer.write(asyncCommand); + + assertThat(asyncCommand.isCompletedExceptionally()).isTrue(); + assertThatThrownBy(asyncCommand::join).hasCauseInstanceOf(RedisException.class); + } + + @Test + void shouldUseDefaultReadRouteForNonScanReadCommands() { + + MasterReplicaChannelWriter writer = new MasterReplicaChannelWriter(connectionProvider, clientResources, clientOptions); + + when(connectionProvider.getConnectionAsync(any(ConnectionIntent.class))) + .thenReturn(CompletableFuture.completedFuture(connection)); + + writer.write(mockCommand(CommandType.GET)); + + verify(connectionProvider).getConnectionAsync(ConnectionIntent.READ); + verify(connectionProvider, never()).getConnectionAsync(any(ConnectionIntent.class), any()); + verify(connectionProvider, never()).getPinnedConnectionAsync(any()); + } + + @Test + void shouldBindScansInTransactionToMaster() { + + MasterReplicaChannelWriter writer = new MasterReplicaChannelWriter(connectionProvider, clientResources, clientOptions); + + when(connectionProvider.getConnectionAsync(any(ConnectionIntent.class))) + .thenReturn(CompletableFuture.completedFuture(connection)); + + writer.write(mockCommand(CommandType.MULTI)); + + Command> scan = scanCommand(); + writer.write(scan); + writer.write(mockCommand(CommandType.EXEC)); + + verify(connectionProvider, times(3)).getConnectionAsync(ConnectionIntent.WRITE); + verify(connectionProvider, never()).getConnectionAsync(any(ConnectionIntent.class), any()); + assertThat(ScanCursorAccessor.getSource(scan.getOutput().get())).isNull(); + } + + @Test + void shouldAssociateSourceWithScanCommandsInBatch() { + + MasterReplicaChannelWriter writer = new MasterReplicaChannelWriter(connectionProvider, clientResources, clientOptions); + + RedisURI uri = RedisURI.create("localhost", 6482); + RedisNodeDescription node = mock(RedisNodeDescription.class); + when(node.getUri()).thenReturn(uri); + + when(connectionProvider.getConnectionAsync(eq(ConnectionIntent.READ), any())).thenAnswer(invocation -> { + Consumer listener = invocation.getArgument(1); + listener.accept(node); + return CompletableFuture.completedFuture(connection); + }); + + Command> scan = scanCommand(); + + writer.write(Arrays.asList(scan, mockCommand(CommandType.GET))); + + assertThat(ScanCursorAccessor.getSource(scan.getOutput().get())).isEqualTo(uri); + verify(connectionProvider).getConnectionAsync(eq(ConnectionIntent.READ), any()); + } + + @Test + void shouldUseDefaultRouteForBatchWithoutScanCommands() { + + MasterReplicaChannelWriter writer = new MasterReplicaChannelWriter(connectionProvider, clientResources, clientOptions); + + when(connectionProvider.getConnectionAsync(any(ConnectionIntent.class))) + .thenReturn(CompletableFuture.completedFuture(connection)); + + writer.write(Collections.singletonList(mockCommand(CommandType.GET))); + + verify(connectionProvider).getConnectionAsync(ConnectionIntent.READ); + verify(connectionProvider, never()).getConnectionAsync(any(ConnectionIntent.class), any()); + } + + private static Command> scanCommand() { + return new Command<>(CommandType.SCAN, new KeyScanOutput<>(StringCodec.UTF8)); + } + private static Command mockCommand(CommandType multi) { return new Command<>(multi, new StatusOutput<>(StringCodec.UTF8)); } diff --git a/src/test/java/io/lettuce/core/masterreplica/MasterReplicaConnectionProviderUnitTests.java b/src/test/java/io/lettuce/core/masterreplica/MasterReplicaConnectionProviderUnitTests.java index 00ced21500..aff1e0fa26 100644 --- a/src/test/java/io/lettuce/core/masterreplica/MasterReplicaConnectionProviderUnitTests.java +++ b/src/test/java/io/lettuce/core/masterreplica/MasterReplicaConnectionProviderUnitTests.java @@ -8,6 +8,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Tag; @@ -22,15 +23,18 @@ import io.lettuce.core.ReadFrom; import io.lettuce.core.RedisChannelHandler; import io.lettuce.core.RedisClient; +import io.lettuce.core.RedisException; import io.lettuce.core.RedisURI; import io.lettuce.core.api.StatefulRedisConnection; import io.lettuce.core.api.sync.RedisCommands; import io.lettuce.core.codec.StringCodec; import io.lettuce.core.models.role.RedisInstance; +import io.lettuce.core.models.role.RedisNodeDescription; import io.lettuce.core.protocol.ConnectionIntent; /** * @author Mark Paluch + * @author Sanghun Lee */ @Tag(UNIT_TEST) @ExtendWith(MockitoExtension.class) @@ -90,4 +94,67 @@ void shouldUseDirectConnectionForSingleReadSelection() { verify(nodeConnectionMock, never()).isOpen(); } + @Test + void shouldReportSelectedNodeToListener() { + + when(clientMock.connectAsync(eq(StringCodec.UTF8), any())) + .thenReturn(ConnectionFuture.completed(null, nodeConnectionMock)); + + sut.setReadFrom(ReadFrom.ANY); + + AtomicReference selectedNode = new AtomicReference<>(); + + StatefulRedisConnection connection = sut.getConnectionAsync(ConnectionIntent.READ, selectedNode::set) + .join(); + + assertThat(connection).isSameAs(nodeConnectionMock); + assertThat(selectedNode.get()).isNotNull(); + assertThat(selectedNode.get().getUri()).isEqualTo(RedisURI.create("localhost", 1)); + } + + @Test + void shouldReportSelectedNodeForMultiNodeSelection() { + + sut.setKnownNodes(Arrays.asList( + new RedisMasterReplicaNode("localhost", 1, RedisURI.create("localhost", 1), RedisInstance.Role.UPSTREAM), + new RedisMasterReplicaNode("localhost", 2, RedisURI.create("localhost", 2), RedisInstance.Role.REPLICA))); + + when(clientMock.connectAsync(eq(StringCodec.UTF8), any())) + .thenReturn(ConnectionFuture.completed(null, nodeConnectionMock)); + when(nodeConnectionMock.isOpen()).thenReturn(true); + + sut.setReadFrom(ReadFrom.ANY); + + AtomicReference selectedNode = new AtomicReference<>(); + + StatefulRedisConnection connection = sut.getConnectionAsync(ConnectionIntent.READ, selectedNode::set) + .join(); + + assertThat(connection).isSameAs(nodeConnectionMock); + assertThat(selectedNode.get()).isNotNull(); + assertThat(selectedNode.get().getUri().getPort()).isIn(1, 2); + } + + @Test + void shouldConnectPinnedConnectionToKnownNode() { + + when(clientMock.connectAsync(eq(StringCodec.UTF8), any())) + .thenReturn(ConnectionFuture.completed(null, nodeConnectionMock)); + + StatefulRedisConnection connection = sut.getPinnedConnectionAsync(RedisURI.create("localhost", 1)) + .join(); + + assertThat(connection).isSameAs(nodeConnectionMock); + } + + @Test + void shouldFailPinnedConnectionWhenNodeIsNoLongerKnown() { + + CompletableFuture> future = sut + .getPinnedConnectionAsync(RedisURI.create("localhost", 42)); + + assertThat(future).isCompletedExceptionally(); + assertThatThrownBy(future::join).hasCauseInstanceOf(RedisException.class).hasMessageContaining("no longer available"); + } + } diff --git a/src/test/java/io/lettuce/core/masterreplica/MasterReplicaStickyScanIntegrationTests.java b/src/test/java/io/lettuce/core/masterreplica/MasterReplicaStickyScanIntegrationTests.java new file mode 100644 index 0000000000..821da19f95 --- /dev/null +++ b/src/test/java/io/lettuce/core/masterreplica/MasterReplicaStickyScanIntegrationTests.java @@ -0,0 +1,364 @@ +/* + * Copyright 2020-Present, Redis Ltd. and Contributors + * All rights reserved. + * + * Licensed under the MIT License. + */ +package io.lettuce.core.masterreplica; + +import static io.lettuce.TestTags.INTEGRATION_TEST; +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +import java.time.Duration; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; + +import io.lettuce.core.AbstractRedisClientTest; +import io.lettuce.core.KeyScanCursor; +import io.lettuce.core.MapScanCursor; +import io.lettuce.core.ReadFrom; +import io.lettuce.core.RedisURI; +import io.lettuce.core.ScanArgs; +import io.lettuce.core.ScoredValueScanCursor; +import io.lettuce.core.ValueScanCursor; +import io.lettuce.core.api.sync.RedisCommands; +import io.lettuce.core.codec.StringCodec; +import io.lettuce.core.models.role.RedisInstance; +import io.lettuce.core.models.role.RoleParser; +import io.lettuce.test.TestFutures; +import io.lettuce.test.WithPassword; +import io.lettuce.test.settings.TestSettings; + +/** + * Integration tests for sticky {@code SCAN} cursors on Master/Replica connections, + * #3287. + * + * {@code SCAN} family cursors are node-local. When {@link ReadFrom} spreads reads across nodes, continuation requests (cursor + * != 0) must be routed to the node that served the initial request, otherwise the iteration yields undefined results + * (missing/duplicate elements). These tests iterate by passing the returned cursor object back into the next call (the + * canonical usage, also used by {@code ScanIterator}/{@code ScanStream}), which is the carrier for the cursor's source node. + * + * Matrix: {SCAN, HSCAN, SSCAN, ZSCAN} x {sync, async, reactive} under {@link ReadFrom#ANY}, plus a {@link ReadFrom#UPSTREAM} + * control per command (single-node iteration, must stay green). Runs against the standard master/replica pair provided by the + * project's docker-env ({@link TestSettings#port(int) port(3)}/{@code port(4)}). + * + * @author Sanghun Lee + */ +@Tag(INTEGRATION_TEST) +class MasterReplicaStickyScanIntegrationTests extends AbstractRedisClientTest { + + enum Api { + SYNC, ASYNC, REACTIVE + } + + private static final String HASH_KEY = "sticky-3287-hash"; + + private static final String SET_KEY = "sticky-3287-set"; + + private static final String ZSET_KEY = "sticky-3287-zset"; + + private static final String KEY_PREFIX = "sticky-3287-key-"; + + private static final String KEY_MATCH = KEY_PREFIX + "*"; + + // > *-max-listpack-entries (default 128) so collections are hashtable-encoded and the SCAN family actually + // paginates. Listpack-encoded collections return everything with cursor 0 and would hide the bug. + private static final int ELEMENT_COUNT = 1000; + + private static final int PAGE_SIZE = 10; + + private static final int MAX_CALLS = 100_000; + + private RedisURI upstream; + + private RedisURI replica; + + private RedisCommands connection1; + + private RedisCommands connection2; + + private StatefulRedisMasterReplicaConnection connection; + + @BeforeEach + void before() { + + RedisURI node1 = RedisURI.Builder.redis(host, TestSettings.port(3)).build(); + RedisURI node2 = RedisURI.Builder.redis(host, TestSettings.port(4)).build(); + + connection1 = client.connect(node1).sync(); + connection2 = client.connect(node2).sync(); + + RedisInstance node1Instance = RoleParser.parse(this.connection1.role()); + RedisInstance node2Instance = RoleParser.parse(this.connection2.role()); + + if (node1Instance.getRole().isUpstream() && node2Instance.getRole().isReplica()) { + upstream = node1; + replica = node2; + } else if (node2Instance.getRole().isUpstream() && node1Instance.getRole().isReplica()) { + upstream = node2; + replica = node1; + } else { + assumeTrue(false, + String.format("Cannot run the test because I don't have a distinct master and replica but %s and %s", + node1Instance, node2Instance)); + } + + WithPassword.enableAuthentication(this.connection1); + this.connection1.auth(passwd); + this.connection1.configSet("masterauth", passwd.toString()); + + WithPassword.enableAuthentication(this.connection2); + this.connection2.auth(passwd); + this.connection2.configSet("masterauth", passwd.toString()); + + upstream.setAuthentication(passwd); + replica.setAuthentication(passwd); + + connection = MasterReplica.connect(client, StringCodec.UTF8, Arrays.asList(upstream, replica)); + + seed(); + } + + @AfterEach + void after() { + + if (connection1 != null) { + WithPassword.disableAuthentication(connection1); + connection1.configSet("masterauth", ""); + connection1.configRewrite(); + connection1.getStatefulConnection().close(); + } + + if (connection2 != null) { + WithPassword.disableAuthentication(connection2); + connection2.configSet("masterauth", ""); + connection2.configRewrite(); + connection2.getStatefulConnection().close(); + } + + if (connection != null) { + connection.close(); + } + } + + // seed on master, then wait for the replica to catch up before spreading reads across nodes. + private void seed() { + + connection.setReadFrom(ReadFrom.UPSTREAM); + + RedisCommands sync = connection.sync(); + sync.flushall(); + + Map hash = new HashMap<>(); + Map keys = new HashMap<>(); + String[] members = new String[ELEMENT_COUNT]; + Object[] scoredValues = new Object[ELEMENT_COUNT * 2]; + + for (int i = 0; i < ELEMENT_COUNT; i++) { + hash.put("field-" + i, "value-" + i); + keys.put(KEY_PREFIX + i, "value-" + i); + members[i] = "member-" + i; + scoredValues[i * 2] = (double) i; + scoredValues[i * 2 + 1] = "member-" + i; + } + + sync.hset(HASH_KEY, hash); + sync.mset(keys); + sync.sadd(SET_KEY, members); + sync.zadd(ZSET_KEY, scoredValues); + + sync.waitForReplication(1, 5000); + } + + // --- ReadFrom.ANY: reads spread across master + replica. Without sticky cursors this is #3287 (undefined results). + + @ParameterizedTest + @EnumSource(Api.class) + void scanWithReadFromAnyReturnsAllKeys(Api api) { + connection.setReadFrom(ReadFrom.ANY); + assertThat(scanAll(api)).as("full SCAN under ReadFrom.ANY must return every key").hasSize(ELEMENT_COUNT); + } + + @ParameterizedTest + @EnumSource(Api.class) + void hscanWithReadFromAnyReturnsAllFields(Api api) { + connection.setReadFrom(ReadFrom.ANY); + assertThat(hscanAll(api)).as("full HSCAN under ReadFrom.ANY must return every field").hasSize(ELEMENT_COUNT); + } + + @ParameterizedTest + @EnumSource(Api.class) + void sscanWithReadFromAnyReturnsAllMembers(Api api) { + connection.setReadFrom(ReadFrom.ANY); + assertThat(sscanAll(api)).as("full SSCAN under ReadFrom.ANY must return every member").hasSize(ELEMENT_COUNT); + } + + @ParameterizedTest + @EnumSource(Api.class) + void zscanWithReadFromAnyReturnsAllMembers(Api api) { + connection.setReadFrom(ReadFrom.ANY); + assertThat(zscanAll(api)).as("full ZSCAN under ReadFrom.ANY must return every member").hasSize(ELEMENT_COUNT); + } + + // --- ReadFrom.UPSTREAM controls: single-node iteration, green today, guards against regressions. + + @Test + void scanWithReadFromUpstreamReturnsAllKeys() { + connection.setReadFrom(ReadFrom.UPSTREAM); + assertThat(scanAll(Api.SYNC)).hasSize(ELEMENT_COUNT); + } + + @Test + void hscanWithReadFromUpstreamReturnsAllFields() { + connection.setReadFrom(ReadFrom.UPSTREAM); + assertThat(hscanAll(Api.SYNC)).hasSize(ELEMENT_COUNT); + } + + @Test + void sscanWithReadFromUpstreamReturnsAllMembers() { + connection.setReadFrom(ReadFrom.UPSTREAM); + assertThat(sscanAll(Api.SYNC)).hasSize(ELEMENT_COUNT); + } + + @Test + void zscanWithReadFromUpstreamReturnsAllMembers() { + connection.setReadFrom(ReadFrom.UPSTREAM); + assertThat(zscanAll(Api.SYNC)).hasSize(ELEMENT_COUNT); + } + + // --- full iterations, passing the returned cursor object back into the next call. + + private Set scanAll(Api api) { + + ScanArgs args = ScanArgs.Builder.matches(KEY_MATCH).limit(PAGE_SIZE); + Set seen = new HashSet<>(); + KeyScanCursor cursor = null; + int calls = 0; + + do { + cursor = scanNext(api, cursor, args); + seen.addAll(cursor.getKeys()); + } while (!cursor.isFinished() && ++calls < MAX_CALLS); + + return seen; + } + + private KeyScanCursor scanNext(Api api, KeyScanCursor cursor, ScanArgs args) { + switch (api) { + case SYNC: + return cursor == null ? connection.sync().scan(args) : connection.sync().scan(cursor, args); + case ASYNC: + return TestFutures + .getOrTimeout(cursor == null ? connection.async().scan(args) : connection.async().scan(cursor, args)); + case REACTIVE: + return (cursor == null ? connection.reactive().scan(args) : connection.reactive().scan(cursor, args)) + .block(Duration.ofSeconds(10)); + default: + throw new IllegalStateException("Unsupported API " + api); + } + } + + private Set hscanAll(Api api) { + + ScanArgs args = ScanArgs.Builder.limit(PAGE_SIZE); + Set seen = new HashSet<>(); + MapScanCursor cursor = null; + int calls = 0; + + do { + cursor = hscanNext(api, cursor, args); + seen.addAll(cursor.getMap().keySet()); + } while (!cursor.isFinished() && ++calls < MAX_CALLS); + + return seen; + } + + private MapScanCursor hscanNext(Api api, MapScanCursor cursor, ScanArgs args) { + switch (api) { + case SYNC: + return cursor == null ? connection.sync().hscan(HASH_KEY, args) + : connection.sync().hscan(HASH_KEY, cursor, args); + case ASYNC: + return TestFutures.getOrTimeout(cursor == null ? connection.async().hscan(HASH_KEY, args) + : connection.async().hscan(HASH_KEY, cursor, args)); + case REACTIVE: + return (cursor == null ? connection.reactive().hscan(HASH_KEY, args) + : connection.reactive().hscan(HASH_KEY, cursor, args)).block(Duration.ofSeconds(10)); + default: + throw new IllegalStateException("Unsupported API " + api); + } + } + + private Set sscanAll(Api api) { + + ScanArgs args = ScanArgs.Builder.limit(PAGE_SIZE); + Set seen = new HashSet<>(); + ValueScanCursor cursor = null; + int calls = 0; + + do { + cursor = sscanNext(api, cursor, args); + seen.addAll(cursor.getValues()); + } while (!cursor.isFinished() && ++calls < MAX_CALLS); + + return seen; + } + + private ValueScanCursor sscanNext(Api api, ValueScanCursor cursor, ScanArgs args) { + switch (api) { + case SYNC: + return cursor == null ? connection.sync().sscan(SET_KEY, args) : connection.sync().sscan(SET_KEY, cursor, args); + case ASYNC: + return TestFutures.getOrTimeout(cursor == null ? connection.async().sscan(SET_KEY, args) + : connection.async().sscan(SET_KEY, cursor, args)); + case REACTIVE: + return (cursor == null ? connection.reactive().sscan(SET_KEY, args) + : connection.reactive().sscan(SET_KEY, cursor, args)).block(Duration.ofSeconds(10)); + default: + throw new IllegalStateException("Unsupported API " + api); + } + } + + private Set zscanAll(Api api) { + + ScanArgs args = ScanArgs.Builder.limit(PAGE_SIZE); + Set seen = new HashSet<>(); + ScoredValueScanCursor cursor = null; + int calls = 0; + + do { + cursor = zscanNext(api, cursor, args); + cursor.getValues().forEach(scoredValue -> seen.add(scoredValue.getValue())); + } while (!cursor.isFinished() && ++calls < MAX_CALLS); + + return seen; + } + + private ScoredValueScanCursor zscanNext(Api api, ScoredValueScanCursor cursor, ScanArgs args) { + switch (api) { + case SYNC: + return cursor == null ? connection.sync().zscan(ZSET_KEY, args) + : connection.sync().zscan(ZSET_KEY, cursor, args); + case ASYNC: + return TestFutures.getOrTimeout(cursor == null ? connection.async().zscan(ZSET_KEY, args) + : connection.async().zscan(ZSET_KEY, cursor, args)); + case REACTIVE: + return (cursor == null ? connection.reactive().zscan(ZSET_KEY, args) + : connection.reactive().zscan(ZSET_KEY, cursor, args)).block(Duration.ofSeconds(10)); + default: + throw new IllegalStateException("Unsupported API " + api); + } + } + +}