Propagate context to lettuce callbacks
ADD
instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceAsyncCommandInstrumentation.java
+/*
+ * 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();
+ }
+ }
+}

MODIFY
instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceAsyncCommandsInstrumentation.java
instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceAsyncCommandsInstrumentation.java
@@ -45,7 +45,10 @@ public class LettuceAsyncCommandsInstrumentation implements TypeInstrumentation
@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();
}

MODIFY
instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceInstrumentationModule.java
instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceInstrumentationModule.java
@@ -20,6 +20,9 @@ public class LettuceInstrumentationModule extends InstrumentationModule {
@Override
public List<TypeInstrumentation> typeInstrumentations() {
- return asList(new LettuceConnectInstrumentation(), new LettuceAsyncCommandsInstrumentation());
+ return asList(
+ new LettuceAsyncCommandInstrumentation(),
+ new LettuceAsyncCommandsInstrumentation(),
+ new LettuceConnectInstrumentation());
}
}

MODIFY
instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceSingletons.java
instrumentation/lettuce/lettuce-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceSingletons.java
@@ -8,6 +8,8 @@ package io.opentelemetry.javaagent.instrumentation.lettuce.v4_0;
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();

MODIFY
instrumentation/lettuce/lettuce-4.0/javaagent/src/test/groovy/LettuceAsyncClientTest.groovy
instrumentation/lettuce/lettuce-4.0/javaagent/src/test/groovy/LettuceAsyncClientTest.groovy
*/
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) {
+ runWithSpan("callback") {
conds.evaluate {
assert res == "TESTVAL"
}
}
}
+ }
when:
+ 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) {
+ 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) {
+ runWithSpan("callback2") {
conds.evaluate {
assert input == successStr
}
+ }
return null
}
}
when:
+ runWithSpan("parent") {
RedisFuture<String> redisFuture = asyncCommands.get("NON_EXISTENT_KEY")
- redisFuture.handleAsync(firstStage).thenApply(secondStage)
+ 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) {
+ runWithSpan("callback") {
conds.evaluate {
assert keyRetrieved != null
}
}
}
+ }
when:
+ 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,13 +450,17 @@ 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 ->
+ runWithSpan("callback") {
conds.evaluate {
assert throwable != null
assert throwable instanceof CancellationException
}
+ }
})
when:
@@ -414,10 +471,16 @@ 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"
@@ -425,6 +488,11 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
"lettuce.command.cancelled" true
}
}
+ span(2) {
+ name "callback"
+ kind INTERNAL
+ childOf(span(0))
+ }
}
}
}

ADD
instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceAsyncCommandInstrumentation.java
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.lettuce.v5_0;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import io.lettuce.core.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("io.lettuce.core.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();
+ }
+ }
+}

MODIFY
instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceAsyncCommandsInstrumentation.java
instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceAsyncCommandsInstrumentation.java
@@ -47,7 +47,10 @@ public class LettuceAsyncCommandsInstrumentation implements TypeInstrumentation
@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();
}

