Move command write deduplication to batch write #709

Lettuce now checks for command stack duplicates in CommandHandler.writeBatch(…) by using LinkedHashSet. Duplicates occur usually in batch submissions and the check in the single command path causes additional cost that isn't necessary in the majority of cases.

Using LinkedHashSet as intermediate collection reduces contains cost from O(N^2) to constant time.

src/main/java/com/lambdaworks/redis/protocol/CommandHandler.java

@@ -647,8 +647,25 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise)
             return;
         }
 
+        if (msg instanceof List) {
+
+            List<RedisCommand<K, V, ?>> batch = (List<RedisCommand<K, V, ?>>) msg;
+
+            if (batch.size() == 1) {
+
+                writeSingleCommand(ctx, batch.get(0), promise);
+                return;
+            }
+
+            writeBatch(ctx, batch, promise);
+            return;
+        }
+
         if (msg instanceof Collection) {
-            writeBatch(ctx, (Collection<RedisCommand<K, V, ?>>) msg, promise);
+
+            Collection<RedisCommand<K, V, ?>> batch = (Collection<RedisCommand<K, V, ?>>) msg;
+
+            writeBatch(ctx, batch, promise);
         }
     }
 
@@ -666,52 +683,35 @@ private void writeSingleCommand(ChannelHandlerContext ctx, RedisCommand<K, V, ?>
     private void writeBatch(ChannelHandlerContext ctx, Collection<RedisCommand<K, V, ?>> batch, ChannelPromise promise)
             throws Exception {
 
-        Collection<RedisCommand<K, V, ?>> toWrite = batch;
-        int commandsToWrite = 0;
+        Collection<RedisCommand<K, V, ?>> deduplicated = new LinkedHashSet<>(batch.size(), 1);
 
-        boolean cancelledCommands = false;
-        for (RedisCommand<?, ?, ?> command : batch) {
+        for (RedisCommand<K, V, ?> command : batch) {
 
-            if (!isWriteable(command)) {
-                cancelledCommands = true;
-                break;
+            if (isWriteable(command) && !deduplicated.add(command)) {
+                deduplicated.remove(command);
+                command.completeExceptionally(new RedisException(
+                        "Attempting to write duplicate command that is already enqueued: " + command));
             }
-
-            commandsToWrite++;
         }
 
         try {
-            validateWrite(commandsToWrite);
+            validateWrite(deduplicated.size());
         } catch (Exception e) {
 
-            for (RedisCommand<?, ?, ?> redisCommand : toWrite) {
+            for (RedisCommand<?, ?, ?> redisCommand : deduplicated) {
                 redisCommand.completeExceptionally(e);
             }
 
             promise.setFailure(e);
             return;
         }
 
-        if (cancelledCommands) {
-
-            toWrite = new ArrayList<>(batch.size());
-
-            for (RedisCommand<K, V, ?> command : batch) {
-
-                if (!isWriteable(command)) {
-                    continue;
-                }
-
-                toWrite.add(command);
-            }
-        }
-
-        for (RedisCommand<K, V, ?> command : toWrite) {
+        for (RedisCommand<K, V, ?> command : deduplicated) {
             addToStack(command, promise);
         }
 
-        if (!toWrite.isEmpty()) {
-            ctx.write(toWrite, promise);
+        if (!deduplicated.isEmpty()) {
+            ctx.write(deduplicated, promise);
         }
     }
 
@@ -732,10 +732,6 @@ private void addToStack(RedisCommand<K, V, ?> command, ChannelPromise promise) {
 
                 RedisCommand<K, V, ?> commandToUse = potentiallyWrapLatencyCommand(command);
 
-                if (stack.contains(command)) {
-                    throw new RedisException("Attempting to write duplicate command that is already enqueued: " + command);
-                }
-
                 if (promise.getClass() == VOID_PROMISE_CLASS) {
                     stack.add(commandToUse);
                 } else {

src/test/java/com/lambdaworks/redis/protocol/CommandHandlerTest.java

@@ -498,6 +498,18 @@ public void shouldCancelCommandOnQueueBatchFailure() throws Exception {
         verify(commandMock).completeExceptionally(exception);
     }
 
+    @Test
+    public void shouldFailOnDuplicateCommands() throws Exception {
+
+        Command<String, String, String> commandMock = mock(Command.class);
+
+        ChannelPromise channelPromise = new DefaultChannelPromise(channel, ImmediateEventExecutor.INSTANCE);
+        sut.write(context, Arrays.asList(commandMock, commandMock), channelPromise);
+
+        assertThat(stack).isEmpty();
+        verify(commandMock).completeExceptionally(any(RedisException.class));
+    }
+
     @Test
     public void shouldWriteActiveCommands() throws Exception {
 
@@ -520,17 +532,29 @@ public void shouldNotWriteCancelledCommandBatch() throws Exception {
     }
 
     @Test
-    public void shouldWriteActiveCommandsInBatch() throws Exception {
+    public void shouldWriteSingleActiveCommandsInBatch() throws Exception {
 
         when(promise.isSuccess()).thenReturn(true);
 
         List<Command<String, String, String>> commands = Arrays.asList(command);
         sut.write(context, commands, promise);
 
-        verify(context).write(commands, promise);
+        verify(context).write(command, promise);
         assertThat(stack).hasSize(1);
     }
 
+    @Test
+    public void shouldWriteActiveCommandsInBatch() throws Exception {
+
+        Command<String, String, String> anotherCommand = new Command<>(CommandType.APPEND,
+                new StatusOutput<>(StringCodec.UTF8), null);
+
+        List<Command<String, String, String>> commands = Arrays.asList(command, anotherCommand);
+        sut.write(context, commands, promise);
+
+        verify(context).write(any(Set.class), eq(promise));
+    }
+
     @Test
     @SuppressWarnings("unchecked")
     public void shouldWriteActiveCommandsInMixedBatch() throws Exception {
@@ -543,7 +567,7 @@ public void shouldWriteActiveCommandsInMixedBatch() throws Exception {
 
         sut.write(context, Arrays.asList(command, command2), promise);
 
-        ArgumentCaptor<List> captor = ArgumentCaptor.forClass(List.class);
+        ArgumentCaptor<Collection> captor = ArgumentCaptor.forClass(Collection.class);
         verify(context).write(captor.capture(), any());
 
         assertThat(captor.getValue()).containsOnly(command2);