Support multi with re-auth

Defer the re-auth operation in case there is on-going multi
Tx in lettuce need to be externally synchronised when used in multithreaded env. Since re-auth happens from different thread we need to make sure it does not happen while there is ongoing transaction.

Author: ggivo

src/main/java/io/lettuce/core/ConnectionBuilder.java

@@ -113,17 +113,6 @@ public void apply(RedisURI redisURI) {
         bootstrap.attr(REDIS_URI, redisURI.toString());
     }
 
-    public void registerAuthenticationHandler(RedisCredentialsProvider credentialsProvider, ConnectionState state,
-            Boolean isPubSubConnection) {
-        LettuceAssert.assertState(endpoint != null, "Endpoint must be set");
-        LettuceAssert.assertState(connection != null, "Connection must be set");
-        LettuceAssert.assertState(clientResources != null, "ClientResources must be set");
-
-        RedisAuthenticationHandler authenticationHandler = new RedisAuthenticationHandler(connection, credentialsProvider,
-                state, clientResources.eventBus(), isPubSubConnection);
-        endpoint.registerAuthenticationHandler(authenticationHandler);
-    }
-
     protected List<ChannelHandler> buildHandlers() {
 
         LettuceAssert.assertState(channelGroup != null, "ChannelGroup must be set");

src/main/java/io/lettuce/core/RedisAuthenticationHandler.java

@@ -46,26 +46,18 @@ public class RedisAuthenticationHandler {
 
     private static final InternalLogger log = InternalLoggerFactory.getInstance(RedisAuthenticationHandler.class);
 
-    private final RedisChannelHandler<?, ?> connection;
-
-    private final ConnectionState state;
-
-    private final RedisCommandBuilder<String, String> commandBuilder = new RedisCommandBuilder<>(StringCodec.UTF8);
+    private final StatefulRedisConnectionImpl<?, ?> connection;
 
     private final RedisCredentialsProvider credentialsProvider;
 
     private final AtomicReference<Disposable> credentialsSubscription = new AtomicReference<>();
 
-    private final EventBus eventBus;
-
     private final Boolean isPubSubConnection;
 
-    public RedisAuthenticationHandler(RedisChannelHandler<?, ?> connection, RedisCredentialsProvider credentialsProvider,
-            ConnectionState state, EventBus eventBus, Boolean isPubSubConnection) {
+    public RedisAuthenticationHandler(StatefulRedisConnectionImpl<?, ?> connection,
+            RedisCredentialsProvider credentialsProvider, Boolean isPubSubConnection) {
         this.connection = connection;
-        this.state = state;
         this.credentialsProvider = credentialsProvider;
-        this.eventBus = eventBus;
         this.isPubSubConnection = isPubSubConnection;
     }
 
@@ -125,55 +117,19 @@ protected void onError(Throwable e) {
      * @param credentials the new credentials
      */
     protected void reauthenticate(RedisCredentials credentials) {
-        CharSequence password = CharBuffer.wrap(credentials.getPassword());
-
-        AsyncCommand<String, String, String> authCmd;
-        if (credentials.hasUsername()) {
-            authCmd = new AsyncCommand<>(commandBuilder.auth(credentials.getUsername(), password));
-        } else {
-            authCmd = new AsyncCommand<>(commandBuilder.auth(password));
-        }
-
-        dispatchAuth(authCmd).thenRun(() -> {
-            publishReauthEvent();
-            log.info("Re-authentication succeeded for endpoint {}.", getEpid());
-        }).exceptionally(throwable -> {
-            publishReauthFailedEvent(throwable);
-            log.error("Re-authentication failed for endpoint {}.", getEpid(), throwable);
-            return null;
-        });
-    }
-
-    private AsyncCommand<?, ?, ?> dispatchAuth(RedisCommand<?, ?, ?> authCommand) {
-        AsyncCommand asyncCommand = new AsyncCommand<>(authCommand);
-        RedisCommand<?, ?, ?> dispatched = connection.dispatch(asyncCommand);
-        if (dispatched instanceof AsyncCommand) {
-            return (AsyncCommand<?, ?, ?>) dispatched;
-        }
-        return asyncCommand;
-    }
-
-    private void publishReauthEvent() {
-        eventBus.publish(new ReauthenticateEvent(getEpid()));
-    }
-
-    private void publishReauthFailedEvent(Throwable throwable) {
-        eventBus.publish(new ReauthenticateFailedEvent(getEpid(), throwable));
+        connection.setCredentials(credentials);
     }
 
     protected boolean isSupportedConnection() {
-        if (isPubSubConnection && ProtocolVersion.RESP2 == state.getNegotiatedProtocolVersion()) {
+        if (isPubSubConnection && ProtocolVersion.RESP2 == connection.getConnectionState().getNegotiatedProtocolVersion()) {
             log.warn("Renewable credentials are not supported with RESP2 protocol on a pub/sub connection.");
             return false;
         }
         return true;
     }
 
-    private String getEpid() {
-        if (connection.getChannelWriter() instanceof Endpoint) {
-            return ((Endpoint) connection.getChannelWriter()).getId();
-        }
-        return "unknown";
+    private void publishReauthFailedEvent(Throwable throwable) {
+        connection.getResources().eventBus().publish(new ReauthenticateFailedEvent(throwable));
     }
 
     public static boolean isSupported(ClientOptions clientOptions) {

src/main/java/io/lettuce/core/RedisClient.java

@@ -317,7 +317,10 @@ private <K, V, S> ConnectionFuture<S> connectStatefulAsync(StatefulRedisConnecti
         ConnectionState state = connection.getConnectionState();
         state.apply(redisURI);
         state.setDb(redisURI.getDatabase());
-
+        if (RedisAuthenticationHandler.isSupported(getOptions())) {
+            connection.setAuthenticationHandler(
+                    new RedisAuthenticationHandler(connection, redisURI.getCredentialsProvider(), isPubSub));
+        }
         connectionBuilder.connection(connection);
         connectionBuilder.clientOptions(getOptions());
         connectionBuilder.clientResources(getResources());
@@ -326,11 +329,6 @@ private <K, V, S> ConnectionFuture<S> connectStatefulAsync(StatefulRedisConnecti
         connectionBuilder(getSocketAddressSupplier(redisURI), connectionBuilder, connection.getConnectionEvents(), redisURI);
         connectionBuilder.connectionInitializer(createHandshake(state));
 
-        if (RedisAuthenticationHandler.isSupported(getOptions())) {
-            connectionBuilder.registerAuthenticationHandler(redisURI.getCredentialsProvider(), connection.getConnectionState(),
-                    isPubSub);
-        }
-
         ConnectionFuture<RedisChannelHandler<K, V>> future = initializeChannelAsync(connectionBuilder);
 
         return future.thenApply(channelHandler -> (S) connection);

src/main/java/io/lettuce/core/StatefulRedisConnectionImpl.java

@@ -22,10 +22,14 @@
 import static io.lettuce.core.ClientOptions.DEFAULT_JSON_PARSER;
 import static io.lettuce.core.protocol.CommandType.*;
 
+import java.nio.CharBuffer;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
@@ -37,10 +41,14 @@
 import io.lettuce.core.cluster.api.sync.RedisClusterCommands;
 import io.lettuce.core.codec.RedisCodec;
 import io.lettuce.core.codec.StringCodec;
+import io.lettuce.core.event.connection.ReauthenticateEvent;
+import io.lettuce.core.event.connection.ReauthenticateFailedEvent;
 import io.lettuce.core.json.JsonParser;
 import io.lettuce.core.output.MultiOutput;
 import io.lettuce.core.output.StatusOutput;
 import io.lettuce.core.protocol.*;
+import io.netty.util.internal.logging.InternalLogger;
+import io.netty.util.internal.logging.InternalLoggerFactory;
 import reactor.core.publisher.Mono;
 
 /**
@@ -55,6 +63,8 @@
  */
 public class StatefulRedisConnectionImpl<K, V> extends RedisChannelHandler<K, V> implements StatefulRedisConnection<K, V> {
 
+    private static final InternalLogger logger = InternalLoggerFactory.getInstance(StatefulRedisConnectionImpl.class);
+
     protected final RedisCodec<K, V> codec;
 
     protected final RedisCommands<K, V> sync;
@@ -71,6 +81,14 @@ public class StatefulRedisConnectionImpl<K, V> extends RedisChannelHandler<K, V>
 
     protected MultiOutput<K, V> multi;
 
+    private RedisAuthenticationHandler authHandler;
+
+    private AtomicReference<RedisCredentials> credentialsRef = new AtomicReference<>();
+
+    private final ReentrantLock reAuthSafety = new ReentrantLock();
+
+    private AtomicBoolean inTransaction = new AtomicBoolean(false);
+
     /**
      * Initialize a new connection.
      *
@@ -181,20 +199,38 @@ public boolean isMulti() {
     public <T> RedisCommand<K, V, T> dispatch(RedisCommand<K, V, T> command) {
 
         RedisCommand<K, V, T> toSend = preProcessCommand(command);
-        return super.dispatch(toSend);
+        RedisCommand<K, V, T> result = super.dispatch(toSend);
+        if (toSend.getType() == EXEC || toSend.getType() == DISCARD) {
+            inTransaction.set(false);
+            setCredentials(credentialsRef.getAndSet(null));
+        }
+
+        return result;
     }
 
     @Override
     public Collection<RedisCommand<K, V, ?>> dispatch(Collection<? extends RedisCommand<K, V, ?>> commands) {
 
         List<RedisCommand<K, V, ?>> sentCommands = new ArrayList<>(commands.size());
 
-        commands.forEach(o -> {
+        boolean transactionComplete = false;
+        for (RedisCommand<K, V, ?> o : commands) {
             RedisCommand<K, V, ?> command = preProcessCommand(o);
             sentCommands.add(command);
-        });
+            if (command.getType() == EXEC) {
+                transactionComplete = true;
+            }
+            if (command.getType() == MULTI || command.getType() == DISCARD) {
+                transactionComplete = false;
+            }
+        }
 
-        return super.dispatch(sentCommands);
+        Collection<RedisCommand<K, V, ?>> result = super.dispatch(sentCommands);
+        if (transactionComplete) {
+            inTransaction.set(false);
+            setCredentials(credentialsRef.getAndSet(null));
+        }
+        return result;
     }
 
     // TODO [tihomir.mateev] Refactor to include as part of the Command interface
@@ -273,12 +309,20 @@ protected <T> RedisCommand<K, V, T> preProcessCommand(RedisCommand<K, V, T> comm
 
         if (commandType.equals(MULTI.name())) {
 
+            reAuthSafety.lock();
+            try {
+                inTransaction.set(true);
+            } finally {
+                reAuthSafety.unlock();
+            }
             multi = (multi == null ? new MultiOutput<>(codec) : multi);
 
             if (command instanceof CompleteableCommand) {
                 ((CompleteableCommand<?>) command).onComplete((ignored, e) -> {
                     if (e != null) {
                         multi = null;
+                        inTransaction.set(false);
+                        setCredentials(credentialsRef.getAndSet(null));
                     }
                 });
             }
@@ -318,11 +362,72 @@ public ConnectionState getConnectionState() {
     @Override
     public void activated() {
         super.activated();
+        if (authHandler != null) {
+            authHandler.subscribe();
+        }
     }
 
     @Override
     public void deactivated() {
+        if (authHandler != null) {
+            authHandler.unsubscribe();
+        }
         super.deactivated();
     }
 
+    public void setAuthenticationHandler(RedisAuthenticationHandler handler) {
+        authHandler = handler;
+    }
+
+    public void setCredentials(RedisCredentials credentials) {
+        if (credentials == null) {
+            return;
+        }
+        reAuthSafety.lock();
+        try {
+            credentialsRef.set(credentials);
+            if (!inTransaction.get()) {
+                dispatchAuthCommand(credentialsRef.getAndSet(null));
+            }
+        } finally {
+            reAuthSafety.unlock();
+        }
+    }
+
+    private void dispatchAuthCommand(RedisCredentials credentials) {
+        if (credentials == null) {
+            return;
+        }
+
+        RedisFuture<String> auth;
+        if (credentials.getUsername() != null) {
+            auth = async().auth(credentials.getUsername(), CharBuffer.wrap(credentials.getPassword()));
+        } else {
+            auth = async().auth(CharBuffer.wrap(credentials.getPassword()));
+        }
+        auth.thenRun(() -> {
+            publishReauthEvent();
+            logger.info("Re-authentication succeeded {}.", getEpid());
+        }).exceptionally(throwable -> {
+            publishReauthFailedEvent(throwable);
+            logger.error("Re-authentication failed  {}.", getEpid(), throwable);
+            return null;
+        });
+    }
+
+    private void publishReauthEvent() {
+        getResources().eventBus().publish(new ReauthenticateEvent(getEpid()));
+    }
+
+    private void publishReauthFailedEvent(Throwable throwable) {
+        getResources().eventBus().publish(new ReauthenticateFailedEvent(getEpid(), throwable));
+    }
+
+    private String getEpid() {
+        if (getChannelWriter() instanceof Endpoint) {
+            return ((Endpoint) getChannelWriter()).getId();
+        }
+        return "";
+    }
+
 }

src/main/java/io/lettuce/core/api/StatefulRedisConnection.java

@@ -1,5 +1,6 @@
 package io.lettuce.core.api;
 
+import io.lettuce.core.RedisCredentials;
 import io.lettuce.core.api.async.RedisAsyncCommands;
 import io.lettuce.core.api.push.PushListener;
 import io.lettuce.core.api.reactive.RedisReactiveCommands;