From e328aa0c4b7dd0fca79ab2b46bc7376175d8beb3 Mon Sep 17 00:00:00 2001
From: "xiaofan.chen" <xiaofan.chen@okg.com>
Date: Sat, 17 Aug 2024 10:19:49 +0800
Subject: [PATCH] fix: revert protectMode upon channelActive event, complete
 non-retryable activation command in drainStackUponChannelInactive()

---
 .../lettuce/core/protocol/CommandHandler.java | 13 ++++---
 .../DefaultAutoBatchFlushEndpoint.java        | 36 +++++++++----------
 2 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/src/main/java/io/lettuce/core/protocol/CommandHandler.java b/src/main/java/io/lettuce/core/protocol/CommandHandler.java
index de6cf0f0e..e19021b96 100644
--- a/src/main/java/io/lettuce/core/protocol/CommandHandler.java
+++ b/src/main/java/io/lettuce/core/protocol/CommandHandler.java
@@ -190,13 +190,18 @@ void setBuffer(ByteBuf buffer) {
         return drainCommands(stack);
     }
 
-    private Deque<RedisCommand<?, ?, ?>> drainStack() {
+    private Deque<RedisCommand<?, ?, ?>> drainStackUponChannelInactive() {
         final Deque<RedisCommand<?, ?, ?>> target = new ArrayDeque<>(stack.size());
 
         RedisCommand<?, ?, ?> cmd;
         while ((cmd = stack.poll()) != null) {
-            if (!cmd.isDone() && !ActivationCommand.isActivationCommand(cmd)) {
-                target.add(cmd);
+            if (!cmd.isDone()) {
+                if (!ActivationCommand.isActivationCommand(cmd)) {
+                    target.add(cmd);
+                } else {
+                    cmd.completeExceptionally(
+                            new RedisConnectionException("activation command won't be retried upon channel inactive"));
+                }
             }
         }
 
@@ -379,7 +384,7 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception {
         endpoint.notifyChannelInactive(ctx.channel());
         Deque<RedisCommand<?, ?, ?>> autoBatchFlushRetryableDrainQueuedCommands = UnmodifiableDeque.emptyDeque();
         if (supportsAutoBatchFlush) {
-            autoBatchFlushRetryableDrainQueuedCommands = drainStack();
+            autoBatchFlushRetryableDrainQueuedCommands = drainStackUponChannelInactive();
         } else {
             endpoint.notifyDrainQueuedCommands(this);
         }
diff --git a/src/main/java/io/lettuce/core/protocol/DefaultAutoBatchFlushEndpoint.java b/src/main/java/io/lettuce/core/protocol/DefaultAutoBatchFlushEndpoint.java
index 2bb6d1d8a..f433b2132 100644
--- a/src/main/java/io/lettuce/core/protocol/DefaultAutoBatchFlushEndpoint.java
+++ b/src/main/java/io/lettuce/core/protocol/DefaultAutoBatchFlushEndpoint.java
@@ -133,6 +133,8 @@ protected static void cancelCommandOnEndpointClose(RedisCommand<?, ?, ?> cmd) {
 
     protected volatile @Nonnull ContextualChannel channel = DummyContextualChannelInstances.CHANNEL_CONNECTING;
 
+    private volatile Throwable failedToReconnectReason;
+
     private final Consumer<RedisCommand<?, ?, ?>> callbackOnClose;
 
     private final boolean rejectCommandsWhileDisconnected;
@@ -153,19 +155,17 @@ protected static void cancelCommandOnEndpointClose(RedisCommand<?, ?, ?> cmd) {
 
     private ConnectionFacade connectionFacade;
 
-    private volatile Throwable connectionError;
-
     private final String cachedEndpointId;
 
     protected final UnboundedOfferFirstQueue<Object> taskQueue;
 
     private final boolean canFire;
 
-    private volatile boolean inProtectMode;
+    private volatile EventLoop lastEventLoop = null;
 
-    private volatile Throwable failedToReconnectReason;
+    private volatile Throwable connectionError;
 
-    private volatile EventLoop lastEventLoop = null;
+    private volatile boolean inProtectMode;
 
     private final int writeSpinCount;
 
@@ -316,17 +316,16 @@ private <V, K> void writeAndFlushActivationCommands(ContextualChannel chan,
     @Override
     public void notifyChannelActive(Channel channel) {
         final ContextualChannel contextualChannel = new ContextualChannel(channel, ConnectionContext.State.CONNECTED);
-
-        this.logPrefix = null;
-        this.connectionError = null;
-
         if (!CHANNEL.compareAndSet(this, DummyContextualChannelInstances.CHANNEL_CONNECTING, contextualChannel)) {
             channel.close();
             onUnexpectedState("notifyChannelActive", ConnectionContext.State.CONNECTING);
             return;
         }
 
-        lastEventLoop = channel.eventLoop();
+        this.lastEventLoop = channel.eventLoop();
+        this.connectionError = null;
+        this.inProtectMode = false;
+        this.logPrefix = null;
 
         // Created a synchronize-before with set channel to CHANNEL_CONNECTING,
         if (isClosed()) {
@@ -398,7 +397,7 @@ public void notifyChannelInactive(Channel channel) {
 
     @Override
     public void notifyChannelInactiveAfterWatchdogDecision(Channel channel,
-            Deque<RedisCommand<?, ?, ?>> retryableQueuedCommands) {
+            Deque<RedisCommand<?, ?, ?>> retryablePendingCommands) {
         final ContextualChannel inactiveChan = this.channel;
         if (!inactiveChan.context.initialState.isConnected()) {
             logger.error("[unexpected][{}] notifyChannelInactive: channel initial state not connected", logPrefix());
@@ -446,7 +445,7 @@ public void notifyChannelInactiveAfterWatchdogDecision(Channel channel,
             CHANNEL.set(this, DummyContextualChannelInstances.CHANNEL_ENDPOINT_CLOSED);
         }
         inactiveChan.context
-                .setCloseStatus(new ConnectionContext.CloseStatus(willReconnect, retryableQueuedCommands, exception));
+                .setCloseStatus(new ConnectionContext.CloseStatus(willReconnect, retryablePendingCommands, exception));
         trySetEndpointQuiescence(inactiveChan);
     }
 
@@ -945,11 +944,11 @@ private <K, V, T> RedisCommand<K, V, T> processActivationCommand(RedisCommand<K,
 
     private Throwable validateWrite(ContextualChannel chan, int commands, boolean isActivationCommand) {
         if (isClosed()) {
-            return new RedisException("Connection is closed");
+            return new RedisException("Endpoint is closed");
         }
 
         final Throwable localConnectionErr = connectionError;
-        if (localConnectionErr != null /* different logic of DefaultEndpoint */) {
+        if (localConnectionErr != null /* attention: different logic of DefaultEndpoint */) {
             return localConnectionErr;
         }
 
@@ -961,18 +960,19 @@ private Throwable validateWrite(ContextualChannel chan, int commands, boolean is
 
         final ConnectionContext.State initialState = chan.context.initialState;
         final boolean rejectCommandsWhileDisconnectedLocal = this.rejectCommandsWhileDisconnected || isActivationCommand;
+        final String rejectDesc = isActivationCommand ? "isActivationCommand" : "rejectCommandsWhileDisconnected";
         switch (initialState) {
             case ENDPOINT_CLOSED:
                 return new RedisException("Connection is closed");
             case RECONNECT_FAILED:
-                return failedToReconnectReason;
+                return getFailedToReconnectReason();
             case WILL_RECONNECT:
             case CONNECTING:
-                return rejectCommandsWhileDisconnectedLocal
-                        ? new RedisException("Currently not connected. Commands are rejected.")
+                return rejectCommandsWhileDisconnectedLocal ? new RedisException("Currently not connected and " + rejectDesc)
                         : null;
             case CONNECTED:
-                return !chan.isActive() && rejectCommandsWhileDisconnectedLocal ? new RedisException("Channel is closed")
+                return !chan.isActive() && rejectCommandsWhileDisconnectedLocal
+                        ? new RedisException("Channel is inactive and " + rejectDesc)
                         : null;
             default:
                 throw new IllegalStateException("unexpected state: " + initialState);