[automatic failover] Refactor Circuit Breaker Command Tracking

src/main/java/io/lettuce/core/failover/DatabaseCommandTracker.java

@@ -1,17 +1,19 @@
 package io.lettuce.core.failover;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import io.lettuce.core.RedisCommandTimeoutException;
 import io.lettuce.core.failover.api.RedisCircuitBreakerException;
 import io.lettuce.core.protocol.CommandHandler;
-import io.lettuce.core.protocol.CompleteableCommand;
+import io.lettuce.core.protocol.CommandWrapper;
 import io.lettuce.core.protocol.RedisCommand;
 import io.netty.channel.Channel;
+import io.netty.util.concurrent.Future;
+import io.netty.util.concurrent.GenericFutureListener;
 
 /**
  * Tracks command results and records them to the circuit breaker.
@@ -86,6 +88,13 @@ public <K, V, T> RedisCommand<K, V, T> write(RedisCommand<K, V, T> command) {
             return command;
         }
 
+        @SuppressWarnings("unchecked")
+        CircuitBreakerAwareCommand<K, V, T> circuitBreakerAwareCommand = CommandWrapper.unwrap(command,
+                CircuitBreakerAwareCommand.class);
+        if (circuitBreakerAwareCommand == null) {
+            command = new CircuitBreakerAwareCommand<>(command, circuitBreaker.getGeneration());
+        }
+
         RedisCommand<K, V, T> result;
         try {
             // Delegate to parent
@@ -95,9 +104,6 @@ public <K, V, T> RedisCommand<K, V, T> write(RedisCommand<K, V, T> command) {
             throw e;
         }
 
-        CircuitBreakerGeneration generation = circuitBreaker.getGeneration();
-        // Attach completion callback to track success/failure
-        attachRecorder(generation, result);
         return result;
     }
 
@@ -109,36 +115,59 @@ public <K, V, T> RedisCommand<K, V, T> write(RedisCommand<K, V, T> command) {
             commands.forEach(c -> c.completeExceptionally(RedisCircuitBreakerException.INSTANCE));
             return (Collection) commands;
         }
+
+        Collection<RedisCommand<K, V, ?>> circuitBreakerAwareCommands = new ArrayList<>(commands.size());
+        for (RedisCommand<K, V, ?> command : commands) {
+            @SuppressWarnings("unchecked")
+            CircuitBreakerAwareCommand<K, V, ?> circuitBreakerAwareCommand = CommandWrapper.unwrap(command,
+                    CircuitBreakerAwareCommand.class);
+            if (circuitBreakerAwareCommand == null) {
+                command = new CircuitBreakerAwareCommand<>(command, circuitBreaker.getGeneration());
+            }
+            circuitBreakerAwareCommands.add(command);
+        }
+
+        commands = circuitBreakerAwareCommands;
+
         Collection<RedisCommand<K, V, ?>> result;
         try {
             // Delegate to parent
             result = commandWriter.writeMany(commands);
         } catch (Exception e) {
-            // TODO: here not sure we should record exception for each command or just once for the batch
-            circuitBreaker.getGeneration().recordResult(e);
+            commands.forEach(c -> circuitBreaker.getGeneration().recordResult(e));
             throw e;
         }
 
-        // Attach completion callbacks to track success/failure for each command
-        CircuitBreakerGeneration generation = circuitBreaker.getGeneration();
-        for (RedisCommand<K, V, ?> command : result) {
-            attachRecorder(generation, command);
-        }
         return result;
     }
 
-    private <K, V> void attachRecorder(CircuitBreakerGeneration generation, RedisCommand<K, V, ?> command) {
-        if (command instanceof CompleteableCommand) {
-            @SuppressWarnings("unchecked")
-            CompleteableCommand<Object> completeable = (CompleteableCommand<Object>) command;
-            completeable.onComplete((o, e) -> {
-                // Only record timeout exceptions
-                // Other exceptions are tracked inside the pipeline vioa MultiDbOutboundHandler
-                if (e instanceof RedisCommandTimeoutException) {
-                    generation.recordResult(e);
-                }
-            });
+    static class CircuitBreakerAwareCommand<K, V, T> extends CommandWrapper<K, V, T>
+            implements GenericFutureListener<Future<Void>> {
+
+        private final CircuitBreakerGeneration generation;
+
+        public CircuitBreakerAwareCommand(RedisCommand<K, V, T> command, CircuitBreakerGeneration generation) {
+            super(command);
+            this.generation = generation;
+        }
+
+        @Override
+        protected void doBeforeComplete() {
+            generation.recordResult(null);
+        }
+
+        @Override
+        protected void doBeforeError(Throwable throwable) {
+            generation.recordResult(throwable);
         }
+
+        @Override
+        public void operationComplete(Future<Void> future) throws Exception {
+            if (!future.isSuccess()) {
+                generation.recordResult(future.cause());
+            }
+        }
+
     }
 
 }

src/main/java/io/lettuce/core/failover/MultiDbOutboundHandler.java

@@ -1,15 +1,11 @@
 package io.lettuce.core.failover;
 
-import java.util.List;
+import java.util.Collection;
 
-import io.lettuce.core.RedisCommandTimeoutException;
-import io.lettuce.core.protocol.CompleteableCommand;
-import io.lettuce.core.protocol.RedisCommand;
+import io.lettuce.core.failover.DatabaseCommandTracker.CircuitBreakerAwareCommand;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelOutboundHandlerAdapter;
 import io.netty.channel.ChannelPromise;
-import io.netty.util.concurrent.Future;
-import io.netty.util.concurrent.GenericFutureListener;
 
 /**
  * Channel handler for tracking command results and recording them to the circuit breaker.
@@ -39,15 +35,18 @@ public MultiDbOutboundHandler(CircuitBreaker circuitBreaker) {
     public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
         if (circuitBreaker != null) {
 
-            if (msg instanceof RedisCommand<?, ?, ?>) {
-                promise.addListener(recorder(circuitBreaker.getGeneration(), (RedisCommand<?, ?, ?>) msg));
+            if (msg instanceof CircuitBreakerAwareCommand<?, ?, ?>) {
+                promise.addListener(((CircuitBreakerAwareCommand<?, ?, ?>) msg));
             }
 
-            if (msg instanceof List) {
-                List<?> list = (List<?>) msg;
-                for (Object o : list) {
-                    if (o instanceof RedisCommand<?, ?, ?>) {
-                        promise.addListener(recorder(circuitBreaker.getGeneration(), (RedisCommand<?, ?, ?>) o));
+            if (msg instanceof Collection) {
+                Collection<?> collection = (Collection<?>) msg;
+                for (Object cmd : collection) {
+                    if (cmd instanceof CircuitBreakerAwareCommand<?, ?, ?>) {
+                        // we want to use the generation that is active at the time of write.
+                        // Using a lamda which acces directly to circuitBreaker.getGeneration() here would delay capturing the
+                        // generation and use a potentially wrong generation, a later one.
+                        promise.addListener(((CircuitBreakerAwareCommand<?, ?, ?>) cmd));
                     }
                 }
             }
@@ -57,31 +56,4 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise)
 
     }
 
-    private GenericFutureListener<? extends Future<? super Void>> recorder(CircuitBreakerGeneration generation,
-            RedisCommand<?, ?, ?> command) {
-        return future -> {
-            // Record failure if write failed
-            if (!future.isSuccess()) {
-                generation.recordResult(future.cause());
-            } else {
-                // Track command completion if write succeeded
-                recordOnCommandComplete(generation, command);
-            }
-        };
-    }
-
-    private void recordOnCommandComplete(CircuitBreakerGeneration generation, RedisCommand<?, ?, ?> command) {
-        // Attach completion callback to track success/failure
-        if (command instanceof CompleteableCommand) {
-            CompleteableCommand<?> completeable = (CompleteableCommand<?>) command;
-            completeable.onComplete((o, e) -> {
-                // Record failures except for RedisCommandTimeoutException
-                // Timeouts are recorded by DatabaseCommandTracker
-                if (!(e instanceof RedisCommandTimeoutException)) {
-                    generation.recordResult(e);
-                }
-            });
-        }
-    }
-
 }

src/main/java/io/lettuce/core/protocol/CommandWrapper.java

@@ -64,13 +64,24 @@ public void complete() {
 
         if (consumers != COMPLETE && ONCOMPLETE.compareAndSet(this, consumers, COMPLETE)) {
 
+            doBeforeComplete();
+
             command.complete();
 
             doOnComplete();
             notifyConsumers(consumers);
         }
     }
 
+    /**
+     * Callback method called after successful completion and before notifying downstream consumers.
+     *
+     * @since 7.4
+     */
+    protected void doBeforeComplete() {
+
+    }
+
     /**
      * Callback method called after successful completion and before notifying downstream consumers.
      *
@@ -111,9 +122,11 @@ public void cancel() {
         Object[] consumers = ONCOMPLETE.get(this);
 
         if (consumers != COMPLETE && ONCOMPLETE.compareAndSet(this, consumers, COMPLETE)) {
+            CancellationException exception = new CancellationException();
+
+            doBeforeError(exception);
 
             command.cancel();
-            CancellationException exception = new CancellationException();
             doOnError(exception);
             notifyBiConsumer(consumers, exception);
         }
@@ -128,6 +141,7 @@ public boolean completeExceptionally(Throwable throwable) {
         boolean result = false;
         if (consumers != COMPLETE && ONCOMPLETE.compareAndSet(this, consumers, COMPLETE)) {
 
+            doBeforeError(throwable);
             result = command.completeExceptionally(throwable);
             doOnError(throwable);
             notifyBiConsumer(consumers, throwable);
@@ -136,6 +150,16 @@ public boolean completeExceptionally(Throwable throwable) {
         return result;
     }
 
+    /**
+     * Callback method called after error completion and before notifying downstream consumers.
+     *
+     * @param throwable
+     * @since 7.4
+     */
+    protected void doBeforeError(Throwable throwable) {
+
+    }
+
     /**
      * Callback method called after error completion and before notifying downstream consumers.
      *