Trace Lettuce batch writes, stop orphan reactive traces and unfinished spans of done commands

CHANGES.md

@@ -9,6 +9,10 @@ Release Notes.
   `Statement.RETURN_GENERATED_KEYS` or column names (apache/skywalking#14138).
 * Fix the `NullPointerException` thrown by the Lettuce plugin and the lost trace context of reactive commands when
   Lettuce tracing is enabled, and when a reactive command is cancelled before it is written (apache/skywalking#14178).
+* Trace the commands that Lettuce writes as a batch, such as `StatefulConnection#dispatch(Collection)`, with one
+  `Lettuce/BATCH_WRITE` span. Stop the Lettuce plugin from starting a trace of its own for every reactive command
+  created without an active trace, and from creating a span that is never finished for a command that is already
+  cancelled or timed out when it is written.
 
 All issues and pull requests are [here](https://github.com/apache/skywalking/milestone/267?closed=1)
 

apm-sniffer/apm-sdk-plugin/lettuce-plugins/lettuce-5.x-6.4.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/RedisReactiveCreatePublisherMethodInterceptorV5.java

@@ -39,7 +39,7 @@
  *
  * <p>If the Reactor Context does not already contain a snapshot, this interceptor
  * captures the current active context and writes it into the subscriber context
- * as a fallback propagation mechanism.</p>
+ * as a fallback propagation mechanism. Nothing is captured when no context is active.</p>
  */
 public class RedisReactiveCreatePublisherMethodInterceptorV5 implements InstanceMethodsAroundInterceptorV2 {
 
@@ -56,6 +56,12 @@ public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allA
         if (!(ret instanceof Mono) && !(ret instanceof Flux)) {
             return ret;
         }
+        // Without an active trace there is nothing to capture, and a local span would start a trace of its own.
+        // A snapshot put into the Reactor Context by another plugin, such as spring-webflux, is still continued
+        // when the command is written.
+        if (!ContextManager.isActive()) {
+            return ret;
+        }
         final AbstractSpan localSpan = ContextManager.createLocalSpan("Lettuce/Reactive/" + method.getName());
         localSpan.setComponent(ComponentsDefine.LETTUCE);
         SpanLayer.asCache(localSpan);

apm-sniffer/apm-sdk-plugin/lettuce-plugins/lettuce-5.x-6.4.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/define/RedisChannelWriterInstrumentationV5.java

@@ -18,6 +18,7 @@
 
 package org.apache.skywalking.apm.plugin.lettuce.v5.define;
 
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import net.bytebuddy.description.method.MethodDescription;
@@ -50,8 +51,8 @@ public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
             new InstanceMethodsInterceptPoint() {
                 @Override
                 public ElementMatcher<MethodDescription> getMethodsMatcher() {
-                    return named("write").and(takesArgument(0, named("io.lettuce.core.protocol.RedisCommand")).or(
-                        takesArgument(0, List.class)));
+                    return named("write").and(takesArgument(0, named("io.lettuce.core.protocol.RedisCommand"))
+                        .or(takesArgument(0, Collection.class)));
                 }
 
                 @Override

apm-sniffer/apm-sdk-plugin/lettuce-plugins/lettuce-6.5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v65/RedisReactiveCreatePublisherMethodInterceptorV65.java

@@ -38,7 +38,7 @@
  *
  * <p>If the Reactor Context does not already contain a snapshot, this interceptor
  * captures the current active context and writes it into the subscriber context
- * as a fallback propagation mechanism.</p>
+ * as a fallback propagation mechanism. Nothing is captured when no context is active.</p>
  */
 public class RedisReactiveCreatePublisherMethodInterceptorV65 implements InstanceMethodsAroundInterceptorV2 {
 
@@ -55,6 +55,12 @@ public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allA
         if (!(ret instanceof Mono) && !(ret instanceof Flux)) {
             return ret;
         }
+        // Without an active trace there is nothing to capture, and a local span would start a trace of its own.
+        // A snapshot put into the Reactor Context by another plugin, such as spring-webflux, is still continued
+        // when the command is written.
+        if (!ContextManager.isActive()) {
+            return ret;
+        }
 
         final AbstractSpan localSpan = ContextManager.createLocalSpan("Lettuce/Reactive/" + method.getName());
         localSpan.setComponent(ComponentsDefine.LETTUCE);

apm-sniffer/apm-sdk-plugin/lettuce-plugins/lettuce-6.5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v65/define/RedisChannelWriterInstrumentationV65.java

@@ -18,6 +18,7 @@
 
 package org.apache.skywalking.apm.plugin.lettuce.v65.define;
 
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import net.bytebuddy.description.method.MethodDescription;
@@ -50,8 +51,8 @@ public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
             new InstanceMethodsInterceptPoint() {
                 @Override
                 public ElementMatcher<MethodDescription> getMethodsMatcher() {
-                    return named("write").and(takesArgument(0, named("io.lettuce.core.protocol.RedisCommand")).or(
-                        takesArgument(0, List.class)));
+                    return named("write").and(takesArgument(0, named("io.lettuce.core.protocol.RedisCommand"))
+                        .or(takesArgument(0, Collection.class)));
                 }
 
                 @Override

apm-sniffer/apm-sdk-plugin/lettuce-plugins/lettuce-common/src/main/java/org/apache/skywalking/apm/plugin/lettuce/common/RedisChannelWriterInterceptor.java

@@ -37,7 +37,6 @@
 import java.lang.reflect.Method;
 import java.nio.ByteBuffer;
 import java.util.Collection;
-import java.util.List;
 import java.util.Optional;
 
 @SuppressWarnings("unchecked")
@@ -51,10 +50,17 @@ public class RedisChannelWriterInterceptor implements InstanceMethodsAroundInter
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, MethodInterceptResult result) {
         String peer = (String) objInst.getSkyWalkingDynamicField();
-        RedisCommand<?, ?, ?> spanCarrierCommand = getSpanCarrierCommand(allArguments[0]);
+        RedisCommand<?, ?, ?> writtenCommand = getWrittenCommand(allArguments[0]);
+        RedisCommand<?, ?, ?> spanCarrierCommand = getSpanCarrierCommand(writtenCommand);
         if (spanCarrierCommand == null) {
             return;
         }
+        // A command that is already done, such as one cancelled or timed out while its write was delayed, is dropped
+        // by the CommandHandler, or gets no terminal call that reaches the span carrier, so its span would never be
+        // finished.
+        if (isDone(writtenCommand, spanCarrierCommand)) {
+            return;
+        }
         EnhancedInstance enhancedCommand = (EnhancedInstance) spanCarrierCommand;
 
         RedisCommandEnhanceInfo redisCommandEnhanceInfo = (RedisCommandEnhanceInfo) enhancedCommand.getSkyWalkingDynamicField();
@@ -64,10 +70,11 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
         }
 
         // command has been handle by another channel writer (cluster or sentinel case)
-        if (redisCommandEnhanceInfo.getSpan() != null) {
+        AbstractSpan existingSpan = redisCommandEnhanceInfo.getSpan();
+        if (existingSpan != null) {
             //set peer in last channel writer (delegate)
             if (peer != null) {
-                redisCommandEnhanceInfo.getSpan().setPeer(peer);
+                existingSpan.setPeer(peer);
             }
             return;
         }
@@ -112,6 +119,17 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
         }
 
         enhancedCommand.setSkyWalkingDynamicField(redisCommandEnhanceInfo.setSpan(span));
+
+        // The command may have been cancelled or timed out before the span was set, when its terminal interceptor
+        // had no span to finish.
+        if (isDone(writtenCommand, spanCarrierCommand)) {
+            AbstractSpan unfinishedSpan = redisCommandEnhanceInfo.takeSpan();
+            if (unfinishedSpan != null) {
+                enhancedCommand.setSkyWalkingDynamicField(null);
+                unfinishedSpan.errorOccurred();
+                unfinishedSpan.asyncFinish();
+            }
+        }
     }
 
     private String getArgsKey(RedisCommand<?, ?, ?> redisCommand) {
@@ -141,34 +159,56 @@ public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allA
     @Override
     public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Throwable t) {
         RedisCommand<?, ?, ?> redisCommand = getSpanCarrierCommand(allArguments[0]);
-        if (redisCommand instanceof EnhancedInstance && ((EnhancedInstance) redisCommand).getSkyWalkingDynamicField() != null) {
+        if (redisCommand instanceof EnhancedInstance) {
             EnhancedInstance enhancedRedisCommand = (EnhancedInstance) redisCommand;
-            AbstractSpan abstractSpan = ((RedisCommandEnhanceInfo) enhancedRedisCommand.getSkyWalkingDynamicField()).getSpan();
-            enhancedRedisCommand.setSkyWalkingDynamicField(null);
-            abstractSpan.log(t);
-            abstractSpan.asyncFinish();
+            AbstractSpan abstractSpan = RedisCommandEnhanceInfo.takeSpan(enhancedRedisCommand);
+            if (abstractSpan != null) {
+                enhancedRedisCommand.setSkyWalkingDynamicField(null);
+                abstractSpan.log(t);
+                abstractSpan.asyncFinish();
+            }
         }
     }
 
     protected String getCommandName(ProtocolKeyword protocol) {
         return "UNKNOWN";
     }
 
-    static RedisCommand<?, ?, ?> getSpanCarrierCommand(Object o) {
-        RedisCommand<?, ?, ?> command = null;
+    /**
+     * @return the written command. For a batch write, it is the last command that is not done, which carries the span
+     * of the whole batch, because a done command is dropped while the others are still sent. A command of the batch
+     * that is cancelled while the batch is being written can still leave it without a span, or with two.
+     */
+    private static RedisCommand<?, ?, ?> getWrittenCommand(Object o) {
         if (o instanceof RedisCommand) {
-            command = (RedisCommand<?, ?, ?>) o;
-        } else if (o instanceof List) {
-            List<?> list = (List<?>) o;
-            command = list.isEmpty() ? null : (RedisCommand<?, ?, ?>) list.get(list.size() - 1);
-        } else if (o instanceof Collection) {
-            Collection<RedisCommand<?, ?, ?>> redisCommands = (Collection<RedisCommand<?, ?, ?>>) o;
-            RedisCommand<?, ?, ?> last = null;
-            for (RedisCommand<?, ?, ?> redisCommand : redisCommands) {
-                last = redisCommand;
+            return (RedisCommand<?, ?, ?>) o;
+        }
+        if (!(o instanceof Collection)) {
+            return null;
+        }
+        RedisCommand<?, ?, ?> lastCommand = null;
+        RedisCommand<?, ?, ?> lastCommandNotDone = null;
+        boolean anyCommandNotDone = false;
+        for (RedisCommand<?, ?, ?> redisCommand : (Collection<RedisCommand<?, ?, ?>>) o) {
+            lastCommand = redisCommand;
+            if (redisCommand == null || !isDone(redisCommand, getSpanCarrierCommand(redisCommand))) {
+                lastCommandNotDone = redisCommand;
+                anyCommandNotDone = true;
             }
-            command = last;
         }
+        return anyCommandNotDone ? lastCommandNotDone : lastCommand;
+    }
+
+    /**
+     * The CommandHandler checks the written command, while a wrapper of the span carrier, such as an AsyncCommand
+     * whose cancellation cancels the carrier first, or a ClusterCommand, can look not done after the carrier is.
+     */
+    private static boolean isDone(RedisCommand<?, ?, ?> writtenCommand, RedisCommand<?, ?, ?> spanCarrierCommand) {
+        return writtenCommand.isDone() || spanCarrierCommand != null && spanCarrierCommand.isDone();
+    }
+
+    static RedisCommand<?, ?, ?> getSpanCarrierCommand(Object o) {
+        RedisCommand<?, ?, ?> command = getWrittenCommand(o);
         if (command instanceof DecoratedCommand) {
             while (command instanceof DecoratedCommand) {
                 DecoratedCommand<?, ?, ?> wrapper = (DecoratedCommand<?, ?, ?>) command;

apm-sniffer/apm-sdk-plugin/lettuce-plugins/lettuce-common/src/main/java/org/apache/skywalking/apm/plugin/lettuce/common/RedisCommandCancelMethodInterceptor.java

@@ -36,7 +36,7 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
 
     @Override
     public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Object ret) {
-        AbstractSpan span = RedisCommandEnhanceInfo.getSpan(objInst);
+        AbstractSpan span = RedisCommandEnhanceInfo.takeSpan(objInst);
         if (span != null) {
             span.errorOccurred();
             span.tag(new StringTag(CANCEL_SIGNAL_TAG), COMMAND_CANCEL_VALUE);

apm-sniffer/apm-sdk-plugin/lettuce-plugins/lettuce-common/src/main/java/org/apache/skywalking/apm/plugin/lettuce/common/RedisCommandCompleteExceptionallyMethodInterceptor.java

@@ -33,7 +33,7 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
 
     @Override
     public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Object ret) {
-        AbstractSpan span = RedisCommandEnhanceInfo.getSpan(objInst);
+        AbstractSpan span = RedisCommandEnhanceInfo.takeSpan(objInst);
         if (span != null) {
             Throwable t = (Throwable) allArguments[0];
             span.log(t);

apm-sniffer/apm-sdk-plugin/lettuce-plugins/lettuce-common/src/main/java/org/apache/skywalking/apm/plugin/lettuce/common/RedisCommandCompleteMethodInterceptor.java

@@ -33,7 +33,7 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
 
     @Override
     public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Object ret) {
-        AbstractSpan span = RedisCommandEnhanceInfo.getSpan(objInst);
+        AbstractSpan span = RedisCommandEnhanceInfo.takeSpan(objInst);
         if (span != null) {
             span.asyncFinish();
             objInst.setSkyWalkingDynamicField(null);

apm-sniffer/apm-sdk-plugin/lettuce-plugins/lettuce-common/src/main/java/org/apache/skywalking/apm/plugin/lettuce/common/RedisCommandEnhanceInfo.java

@@ -18,6 +18,7 @@
 
 package org.apache.skywalking.apm.plugin.lettuce.common;
 
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
 import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
@@ -37,11 +38,11 @@
  */
 class RedisCommandEnhanceInfo {
 
-    private AbstractSpan span;
+    private final AtomicReference<AbstractSpan> span = new AtomicReference<>();
     private ContextSnapshot snapshot;
 
     public AbstractSpan getSpan() {
-        return span;
+        return span.get();
     }
 
     /**
@@ -53,8 +54,22 @@ static AbstractSpan getSpan(EnhancedInstance command) {
         return enhanceInfo == null ? null : enhanceInfo.getSpan();
     }
 
+    /**
+     * @return the span of the command, which the caller must finish, or null if the command has no span or another
+     * caller has taken it. A command can be cancelled while its write is traced, so more than one caller may try to
+     * finish the span, and only the first one gets it.
+     */
+    static AbstractSpan takeSpan(EnhancedInstance command) {
+        RedisCommandEnhanceInfo enhanceInfo = (RedisCommandEnhanceInfo) command.getSkyWalkingDynamicField();
+        return enhanceInfo == null ? null : enhanceInfo.takeSpan();
+    }
+
+    AbstractSpan takeSpan() {
+        return span.getAndSet(null);
+    }
+
     public RedisCommandEnhanceInfo setSpan(AbstractSpan span) {
-        this.span = span;
+        this.span.set(span);
         return this;
     }
 

apm-sniffer/apm-sdk-plugin/lettuce-plugins/lettuce-common/src/test/java/org/apache/skywalking/apm/plugin/lettuce/common/RedisChannelWriterInterceptorTest.java

@@ -111,6 +111,20 @@ public void setSkyWalkingDynamicField(Object value) {
         }
     }
 
+    private static class CancelledWhileSpanIsCreatedCommand<K, V, T> extends MockRedisCommand<K, V, T> {
+        public CancelledWhileSpanIsCreatedCommand(ProtocolKeyword type, CommandOutput<K, V, T> output) {
+            super(type, output);
+        }
+
+        @Override
+        public void setSkyWalkingDynamicField(Object value) {
+            super.setSkyWalkingDynamicField(value);
+            if (value != null) {
+                cancel();
+            }
+        }
+    }
+
     @Before
     public void setUp() {
         LettucePluginConfig.Plugin.Lettuce.TRACE_REDIS_PARAMETERS = true;
@@ -141,6 +155,45 @@ public void testInterceptor() {
         assertThat(SpanHelper.getPeer(spans.get(0)), is(PEER));
     }
 
+    @Test
+    public void testDoneCommandHasNoSpan() {
+        MockRedisCommand<?, ?, ?> redisCommand = new MockRedisCommand<>(CommandType.SET, null);
+        redisCommand.cancel();
+        interceptor.beforeMethod(mockRedisChannelWriterInstance, null, new Object[]{redisCommand}, null, null);
+
+        assertNull(redisCommand.getSkyWalkingDynamicField());
+        assertThat(segmentStorage.getTraceSegments().size(), is(0));
+    }
+
+    @Test
+    public void testCommandCancelledWhileSpanIsCreatedHasItsSpanFinished() {
+        CancelledWhileSpanIsCreatedCommand<?, ?, ?> redisCommand = new CancelledWhileSpanIsCreatedCommand<>(CommandType.SET, null);
+        // The cancellation interceptor found no span, and the command is done by the time the writer has set the span.
+        new RedisCommandCancelMethodInterceptor().afterMethod(redisCommand, null, null, null, null);
+        interceptor.beforeMethod(mockRedisChannelWriterInstance, null, new Object[]{redisCommand}, null, null);
+
+        assertNull(redisCommand.getSkyWalkingDynamicField());
+        assertThat(segmentStorage.getTraceSegments().size(), is(1));
+        List<AbstractTracingSpan> spans = SegmentHelper.getSpans(segmentStorage.getTraceSegments().get(0));
+        assertThat(spans.size(), is(1));
+        assertThat(SpanHelper.getErrorOccurred(spans.get(0)), is(true));
+    }
+
+    @Test
+    public void testSpanIsTakenByOneTerminalInterceptor() {
+        MockRedisCommand<?, ?, ?> redisCommand = new MockRedisCommand<>(CommandType.SET, null);
+        interceptor.beforeMethod(mockRedisChannelWriterInstance, null, new Object[]{redisCommand}, null, null);
+        Object enhanceInfo = redisCommand.getSkyWalkingDynamicField();
+
+        redisCommandCompleteMethodInterceptor.afterMethod(redisCommand, null, null, null, null);
+        // A concurrent cancellation read the enhance info before the completion cleared it.
+        redisCommand.setSkyWalkingDynamicField(enhanceInfo);
+        new RedisCommandCancelMethodInterceptor().afterMethod(redisCommand, null, null, null, null);
+
+        assertThat(segmentStorage.getTraceSegments().size(), is(1));
+        assertThat(SpanHelper.getErrorOccurred(SegmentHelper.getSpans(segmentStorage.getTraceSegments().get(0)).get(0)), is(false));
+    }
+
     @Test
     public void testGetSpanCarrierCommand() throws Exception {
         Command<?, ?, ?> command = new Command<>(CommandType.SET, null, null);
@@ -159,6 +212,13 @@ public void testGetSpanCarrierCommand() throws Exception {
         RedisCommand<?, ?, ?> last =
             (RedisCommand<?, ?, ?>) getSpanCarrierCommandMethod.invoke(interceptor, list);
         assertEquals(command, last);
+        Command<?, ?, ?> cancelled = new Command<>(CommandType.SET, null, null);
+        cancelled.cancel();
+        list.add(cancelled);
+        RedisCommand<?, ?, ?> lastNotDone =
+            (RedisCommand<?, ?, ?>) getSpanCarrierCommandMethod.invoke(interceptor, list);
+        assertEquals(command, lastNotDone);
+        list.remove(cancelled);
         RedisCommand<?, ?, ?> nullValue1 =
             (RedisCommand<?, ?, ?>) getSpanCarrierCommandMethod.invoke(interceptor, (Object) null);
         assertNull(nullValue1);