Route Master/Replica SCAN continuations to the issuing node #3287

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<K, V> extends BaseRedisCommandBuilder<K, V> {
@@ -1739,6 +1740,7 @@ Command<K, V, MapScanCursor<K, V>> hscan(K key, ScanCursor scanCursor, ScanArgs
         scanArgs(scanCursor, scanArgs, args);
 
         MapScanOutput<K, V> output = new MapScanOutput<>(codec);
+        associateSourceHint(output, scanCursor);
         return createCommand(HSCAN, output, args);
     }
 
@@ -1753,6 +1755,7 @@ Command<K, V, KeyScanCursor<K>> hscanNovalues(K key, ScanCursor scanCursor, Scan
         args.add(NOVALUES);
 
         KeyScanOutput<K, V> output = new KeyScanOutput<>(codec);
+        associateSourceHint(output, scanCursor);
         return createCommand(HSCAN, output, args);
     }
 
@@ -1809,6 +1812,7 @@ Command<K, V, StreamScanCursor> hscanStreaming(KeyValueStreamingChannel<K, V> ch
         scanArgs(scanCursor, scanArgs, args);
 
         KeyValueScanStreamingOutput<K, V> output = new KeyValueScanStreamingOutput<>(codec, channel);
+        associateSourceHint(output, scanCursor);
         return createCommand(HSCAN, output, args);
     }
 
@@ -1825,6 +1829,7 @@ Command<K, V, StreamScanCursor> hscanNoValuesStreaming(KeyStreamingChannel<K> ch
         args.add(NOVALUES);
 
         KeyScanStreamingOutput<K, V> output = new KeyScanStreamingOutput<>(codec, channel);
+        associateSourceHint(output, scanCursor);
         return createCommand(HSCAN, output, args);
     }
 
@@ -2644,6 +2649,7 @@ Command<K, V, KeyScanCursor<K>> scan(ScanCursor scanCursor, ScanArgs scanArgs) {
         scanArgs(scanCursor, scanArgs, args);
 
         KeyScanOutput<K, V> output = new KeyScanOutput<>(codec);
+        associateSourceHint(output, scanCursor);
         return createCommand(SCAN, output, args);
     }
 
@@ -2658,6 +2664,18 @@ protected void scanArgs(ScanCursor scanCursor, ScanArgs scanArgs, CommandArgs<K,
         }
     }
 
