Propagate context to lettuce callbacks (#3839)

This commit is contained in:
Lauri Tulmin 2021-08-17 10:02:23 +03:00 committed by GitHub
parent c54d192a1a
commit 8c175d4fce
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
16 changed files with 580 additions and 92 deletions

View File

@ -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();
}
}
}

View File

@ -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();
}

View File

@ -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());
}
}

View File

@ -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();

View File

@ -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,9 +232,11 @@ 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)
}
@ -226,30 +244,50 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
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,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))
}
}
}
}

View File

@ -0,0 +1,66 @@
/*
* 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();
}
}
}

View File

@ -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();
}

View File

@ -30,6 +30,7 @@ public class LettuceInstrumentationModule extends InstrumentationModule {
@Override
public List<TypeInstrumentation> typeInstrumentations() {
return asList(
new LettuceAsyncCommandInstrumentation(),
new LettuceAsyncCommandsInstrumentation(),
new LettuceClientInstrumentation(),
new LettuceReactiveCommandsInstrumentation());

View File

@ -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();

View File

@ -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 io.lettuce.core.ClientOptions
@ -185,29 +186,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_STATEMENT.key" "GET TESTKEY"
"$SemanticAttributes.DB_OPERATION.key" "GET"
}
}
span(2) {
name "callback"
kind INTERNAL
childOf(span(0))
}
}
}
}
@ -221,9 +237,11 @@ 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)
}
@ -231,30 +249,50 @@ class LettuceAsyncClientTest extends AgentInstrumentationSpecification {
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.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) {
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_STATEMENT.key" "RANDOMKEY"
"$SemanticAttributes.DB_OPERATION.key" "RANDOMKEY"
}
}
span(2) {
name "callback"
kind INTERNAL
childOf(span(0))
}
}
}
}
@ -401,12 +454,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
}
}
})
@ -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))
}
}
}
}

View File

@ -4,6 +4,7 @@
*/
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) {
conds.evaluate {
assert res == "OK"
runWithSpan("callback") {
conds.evaluate {
assert res == "OK"
}
}
}
}
when:
reactiveCommands.set("TESTSETKEY", "TESTSETVAL").subscribe(consumer)
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:
reactiveCommands.get("NON_EXISTENT_KEY").defaultIfEmpty(defaultVal).subscribe {
res ->
conds.evaluate {
assert res == defaultVal
}
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))
}
}
}

View File

@ -0,0 +1,64 @@
/*
* 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();
}
}
}

View File

@ -6,7 +6,7 @@
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());
}
}

View File

@ -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
}
}

View File

@ -6,6 +6,7 @@
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) {
conds.evaluate {
assert res == "TESTVAL"
runWithCallbackSpan("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, 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,9 +238,11 @@ abstract class AbstractLettuceAsyncClientTest extends InstrumentationSpecificati
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
runWithCallbackSpan("callback1") {
conds.evaluate {
assert res == null
assert throwable == null
}
}
return (res == null ? successStr : res)
}
@ -220,24 +250,34 @@ abstract class AbstractLettuceAsyncClientTest extends InstrumentationSpecificati
Function<String, Object> secondStage = new Function<String, Object>() {
@Override
Object apply(String input) {
conds.evaluate {
assert input == successStr
runWithCallbackSpan("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.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) {
conds.evaluate {
assert keyRetrieved != null
runWithCallbackSpan("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, 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))
}
}
}
}
}

View File

@ -6,6 +6,7 @@
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) {
conds.evaluate {
assert res == "OK"
runWithSpan("callback") {
conds.evaluate {
assert res == "OK"
}
}
}
}
when:
reactiveCommands.set("TESTSETKEY", "TESTSETVAL").subscribe(consumer)
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:
reactiveCommands.get("NON_EXISTENT_KEY").defaultIfEmpty(defaultVal).subscribe {
res ->
conds.evaluate {
assert res == defaultVal
}
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))
}
}
}