Inspect Pub/Sub responses for interleaved messages #724

Lettuce now inspects Redis responses via PubSubCommandHandler and ReplayOutput whether a received response is a Pub/Sub message or whether the response belongs to a command on the protocol stack. Introspection is required as Redis responses may contain interleaved messages that do not belong to a command or may arrive before the command response.

Previously, interleaved messages could get used to complete commands on the protocol stack which causes a defunct protocol state.

src/main/java/io/lettuce/core/output/ReplayOutput.java

@@ -0,0 +1,197 @@
+/*
+ * Copyright 2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.lettuce.core.output;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+
+import io.lettuce.core.codec.RedisCodec;
+import io.lettuce.core.codec.StringCodec;
+
+/**
+ * Replayable {@link CommandOutput} capturing output signals to replay these on a target {@link CommandOutput}. Replay is useful
+ * when the response requires inspection prior to dispatching the actual output to a command target.
+ *
+ * @author Mark Paluch
+ * @since 5.0.3
+ */
+public class ReplayOutput<K, V> extends CommandOutput<K, V, List<ReplayOutput.Signal>> {
+
+    /**
+     * Initialize a new instance that encodes and decodes keys and values using the supplied codec.
+     */
+    public ReplayOutput() {
+        super((RedisCodec) StringCodec.ASCII, new ArrayList<>());
+    }
+
+    @Override
+    public void set(ByteBuffer bytes) {
+        output.add(new BulkString(bytes));
+    }
+
+    @Override
+    public void set(long integer) {
+        output.add(new Integer(integer));
+    }
+
+    @Override
+    public void setError(ByteBuffer error) {
+        error.mark();
+        output.add(new ErrorBytes(error));
+        error.reset();
+        super.setError(error);
+    }
+
+    @Override
+    public void setError(String error) {
+        output.add(new ErrorString(error));
+        super.setError(error);
+    }
+
+    @Override
+    public void complete(int depth) {
+        output.add(new Complete(depth));
+    }
+
+    @Override
+    public void multi(int count) {
+        output.add(new Multi(count));
+    }
+
+    /**
+     * Replay all captured signals on a {@link CommandOutput}.
+     *
+     * @param target the target {@link CommandOutput}.
+     */
+    public void replay(CommandOutput<?, ?, ?> target) {
+
+        for (Signal signal : output) {
+            signal.replay(target);
+        }
+    }
+
+    /**
+     * Encapsulates a replayable decoding signal.
+     */
+    public static abstract class Signal {
+
+        /**
+         * Replay the signal on a {@link CommandOutput}.
+         *
+         * @param target
+         */
+        protected abstract void replay(CommandOutput<?, ?, ?> target);
+    }
+
+    abstract static class BulkStringSupport extends Signal {
+
+        final ByteBuffer message;
+
+        BulkStringSupport(ByteBuffer message) {
+
+            if (message != null) {
+
+                // need to copy the buffer to prevent buffer lifecycle mismatch
+                this.message = ByteBuffer.allocate(message.remaining());
+                this.message.put(message);
+                this.message.rewind();
+            } else {
+                this.message = null;
+            }
+        }
+    }
+
+    public static class BulkString extends BulkStringSupport {
+
+        BulkString(ByteBuffer message) {
+            super(message);
+        }
+
+        @Override
+        protected void replay(CommandOutput<?, ?, ?> target) {
+            target.set(message);
+        }
+    }
+
+    static class Integer extends Signal {
+
+        final long message;
+
+        Integer(long message) {
+            this.message = message;
+        }
+
+        @Override
+        protected void replay(CommandOutput<?, ?, ?> target) {
+            target.set(message);
+        }
+    }
+
+    public static class ErrorBytes extends BulkStringSupport {
+
+        ErrorBytes(ByteBuffer message) {
+            super(message);
+        }
+
+        @Override
+        protected void replay(CommandOutput<?, ?, ?> target) {
+            target.setError(message);
+        }
+    }
+
+    static class ErrorString extends Signal {
+
+        final String message;
+
+        ErrorString(String message) {
+            this.message = message;
+        }
+
+        @Override
+        protected void replay(CommandOutput<?, ?, ?> target) {
+            target.setError(message);
+        }
+    }
+
+    static class Multi extends Signal {
+
+        final int count;
+
+        Multi(int count) {
+            this.count = count;
+        }
+
+        @Override
+        protected void replay(CommandOutput<?, ?, ?> target) {
+            target.multi(count);
+        }
+    }
+
+    static class Complete extends Signal {
+
+        final int depth;
+
+        public Complete(int depth) {
+            this.depth = depth;
+        }
+
+        @Override
+        protected void replay(CommandOutput<?, ?, ?> target) {
+            target.complete(depth);
+        }
+    }
+}

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