MODIFY
instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceInstrumentationModule.java
instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceInstrumentationModule.java
@@ -30,6 +30,7 @@ public class LettuceInstrumentationModule extends InstrumentationModule {
@Override
public List<TypeInstrumentation> typeInstrumentations() {
return asList(
+ new LettuceAsyncCommandInstrumentation(),
new LettuceAsyncCommandsInstrumentation(),
new LettuceClientInstrumentation(),
new LettuceReactiveCommandsInstrumentation());

MODIFY
instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceSingletons.java
instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceSingletons.java
@@ -8,6 +8,8 @@ package io.opentelemetry.javaagent.instrumentation.lettuce.v5_0;
import io.lettuce.core.RedisURI;
import io.lettuce.core.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-5.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-v5_0-context-key");
+
static {
DbAttributesExtractor<RedisCommand<?, ?, ?>, Void> attributesExtractor =
new LettuceDbAttributesExtractor();

MODIFY
instrumentation/lettuce/lettuce-5.0/javaagent/src/test/groovy/LettuceAsyncClientTest.groovy
instrumentation/lettuce/lettuce-5.0/javaagent/src/test/groovy/LettuceAsyncClientTest.groovy
*/
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 io.lettuce.core.ClientOptions
@@ -185,29 +186,44 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
Consumer<String> consumer = new Consumer<String>() {
@Override
void accept(String res) {
+ runWithSpan("callback") {
conds.evaluate {
assert res == "TESTVAL"
}
}
}
+ }
when:
+ 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_STATEMENT.key" "GET TESTKEY"
"$SemanticAttributes.DB_OPERATION.key" "GET"
}
}
+ span(2) {
+ name "callback"
+ kind INTERNAL
+ childOf(span(0))
+ }
}
}
}
@@ -221,40 +237,62 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
BiFunction<String, Throwable, String> firstStage = new BiFunction<String, Throwable, String>() {
@Override
String apply(String res, Throwable throwable) {
+ 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) {
+ runWithSpan("callback2") {
conds.evaluate {
assert input == successStr
}
+ }
return null
}
}
when:
+ runWithSpan("parent") {
RedisFuture<String> redisFuture = asyncCommands.get("NON_EXISTENT_KEY")
redisFuture.handleAsync(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_STATEMENT.key" "GET NON_EXISTENT_KEY"
"$SemanticAttributes.DB_OPERATION.key" "GET"
}
}
+ span(2) {
+ name "callback1"
+ kind INTERNAL
+ childOf(span(0))
+ }
+ span(3) {
+ name "callback2"
+ kind INTERNAL
+ childOf(span(0))
+ }
}
}
}
@@ -265,29 +303,44 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
BiConsumer<String, Throwable> biConsumer = new BiConsumer<String, Throwable>() {
@Override
void accept(String keyRetrieved, Throwable throwable) {
+ runWithSpan("callback") {
conds.evaluate {
assert keyRetrieved != null
}
}
}
+ }
when:
+ 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_STATEMENT.key" "RANDOMKEY"
"$SemanticAttributes.DB_OPERATION.key" "RANDOMKEY"
}
}
+ span(2) {
+ name "callback"
+ kind INTERNAL
+ childOf(span(0))
+ }
}
}
}
@@ -401,13 +454,17 @@ 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 ->
+ runWithSpan("callback") {
conds.evaluate {
assert throwable != null
assert throwable instanceof CancellationException
}
+ }
})
when:
@@ -418,10 +475,16 @@ 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_STATEMENT.key" "SADD SKEY ? ?"
@@ -429,6 +492,11 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
"lettuce.command.cancelled" true
}
}
+ span(2) {
+ name "callback"
+ kind INTERNAL
+ childOf(span(0))
+ }
}
}
}

