Propagate context to lettuce callbacks

instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceAsyncCommandInstrumentation.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.lettuce.v4_0;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.lambdaworks.redis.protocol.AsyncCommand;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class LettuceAsyncCommandInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.lambdaworks.redis.protocol.AsyncCommand");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isConstructor(), LettuceAsyncCommandInstrumentation.class.getName() + "$SaveContextAdvice");
+    transformer.applyAdviceToMethod(
+        named("complete").or(named("completeExceptionally")).or(named("cancel")),
+        LettuceAsyncCommandInstrumentation.class.getName() + "$RestoreContextAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class SaveContextAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void saveContext(@Advice.This AsyncCommand<?, ?, ?> asyncCommand) {
+      Context context = Java8BytecodeBridge.currentContext();
+      // get the context that submitted this command and attach it, it will be used to run callbacks
+      context = context.get(LettuceSingletons.COMMAND_CONTEXT_KEY);
+      InstrumentationContext.get(AsyncCommand.class, Context.class).put(asyncCommand, context);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class RestoreContextAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.This AsyncCommand<?, ?, ?> asyncCommand, @Advice.Local("otelScope") Scope scope) {
+      Context context =
+          InstrumentationContext.get(AsyncCommand.class, Context.class).get(asyncCommand);
+      scope = context.makeCurrent();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.Local("otelScope") Scope scope) {
+      scope.close();
+    }
+  }
+}

instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceAsyncCommandsInstrumentation.java

@@ -45,7 +45,10 @@ public static void onEnter(
         @Advice.Argument(0) RedisCommand<?, ?, ?> command,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
-      context = instrumenter().start(currentContext(), command);
+      Context parentContext = currentContext();
+      context = instrumenter().start(parentContext, command);
+      // remember the context that called dispatch, it is used in LettuceAsyncCommandInstrumentation
+      context = context.with(LettuceSingletons.COMMAND_CONTEXT_KEY, parentContext);
       scope = context.makeCurrent();
     }
 

instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceInstrumentationModule.java

@@ -20,6 +20,9 @@ public LettuceInstrumentationModule() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return asList(new LettuceConnectInstrumentation(), new LettuceAsyncCommandsInstrumentation());
+    return asList(
+        new LettuceAsyncCommandInstrumentation(),
+        new LettuceAsyncCommandsInstrumentation(),
+        new LettuceConnectInstrumentation());
   }
 }

instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceSingletons.java

@@ -8,6 +8,8 @@
 import com.lambdaworks.redis.RedisURI;
 import com.lambdaworks.redis.protocol.RedisCommand;
 import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
@@ -19,9 +21,11 @@ public final class LettuceSingletons {
   private static final String INSTRUMENTATION_NAME = "io.opentelemetry.lettuce-4.0";
 
   private static final Instrumenter<RedisCommand<?, ?, ?>, Void> INSTRUMENTER;
-
   private static final Instrumenter<RedisURI, Void> CONNECT_INSTRUMENTER;
 
+  public static final ContextKey<Context> COMMAND_CONTEXT_KEY =
+      ContextKey.named("opentelemetry-lettuce-v4_0-context-key");
+
   static {
     DbAttributesExtractor<RedisCommand<?, ?, ?>, Void> attributesExtractor =
         new LettuceDbAttributesExtractor();

instrumentation/lettuce/lettuce-4.0/javaagent/src/test/groovy/LettuceAsyncClientTest.groovy

@@ -4,6 +4,7 @@
  */
 
 import static io.opentelemetry.api.trace.SpanKind.CLIENT
+import static io.opentelemetry.api.trace.SpanKind.INTERNAL
 import static io.opentelemetry.api.trace.StatusCode.ERROR
 
 import com.lambdaworks.redis.ClientOptions
@@ -180,29 +181,44 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
     Consumer<String> consumer = new Consumer<String>() {
       @Override
       void accept(String res) {
-        conds.evaluate {
-          assert res == "TESTVAL"
+        runWithSpan("callback") {
+          conds.evaluate {
+            assert res == "TESTVAL"
+          }
         }
       }
     }
 
     when:
-    RedisFuture<String> redisFuture = asyncCommands.get("TESTKEY")
-    redisFuture.thenAccept(consumer)
+    runWithSpan("parent") {
+      RedisFuture<String> redisFuture = asyncCommands.get("TESTKEY")
+      redisFuture.thenAccept(consumer)
+    }
 
     then:
     conds.await()
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
         span(0) {
+          name "parent"
+          kind INTERNAL
+          hasNoParent()
+        }
+        span(1) {
           name "GET"
           kind CLIENT
+          childOf(span(0))
           attributes {
             "${SemanticAttributes.DB_SYSTEM.key}" "redis"
             "${SemanticAttributes.DB_OPERATION.key}" "GET"
             "${SemanticAttributes.DB_STATEMENT.key}" "GET"
           }
         }
+        span(2) {
+          name "callback"
+          kind INTERNAL
+          childOf(span(0))
+        }
       }
     }
   }