@@ -411,7 +411,7 @@ private void addToStack(RedisCommand<?, ?, ?> command, ChannelPromise promise) {
 
             if (command.getOutput() == null) {
                 // fire&forget commands are excluded from metrics
-                command.complete();
+                complete(command);
             }
 
             RedisCommand<?, ?, ?> redisCommand = potentiallyWrapLatencyCommand(command);
@@ -552,27 +552,55 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf buffer) throws Interrup
                 onProtectedMode(command.getOutput().getError());
             } else {
 
-                stack.poll();
+                if (canComplete(command)) {
+                    stack.poll();
 
-                try {
-                    command.complete();
-                } catch (Exception e) {
-                    logger.warn("{} Unexpected exception during request: {}", logPrefix, e.toString(), e);
+                    try {
+                        complete(command);
+                    } catch (Exception e) {
+                        logger.warn("{} Unexpected exception during request: {}", logPrefix, e.toString(), e);
+                    }
                 }
             }
 
-            afterComplete(ctx, command);
+            afterDecode(ctx, command);
         }
 
         if (buffer.refCnt() != 0) {
             buffer.discardReadBytes();
         }
     }
 
+    /**
+     * Decoding hook: Can the buffer be decoded to a command.
+     *
+     * @param buffer
+     * @return
+     */
     protected boolean canDecode(ByteBuf buffer) {
         return !stack.isEmpty() && buffer.isReadable();
     }
 
+    /**
+     * Decoding hook: Can the command be completed.
+     *
+     * @param command
+     * @return
+     */
+    protected boolean canComplete(RedisCommand<?, ?, ?> command) {
+        return true;
+    }
+
+    /**
+     * Decoding hook: Complete a command.
+     *
+     * @param command
+     * @see RedisCommand#complete()
+     */
+    protected void complete(RedisCommand<?, ?, ?> command) {
+        command.complete();
+    }
+
     private boolean decode(ChannelHandlerContext ctx, ByteBuf buffer, RedisCommand<?, ?, ?> command) {
 
         if (latencyMetricsEnabled && command instanceof WithLatency) {
@@ -596,7 +624,7 @@ private boolean decode(ChannelHandlerContext ctx, ByteBuf buffer, RedisCommand<?
 
     private boolean decode0(ChannelHandlerContext ctx, ByteBuf buffer, RedisCommand<?, ?, ?> command) {
 
-        if (!decode(buffer, command, command.getOutput())) {
+        if (!decode(buffer, command, getCommandOutput(command))) {
 
             if (command instanceof DemandAware.Sink) {
 
@@ -616,6 +644,17 @@ private boolean decode0(ChannelHandlerContext ctx, ByteBuf buffer, RedisCommand<
         return true;
     }
 
+    /**
+     * Decoding hook: Retrieve {@link CommandOutput} for {@link RedisCommand} decoding.
+     *
+     * @param command
+     * @return
+     * @see RedisCommand#getOutput()
+     */
+    protected CommandOutput<?, ?, ?> getCommandOutput(RedisCommand<?, ?, ?> command) {
+        return command.getOutput();
+    }
+
     protected boolean decode(ByteBuf buffer, CommandOutput<?, ?, ?> output) {
         return rsm.decode(buffer, output);
     }
@@ -682,7 +721,7 @@ private void onProtectedMode(String message) {
      * @param ctx
      * @param command
      */
-    protected void afterComplete(ChannelHandlerContext ctx, RedisCommand<?, ?, ?> command) {
+    protected void afterDecode(ChannelHandlerContext ctx, RedisCommand<?, ?, ?> command) {
     }
 
     private void recordLatency(WithLatency withLatency, ProtocolKeyword commandType) {