MODIFY
instrumentation/lettuce/lettuce-5.0/javaagent/src/test/groovy/LettuceReactiveClientTest.groovy
instrumentation/lettuce/lettuce-5.0/javaagent/src/test/groovy/LettuceReactiveClientTest.groovy
*/
import static io.opentelemetry.api.trace.SpanKind.CLIENT
+import static io.opentelemetry.api.trace.SpanKind.INTERNAL
import io.lettuce.core.ClientOptions
import io.lettuce.core.RedisClient
@@ -73,28 +74,43 @@ class LettuceReactiveClientTest extends AgentInstrumentationSpecification {
Consumer<String> consumer = new Consumer<String>() {
@Override
void accept(String res) {
+ runWithSpan("callback") {
conds.evaluate {
assert res == "OK"
}
}
}
+ }
when:
+ runWithSpan("parent") {
reactiveCommands.set("TESTSETKEY", "TESTSETVAL").subscribe(consumer)
+ }
then:
conds.await()
assertTraces(1) {
- trace(0, 1) {
+ trace(0, 3) {
span(0) {
+ name "parent"
+ kind INTERNAL
+ hasNoParent()
+ }
+ span(1) {
name "SET"
kind CLIENT
+ childOf(span(0))
attributes {
"$SemanticAttributes.DB_SYSTEM.key" "redis"
"$SemanticAttributes.DB_STATEMENT.key" "SET TESTSETKEY ?"
"$SemanticAttributes.DB_OPERATION.key" "SET"
}
}
+ span(2) {
+ name "callback"
+ kind INTERNAL
+ childOf(span(0))
+ }
}
}
}
@@ -131,26 +147,41 @@ class LettuceReactiveClientTest extends AgentInstrumentationSpecification {
final defaultVal = "NOT THIS VALUE"
when:
+ runWithSpan("parent") {
reactiveCommands.get("NON_EXISTENT_KEY").defaultIfEmpty(defaultVal).subscribe {
res ->
+ runWithSpan("callback") {
conds.evaluate {
assert res == defaultVal
}
}
+ }
+ }
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_STATEMENT.key" "GET NON_EXISTENT_KEY"
"$SemanticAttributes.DB_OPERATION.key" "GET"
}
}
+ span(2) {
+ name "callback"
+ kind INTERNAL
+ childOf(span(0))
+ }
}
}

ADD
instrumentation/lettuce/lettuce-5.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_1/LettuceAsyncCommandInstrumentation.java
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.lettuce.v5_1;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import io.lettuce.core.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("io.lettuce.core.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();
+ 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();
+ }
+ }
+}

MODIFY
instrumentation/lettuce/lettuce-5.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_1/LettuceInstrumentationModule.java
instrumentation/lettuce/lettuce-5.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_1/LettuceInstrumentationModule.java
package io.opentelemetry.javaagent.instrumentation.lettuce.v5_1;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
import com.google.auto.service.AutoService;
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -28,6 +28,7 @@ public class LettuceInstrumentationModule extends InstrumentationModule {
@Override
public List<TypeInstrumentation> typeInstrumentations() {
- return singletonList(new DefaultClientResourcesInstrumentation());
+ return asList(
+ new DefaultClientResourcesInstrumentation(), new LettuceAsyncCommandInstrumentation());
}
}