+    /**
+     * Propagate the source-node association from the input {@link ScanCursor} to the cursor carrying the response, so that
+     * drivers routing by cursor source (see {@link ScanCursor#getSource()}) can pin continuation requests to the node that
+     * issued the cursor.
+     */
+    private static void associateSourceHint(ScanOutput<?, ?, ?> output, ScanCursor scanCursor) {
+
+        if (scanCursor.getSource() != null) {
+            output.get().setSource(scanCursor.getSource());
+        }
+    }
+
     Command<K, V, StreamScanCursor> scanStreaming(KeyStreamingChannel<K> channel) {
         notNull(channel);
         LettuceAssert.notNull(channel, "KeyStreamingChannel " + MUST_NOT_BE_NULL);
@@ -2687,6 +2705,7 @@ Command<K, V, StreamScanCursor> scanStreaming(KeyStreamingChannel<K> channel, Sc
         scanArgs(scanCursor, scanArgs, args);
 
         KeyScanStreamingOutput<K, V> output = new KeyScanStreamingOutput<>(codec, channel);
+        associateSourceHint(output, scanCursor);
         return createCommand(SCAN, output, args);
     }
 
@@ -3088,6 +3107,7 @@ Command<K, V, ValueScanCursor<V>> sscan(K key, ScanCursor scanCursor, ScanArgs s
         scanArgs(scanCursor, scanArgs, args);
 
         ValueScanOutput<K, V> output = new ValueScanOutput<>(codec);
+        associateSourceHint(output, scanCursor);
         return createCommand(SSCAN, output, args);
     }
 
@@ -3123,6 +3143,7 @@ Command<K, V, StreamScanCursor> sscanStreaming(ValueStreamingChannel<V> channel,
         scanArgs(scanCursor, scanArgs, args);
 
         ValueScanStreamingOutput<K, V> output = new ValueScanStreamingOutput<>(codec, channel);
+        associateSourceHint(output, scanCursor);
         return createCommand(SSCAN, output, args);
     }
 
@@ -4658,6 +4679,7 @@ Command<K, V, ScoredValueScanCursor<V>> zscan(K key, ScanCursor scanCursor, Scan
         scanArgs(scanCursor, scanArgs, args);
 
         ScoredValueScanOutput<K, V> output = new ScoredValueScanOutput<>(codec);
+        associateSourceHint(output, scanCursor);
         return createCommand(ZSCAN, output, args);
     }
 
@@ -4693,6 +4715,7 @@ Command<K, V, StreamScanCursor> zscanStreaming(ScoredValueStreamingChannel<V> ch
         scanArgs(scanCursor, scanArgs, args);
 
         ScoredValueScanStreamingOutput<K, V> output = new ScoredValueScanStreamingOutput<>(codec, channel);
+        associateSourceHint(output, scanCursor);
         return createCommand(ZSCAN, output, args);
     }
 

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());
+        }
+
     }
 
 }

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);
+    }
+
+}

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 <K, V, T> RedisCommand<K, V, T> write(RedisCommand<K, V, T> command) {
 
         ConnectionIntent connectionIntent = inTransaction ? ConnectionIntent.WRITE
                 : (readOnlyCommands.isReadOnly(command) ? ConnectionIntent.READ : ConnectionIntent.WRITE);
-        CompletableFuture<StatefulRedisConnection<K, V>> future = (CompletableFuture) masterReplicaConnectionProvider
-                .getConnectionAsync(connectionIntent);
+        CompletableFuture<StatefulRedisConnection<K, V>> future = (CompletableFuture) getConnectionAsync(command,
+                connectionIntent);
 
         if (isEndTransaction(command.getType())) {
             inTransaction = false;
@@ -94,6 +103,90 @@ public <K, V, T> RedisCommand<K, V, T> write(RedisCommand<K, V, T> 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<? extends StatefulRedisConnection<?, ?>> 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 <em>continuations</em> 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 <K, V> Consumer<RedisNodeDescription> scanCursorStamper(
+            Collection<? extends RedisCommand<K, V, ?>> commands) {
+
+        List<ScanCursor> cursors = null;
+
+        for (RedisCommand<K, V, ?> 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<ScanCursor> 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 <K, V> void writeCommand(RedisCommand<K, V, ?> command, StatefulRedisConnection<K, V> connection,
             Throwable throwable) {
@@ -131,8 +224,13 @@ private static <K, V> void writeCommand(RedisCommand<K, V, ?> command, StatefulR
         // Currently: Retain order
         ConnectionIntent connectionIntent = inTransaction ? ConnectionIntent.WRITE : getIntent(commands);
 
-        CompletableFuture<StatefulRedisConnection<K, V>> future = (CompletableFuture) masterReplicaConnectionProvider
-                .getConnectionAsync(connectionIntent);
+        Consumer<RedisNodeDescription> scanCursorStamper = connectionIntent == ConnectionIntent.READ
+                ? scanCursorStamper(commands)
+                : null;
+
+        CompletableFuture<StatefulRedisConnection<K, V>> future = (CompletableFuture) (scanCursorStamper != null
+                ? masterReplicaConnectionProvider.getConnectionAsync(connectionIntent, scanCursorStamper)
+                : masterReplicaConnectionProvider.getConnectionAsync(connectionIntent));
 
         for (RedisCommand<K, V, ?> command : commands) {
             if (isEndTransaction(command.getType())) {