Propagate context to lettuce callbacks (#3839)

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))
+ }
  }
  }
  }