RENAME
instrumentation/lettuce/lettuce-5.1/library/src/test/groovy/io/opentelemetry/instrumentation/lettuce/v5_1/LettuceAsyncSyncClientTest.groovy
instrumentation/lettuce/lettuce-5.1/library/src/test/groovy/io/opentelemetry/instrumentation/lettuce/v5_1/LettuceAsyncClientTest.groovy
@@ -9,7 +9,7 @@ import io.lettuce.core.RedisClient
import io.lettuce.core.resource.ClientResources
import io.opentelemetry.instrumentation.test.LibraryTestTrait
-class LettuceAsyncSyncClientTest extends AbstractLettuceAsyncClientTest implements LibraryTestTrait {
+class LettuceAsyncClientTest extends AbstractLettuceAsyncClientTest implements LibraryTestTrait {
@Override
RedisClient createClient(String uri) {
return RedisClient.create(
@@ -18,4 +18,10 @@ class LettuceAsyncSyncClientTest extends AbstractLettuceAsyncClientTest implemen
.build(),
uri)
}
+
+ @Override
+ boolean testCallback() {
+ // context is not propagated into callbacks
+ return false
+ }
}

MODIFY
instrumentation/lettuce/lettuce-5.1/testing/src/main/groovy/io/opentelemetry/instrumentation/lettuce/v5_1/AbstractLettuceAsyncClientTest.groovy
instrumentation/lettuce/lettuce-5.1/testing/src/main/groovy/io/opentelemetry/instrumentation/lettuce/v5_1/AbstractLettuceAsyncClientTest.groovy
package io.opentelemetry.instrumentation.lettuce.v5_1
import static io.opentelemetry.api.trace.SpanKind.CLIENT
+import static io.opentelemetry.api.trace.SpanKind.INTERNAL
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.NetTransportValues.IP_TCP
import io.lettuce.core.ConnectionFuture
@@ -94,6 +95,17 @@ abstract class AbstractLettuceAsyncClientTest extends InstrumentationSpecificati
redisServer.stop()
}
+ boolean testCallback() {
+ return true
+ }
+
+ def <T> T runWithCallbackSpan(String spanName, Closure callback) {
+ if (testCallback()) {
+ return runWithSpan(spanName, callback)
+ }
+ return callback.call()
+ }
+
def "connect using get on ConnectionFuture"() {
setup:
RedisClient testConnectionClient = RedisClient.create(embeddedDbUri)
@@ -166,21 +178,30 @@ abstract class AbstractLettuceAsyncClientTest extends InstrumentationSpecificati
Consumer<String> consumer = new Consumer<String>() {
@Override
void accept(String res) {
+ runWithCallbackSpan("callback") {
conds.evaluate {
assert res == "TESTVAL"
}
}
}
+ }
when:
+ runWithSpan("parent") {
RedisFuture<String> redisFuture = asyncCommands.get("TESTKEY")
redisFuture.thenAccept(consumer)
+ }
then:
conds.await()
assertTraces(1) {
- trace(0, 1) {
+ trace(0, 2 + (testCallback() ? 1 : 0)) {
span(0) {
+ name "parent"
+ kind INTERNAL
+ hasNoParent()
+ }
+ span(1) {
name "GET"
kind CLIENT
attributes {
@@ -197,6 +218,13 @@ abstract class AbstractLettuceAsyncClientTest extends InstrumentationSpecificati
eventName "redis.encode.end"
}
}
+ if (testCallback()) {
+ span(2) {
+ name "callback"
+ kind INTERNAL
+ childOf(span(0))
+ }
+ }
}
}
}
@@ -210,34 +238,46 @@ abstract class AbstractLettuceAsyncClientTest extends InstrumentationSpecificati
BiFunction<String, Throwable, String> firstStage = new BiFunction<String, Throwable, String>() {
@Override
String apply(String res, Throwable throwable) {
+ runWithCallbackSpan("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) {
+ runWithCallbackSpan("callback2") {
conds.evaluate {
assert input == successStr
}
+ }
return null
}
}
when:
+ runWithSpan("parent") {
RedisFuture<String> redisFuture = asyncCommands.get("NON_EXISTENT_KEY")
redisFuture.handleAsync(firstStage).thenApply(secondStage)
+ }
then:
conds.await()
assertTraces(1) {
- trace(0, 1) {
+ trace(0, 2 + (testCallback() ? 2 : 0)) {
span(0) {
+ name "parent"
+ kind INTERNAL
+ hasNoParent()
+ }
+ span(1) {
name "GET"
kind CLIENT
+ childOf(span(0))
attributes {
"${SemanticAttributes.NET_TRANSPORT.key}" IP_TCP
"${SemanticAttributes.NET_PEER_IP.key}" "127.0.0.1"
@@ -252,6 +292,18 @@ abstract class AbstractLettuceAsyncClientTest extends InstrumentationSpecificati
eventName "redis.encode.end"
}
}
+ if (testCallback()) {
+ span(2) {
+ name "callback1"
+ kind INTERNAL
+ childOf(span(0))
+ }
+ span(3) {
+ name "callback2"
+ kind INTERNAL
+ childOf(span(0))
+ }
+ }
}
}
}
@@ -262,23 +314,33 @@ abstract class AbstractLettuceAsyncClientTest extends InstrumentationSpecificati
BiConsumer<String, Throwable> biConsumer = new BiConsumer<String, Throwable>() {
@Override
void accept(String keyRetrieved, Throwable throwable) {
+ runWithCallbackSpan("callback") {
conds.evaluate {
assert keyRetrieved != null
}
}
}
+ }
when:
+ runWithSpan("parent") {
RedisFuture<String> redisFuture = asyncCommands.randomkey()
redisFuture.whenCompleteAsync(biConsumer)
+ }
then:
conds.await()
assertTraces(1) {
- trace(0, 1) {
+ trace(0, 2 + (testCallback() ? 1 : 0)) {
span(0) {
+ name "parent"
+ kind INTERNAL
+ hasNoParent()
+ }
+ span(1) {
name "RANDOMKEY"
kind CLIENT
+ childOf(span(0))
attributes {
"${SemanticAttributes.NET_TRANSPORT.key}" IP_TCP
"${SemanticAttributes.NET_PEER_IP.key}" "127.0.0.1"
@@ -293,6 +355,13 @@ abstract class AbstractLettuceAsyncClientTest extends InstrumentationSpecificati
eventName "redis.encode.end"
}
}
+ if (testCallback()) {
+ span(2) {
+ name "callback"
+ kind INTERNAL
+ childOf(span(0))
+ }
+ }
}
}
}

MODIFY
instrumentation/lettuce/lettuce-5.1/testing/src/main/groovy/io/opentelemetry/instrumentation/lettuce/v5_1/AbstractLettuceReactiveClientTest.groovy
instrumentation/lettuce/lettuce-5.1/testing/src/main/groovy/io/opentelemetry/instrumentation/lettuce/v5_1/AbstractLettuceReactiveClientTest.groovy
package io.opentelemetry.instrumentation.lettuce.v5_1
import static io.opentelemetry.api.trace.SpanKind.CLIENT
+import static io.opentelemetry.api.trace.SpanKind.INTERNAL
import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.NetTransportValues.IP_TCP
import io.lettuce.core.RedisClient
@@ -74,22 +75,32 @@ abstract class AbstractLettuceReactiveClientTest extends InstrumentationSpecific
Consumer<String> consumer = new Consumer<String>() {
@Override
void accept(String res) {
+ runWithSpan("callback") {
conds.evaluate {
assert res == "OK"
}
}
}
+ }
when:
+ runWithSpan("parent") {
reactiveCommands.set("TESTSETKEY", "TESTSETVAL").subscribe(consumer)
+ }
then:
conds.await()
assertTraces(1) {
- trace(0, 1) {
+ trace(0, 3) {
span(0) {
+ name "parent"
+ kind INTERNAL
+ hasNoParent()
+ }
+ span(1) {
name "SET"
kind CLIENT
+ childOf(span(0))
attributes {
"${SemanticAttributes.NET_TRANSPORT.key}" IP_TCP
"${SemanticAttributes.NET_PEER_IP.key}" "127.0.0.1"
@@ -104,6 +115,11 @@ abstract class AbstractLettuceReactiveClientTest extends InstrumentationSpecific
eventName "redis.encode.end"
}
}
+ span(2) {
+ name "callback"
+ kind INTERNAL
+ childOf(span(0))
+ }
}
}
}
@@ -148,20 +164,30 @@ abstract class AbstractLettuceReactiveClientTest extends InstrumentationSpecific
final defaultVal = "NOT THIS VALUE"
when:
+ runWithSpan("parent") {
reactiveCommands.get("NON_EXISTENT_KEY").defaultIfEmpty(defaultVal).subscribe {
res ->
+ runWithSpan("callback") {
conds.evaluate {
assert res == defaultVal
}
}
+ }
+ }
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.NET_TRANSPORT.key}" IP_TCP
"${SemanticAttributes.NET_PEER_IP.key}" "127.0.0.1"
@@ -176,6 +202,11 @@ abstract class AbstractLettuceReactiveClientTest extends InstrumentationSpecific
eventName "redis.encode.end"
}
}
+ span(2) {
+ name "callback"
+ kind INTERNAL
+ childOf(span(0))
+ }
}
}

Propagate context to lettuce callbacks (#3839)