@@ -216,40 +232,62 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
     BiFunction<String, Throwable, String> firstStage = new BiFunction<String, Throwable, String>() {
       @Override
       String apply(String res, Throwable throwable) {
-        conds.evaluate {
-          assert res == null
-          assert throwable == null
+        runWithSpan("callback1") {
+          conds.evaluate {
+            assert res == null
+            assert throwable == null
+          }
         }
         return (res == null ? successStr : res)
       }
     }
     Function<String, Object> secondStage = new Function<String, Object>() {
       @Override
       Object apply(String input) {
-        conds.evaluate {
-          assert input == successStr
+        runWithSpan("callback2") {
+          conds.evaluate {
+            assert input == successStr
+          }
         }
         return null
       }
     }
 
     when:
-    RedisFuture<String> redisFuture = asyncCommands.get("NON_EXISTENT_KEY")
-    redisFuture.handleAsync(firstStage).thenApply(secondStage)
+    runWithSpan("parent") {
+      RedisFuture<String> redisFuture = asyncCommands.get("NON_EXISTENT_KEY")
+      redisFuture.handle(firstStage).thenApply(secondStage)
+    }
 
     then:
     conds.await()
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 4) {
         span(0) {
+          name "parent"
+          kind INTERNAL
+          hasNoParent()
+        }
+        span(1) {
           name "GET"
           kind CLIENT
+          childOf(span(0))
           attributes {
             "${SemanticAttributes.DB_SYSTEM.key}" "redis"
             "${SemanticAttributes.DB_OPERATION.key}" "GET"
             "${SemanticAttributes.DB_STATEMENT.key}" "GET"
           }
         }
+        span(2) {
+          name "callback1"
+          kind INTERNAL
+          childOf(span(0))
+        }
+        span(3) {
+          name "callback2"
+          kind INTERNAL
+          childOf(span(0))
+        }
       }
     }
   }
@@ -260,29 +298,44 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
     BiConsumer<String, Throwable> biConsumer = new BiConsumer<String, Throwable>() {
       @Override
       void accept(String keyRetrieved, Throwable throwable) {
-        conds.evaluate {
-          assert keyRetrieved != null
+        runWithSpan("callback") {
+          conds.evaluate {
+            assert keyRetrieved != null
+          }
         }
       }
     }
 
     when:
-    RedisFuture<String> redisFuture = asyncCommands.randomkey()
-    redisFuture.whenCompleteAsync(biConsumer)
+    runWithSpan("parent") {
+      RedisFuture<String> redisFuture = asyncCommands.randomkey()
+      redisFuture.whenCompleteAsync(biConsumer)
+    }
 
     then:
     conds.await()
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
         span(0) {
+          name "parent"
+          kind INTERNAL
+          hasNoParent()
+        }
+        span(1) {
           name "RANDOMKEY"
           kind CLIENT
+          childOf(span(0))
           attributes {
             "${SemanticAttributes.DB_SYSTEM.key}" "redis"
             "${SemanticAttributes.DB_OPERATION.key}" "RANDOMKEY"
             "${SemanticAttributes.DB_STATEMENT.key}" "RANDOMKEY"
           }
         }
+        span(2) {
+          name "callback"
+          kind INTERNAL
+          childOf(span(0))
+        }
       }
     }
   }
@@ -397,12 +450,16 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
     setup:
     asyncCommands.setAutoFlushCommands(false)
     def conds = new AsyncConditions()
-    RedisFuture redisFuture = asyncCommands.sadd("SKEY", "1", "2")
+    RedisFuture redisFuture = runWithSpan("parent") {
+      asyncCommands.sadd("SKEY", "1", "2")
+    }
     redisFuture.whenCompleteAsync({
       res, throwable ->
-        conds.evaluate {
-          assert throwable != null
-          assert throwable instanceof CancellationException
+        runWithSpan("callback") {
+          conds.evaluate {
+            assert throwable != null
+            assert throwable instanceof CancellationException
+          }
         }
     })
 
@@ -414,17 +471,28 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
     conds.await()
     cancelSuccess == true
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
         span(0) {
+          name "parent"
+          kind INTERNAL
+          hasNoParent()
+        }
+        span(1) {
           name "SADD"
           kind CLIENT
+          childOf(span(0))
           attributes {
             "${SemanticAttributes.DB_SYSTEM.key}" "redis"
             "${SemanticAttributes.DB_OPERATION.key}" "SADD"
             "${SemanticAttributes.DB_STATEMENT.key}" "SADD"
             "lettuce.command.cancelled" true
           }
         }
+        span(2) {
+          name "callback"
+          kind INTERNAL
+          childOf(span(0))
+        }
       }
     }
   }