Fix NPE and lost trace context of reactive Lettuce commands when tracing is enabled

.github/workflows/e2e.yaml

@@ -55,12 +55,8 @@ jobs:
           SW_VERSION=$(./mvnw -q -DforceStdout -N org.apache.maven.plugins:maven-help-plugin:3.2.0:evaluate -Dexpression=project.version)
           ./mvnw -B package -f ./test/e2e/base/pom.xml -Dsw.version=${SW_VERSION}
           echo "::endgroup::"
-      - name: Setup go
-        uses: actions/setup-go@v2
-        with:
-          go-version: '1.24'
       - name: Run E2E Tests
-        uses: apache/skywalking-infra-e2e@8c21e43e241a32a54bdf8eeceb9099eb27e5e9b4
+        uses: apache/skywalking-infra-e2e@56ac8befeb5f910f2dc9339e65d07cd18953f90e
         with:
           e2e-file: ${{ matrix.case.path }}
       - uses: actions/upload-artifact@v4

CHANGES.md

@@ -7,6 +7,8 @@ Release Notes.
 
 * Fix the `NullPointerException` thrown by the `postgresql-8.x` plugin when executing statements prepared with
   `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).
 
 All issues and pull requests are [here](https://github.com/apache/skywalking/milestone/267?closed=1)
 

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

@@ -154,7 +154,7 @@ protected String getCommandName(ProtocolKeyword protocol) {
         return "UNKNOWN";
     }
 
-    private static RedisCommand<?, ?, ?> getSpanCarrierCommand(Object o) {
+    static RedisCommand<?, ?, ?> getSpanCarrierCommand(Object o) {
         RedisCommand<?, ?, ?> command = null;
         if (o instanceof RedisCommand) {
             command = (RedisCommand<?, ?, ?>) o;

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

@@ -36,8 +36,8 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
 
     @Override
     public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Object ret) {
-        if (objInst.getSkyWalkingDynamicField() != null) {
-            AbstractSpan span = ((RedisCommandEnhanceInfo) objInst.getSkyWalkingDynamicField()).getSpan();
+        AbstractSpan span = RedisCommandEnhanceInfo.getSpan(objInst);
+        if (span != null) {
             span.errorOccurred();
             span.tag(new StringTag(CANCEL_SIGNAL_TAG), COMMAND_CANCEL_VALUE);
             span.asyncFinish();
@@ -48,8 +48,8 @@ public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allA
 
     @Override
     public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Throwable t) {
-        if (objInst.getSkyWalkingDynamicField() != null) {
-            AbstractSpan span = ((RedisCommandEnhanceInfo) objInst.getSkyWalkingDynamicField()).getSpan();
+        AbstractSpan span = RedisCommandEnhanceInfo.getSpan(objInst);
+        if (span != null) {
             span.log(t);
         }
     }

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

@@ -33,9 +33,9 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
 
     @Override
     public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Object ret) {
-        if (objInst.getSkyWalkingDynamicField() != null) {
+        AbstractSpan span = RedisCommandEnhanceInfo.getSpan(objInst);
+        if (span != null) {
             Throwable t = (Throwable) allArguments[0];
-            AbstractSpan span = ((RedisCommandEnhanceInfo) objInst.getSkyWalkingDynamicField()).getSpan();
             span.log(t);
             span.asyncFinish();
             objInst.setSkyWalkingDynamicField(null);
@@ -45,8 +45,8 @@ public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allA
 
     @Override
     public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Throwable t) {
-        if (objInst.getSkyWalkingDynamicField() != null) {
-            AbstractSpan span = ((RedisCommandEnhanceInfo) objInst.getSkyWalkingDynamicField()).getSpan();
+        AbstractSpan span = RedisCommandEnhanceInfo.getSpan(objInst);
+        if (span != null) {
             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,8 +33,8 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
 
     @Override
     public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Object ret) {
-        if (objInst.getSkyWalkingDynamicField() != null) {
-            AbstractSpan span = ((RedisCommandEnhanceInfo) objInst.getSkyWalkingDynamicField()).getSpan();
+        AbstractSpan span = RedisCommandEnhanceInfo.getSpan(objInst);
+        if (span != null) {
             span.asyncFinish();
             objInst.setSkyWalkingDynamicField(null);
         }
@@ -44,8 +44,8 @@ public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allA
     @Override
     public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
                                       Class<?>[] argumentsTypes, Throwable t) {
-        if (objInst.getSkyWalkingDynamicField() != null) {
-            AbstractSpan span = ((RedisCommandEnhanceInfo) objInst.getSkyWalkingDynamicField()).getSpan();
+        AbstractSpan span = RedisCommandEnhanceInfo.getSpan(objInst);
+        if (span != null) {
             span.log(t);
         }
     }

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

@@ -20,6 +20,7 @@
 
 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;
 
 /**
  * RedisCommandEnhanceInfo holds SkyWalking tracing data for Lettuce commands
@@ -43,6 +44,15 @@ public AbstractSpan getSpan() {
         return span;
     }
 
+    /**
+     * @return the span of the command, or null if the command has no span. A reactive command only carries the
+     * snapshot until it is written, so a command cancelled before that has no span to finish.
+     */
+    static AbstractSpan getSpan(EnhancedInstance command) {
+        RedisCommandEnhanceInfo enhanceInfo = (RedisCommandEnhanceInfo) command.getSkyWalkingDynamicField();
+        return enhanceInfo == null ? null : enhanceInfo.getSpan();
+    }
+
     public RedisCommandEnhanceInfo setSpan(AbstractSpan span) {
         this.span = span;
         return this;

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

@@ -25,15 +25,17 @@
  * Interceptor for RedisSubscription constructor.
  * <p>
  * This interceptor captures the {@link io.lettuce.core.protocol.RedisCommand} instance
- * at subscription construction time and stores it into SkyWalking dynamic field.
+ * at subscription construction time and stores it into SkyWalking dynamic field. A decorated
+ * command is unwrapped to the command that carries the span, so that the snapshot is found
+ * by {@link RedisChannelWriterInterceptor}.
  * </p>
  */
 public class RedisSubscriptionConstructorInterceptor implements InstanceConstructorInterceptor {
 
     @Override
     public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
         // allArguments[1] is the RedisCommand passed to the RedisSubscription constructor
-        objInst.setSkyWalkingDynamicField(allArguments[1]);
+        objInst.setSkyWalkingDynamicField(RedisChannelWriterInterceptor.getSpanCarrierCommand(allArguments[1]));
     }
 
 }

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

@@ -0,0 +1,227 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.plugin.lettuce.common;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.nullValue;
+import static org.hamcrest.CoreMatchers.sameInstance;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import io.lettuce.core.protocol.Command;
+import io.lettuce.core.protocol.CommandType;
+import io.lettuce.core.protocol.RedisCommand;
+import io.lettuce.core.protocol.TracedCommand;
+import io.lettuce.core.tracing.TraceContext;
+import java.util.List;
+import org.apache.skywalking.apm.agent.core.context.ContextManager;
+import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
+import org.apache.skywalking.apm.agent.core.context.trace.AbstractTracingSpan;
+import org.apache.skywalking.apm.agent.core.context.trace.TraceSegment;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
+import org.apache.skywalking.apm.agent.test.helper.SegmentHelper;
+import org.apache.skywalking.apm.agent.test.tools.AgentServiceRule;
+import org.apache.skywalking.apm.agent.test.tools.SegmentStorage;
+import org.apache.skywalking.apm.agent.test.tools.SegmentStoragePoint;
+import org.apache.skywalking.apm.agent.test.tools.TracingSegmentRunner;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
+import reactor.util.context.Context;
+
+@RunWith(TracingSegmentRunner.class)
+public class RedisReactiveCommandSpanTest {
+
+    private static final String PEER = "192.168.1.12:6379";
+
+    @SegmentStoragePoint
+    private SegmentStorage segmentStorage;
+
+    @Rule
+    public AgentServiceRule serviceRule = new AgentServiceRule();
+
+    private final RedisCommandCompleteMethodInterceptor completeInterceptor = new RedisCommandCompleteMethodInterceptor();
+    private final RedisCommandCancelMethodInterceptor cancelInterceptor = new RedisCommandCancelMethodInterceptor();
+    private final RedisCommandCompleteExceptionallyMethodInterceptor completeExceptionallyInterceptor =
+        new RedisCommandCompleteExceptionallyMethodInterceptor();
+
+    private MockCommand command;
+    private MockTracedCommand tracedCommand;
+    private MockInstance channelWriter;
+    private MockInstance subscription;
+
+    private static class MockInstance implements EnhancedInstance {
+        private Object object;
+
+        @Override
+        public Object getSkyWalkingDynamicField() {
+            return object;
+        }
+
+        @Override
+        public void setSkyWalkingDynamicField(Object value) {
+            this.object = value;
+        }
+    }
+
+    private static class MockCommand extends Command<String, String, String> implements EnhancedInstance {
+        private Object object;
+
+        MockCommand() {
+            super(CommandType.SET, null);
+        }
+
+        @Override
+        public Object getSkyWalkingDynamicField() {
+            return object;
+        }
+
+        @Override
+        public void setSkyWalkingDynamicField(Object value) {
+            this.object = value;
+        }
+    }
+
+    private static class MockTracedCommand extends TracedCommand<String, String, String> implements EnhancedInstance {
+        private Object object;
+
+        MockTracedCommand(MockCommand command) {
+            super(command, TraceContext.EMPTY);
+        }
+
+        @Override
+        public Object getSkyWalkingDynamicField() {
+            return object;
+        }
+
+        @Override
+        public void setSkyWalkingDynamicField(Object value) {
+            this.object = value;
+        }
+    }
+
+    private static class MockSubscriber implements CoreSubscriber<Object> {
+        private final ContextSnapshot snapshot;
+
+        MockSubscriber(ContextSnapshot snapshot) {
+            this.snapshot = snapshot;
+        }
+
+        @Override
+        public Context currentContext() {
+            return Context.of("SKYWALKING_CONTEXT_SNAPSHOT", snapshot);
+        }
+
+        @Override
+        public void onSubscribe(Subscription s) {
+        }
+
+        @Override
+        public void onNext(Object o) {
+        }
+
+        @Override
+        public void onError(Throwable t) {
+        }
+
+        @Override
+        public void onComplete() {
+        }
+    }
+
+    @Before
+    public void setUp() {
+        command = new MockCommand();
+        tracedCommand = new MockTracedCommand(command);
+        channelWriter = new MockInstance();
+        channelWriter.setSkyWalkingDynamicField(PEER);
+        subscription = new MockInstance();
+    }
+
+    @Test
+    public void snapshotOfDecoratedCommandIsContinuedOnSpanCarrierCommand() {
+        subscribe(tracedCommand);
+        write(tracedCommand);
+
+        assertThat(tracedCommand.getSkyWalkingDynamicField(), nullValue());
+
+        completeInterceptor.afterMethod(command, null, null, null, null);
+        // Mono#next cancels the subscription after the command has been completed.
+        completeInterceptor.afterMethod(tracedCommand, null, null, null, null);
+        cancelInterceptor.afterMethod(tracedCommand, null, null, null, null);
+        cancelInterceptor.afterMethod(command, null, null, null, null);
+
+        assertThat(command.getSkyWalkingDynamicField(), nullValue());
+        assertSpansOfContinuedCommand();
+    }
+
+    @Test
+    public void commandWithoutDecoratorStillContinuesSnapshot() {
+        subscribe(command);
+        assertThat(subscription.getSkyWalkingDynamicField(), sameInstance((Object) command));
+
+        write(command);
+        completeInterceptor.afterMethod(command, null, null, null, null);
+
+        assertSpansOfContinuedCommand();
+    }
+
+    @Test
+    public void commandCancelledBeforeWriteHasNoSpanToFinish() {
+        subscribe(command);
+        Throwable error = new IllegalStateException("cancelled");
+
+        cancelInterceptor.afterMethod(command, null, null, null, null);
+        cancelInterceptor.handleMethodException(command, null, null, null, error);
+        completeInterceptor.afterMethod(command, null, null, null, null);
+        completeInterceptor.handleMethodException(command, null, null, null, error);
+        completeExceptionallyInterceptor.afterMethod(command, null, new Object[]{error}, null, null);
+        completeExceptionallyInterceptor.handleMethodException(command, null, null, null, error);
+
+        assertThat(segmentStorage.getTraceSegments().size(), is(1));
+        assertThat(SegmentHelper.getSpans(segmentStorage.getTraceSegments().get(0)).size(), is(1));
+    }
+
+    private void subscribe(RedisCommand<?, ?, ?> redisCommand) {
+        ContextManager.createEntrySpan("/redis", null);
+        ContextSnapshot snapshot = ContextManager.capture();
+        ContextManager.stopSpan();
+
+        new RedisSubscriptionConstructorInterceptor().onConstruct(subscription, new Object[]{null, redisCommand});
+        new RedisSubscriptionSubscribeMethodInterceptor().beforeMethod(
+            subscription, null, new Object[]{new MockSubscriber(snapshot)}, null, null);
+    }
+
+    private void write(Object redisCommand) {
+        new RedisChannelWriterInterceptor().beforeMethod(channelWriter, null, new Object[]{redisCommand}, null, null);
+    }
+
+    private void assertSpansOfContinuedCommand() {
+        List<TraceSegment> segments = segmentStorage.getTraceSegments();
+        assertThat(segments.size(), is(2));
+        List<AbstractTracingSpan> spans = SegmentHelper.getSpans(segments.get(1));
+        assertThat(spans.size(), is(2));
+        assertThat(spans.get(0).getOperationName(), is("Lettuce/UNKNOWN"));
+        assertThat(spans.get(0).isExit(), is(true));
+        assertThat(spans.get(1).getOperationName(), is("RedisReactive/local"));
+        assertThat(segments.get(1).getRelatedGlobalTrace().getId(), is(segments.get(0).getRelatedGlobalTrace().getId()));
+    }
+}