From ffada12fbdfee40bac8646ac1a39b8f1b9043eda Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 9 Jun 2026 17:16:55 +0200 Subject: [PATCH 1/3] Fine tune async propagation for lettuce 5 --- .../lettuce5/AsyncCommandInstrumentation.java | 24 ++++++- .../lettuce5/ConnectionFutureAdvice.java | 16 +++-- .../lettuce5/LettuceAsyncCommandsAdvice.java | 62 +++++++++++++------ .../lettuce5/LettuceInstrumentationUtil.java | 17 +++++ .../test/groovy/Lettuce5ClientTestBase.groovy | 7 --- 5 files changed, 90 insertions(+), 36 deletions(-) diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/AsyncCommandInstrumentation.java b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/AsyncCommandInstrumentation.java index 3e69c09ddcf..414408e3867 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/AsyncCommandInstrumentation.java +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/AsyncCommandInstrumentation.java @@ -2,6 +2,8 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled; import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture; import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.endTaskScope; import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.startTaskScope; @@ -53,8 +55,10 @@ public void methodAdvice(MethodTransformer transformer) { .and(takesArgument(0, named("io.lettuce.core.protocol.RedisCommand")))), getClass().getName() + "$Capture"); transformer.applyAdvice( - isMethod().and(namedOneOf("complete", "completeExceptionally", "onComplete", "encode")), + isMethod().and(namedOneOf("complete", "completeExceptionally", "encode")), getClass().getName() + "$Activate"); + transformer.applyAdvice( + isMethod().and(named("onComplete")), getClass().getName() + "$SuppressAsyncPropagation"); transformer.applyAdvice( isMethod().and(named("cancel")).and(takesArguments(boolean.class)), getClass().getName() + "$Cancel"); @@ -91,4 +95,22 @@ public static void before(@Advice.This AsyncCommand asyncCommand) { InstrumentationContext.get(AsyncCommand.class, State.class), asyncCommand); } } + + public static final class SuppressAsyncPropagation { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static boolean before() { + if (isAsyncPropagationEnabled()) { + setAsyncPropagationEnabled(false); + return true; + } + return false; + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void after(@Advice.Enter final boolean restore) { + if (restore) { + setAsyncPropagationEnabled(true); + } + } + } } diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java index fd9c335af54..c0481b09a89 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java @@ -1,11 +1,9 @@ package datadog.trace.instrumentation.lettuce5; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.lettuce5.LettuceClientDecorator.DECORATE; import datadog.trace.bootstrap.InstrumentationContext; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.lettuce.core.ConnectionFuture; import io.lettuce.core.RedisURI; @@ -14,28 +12,29 @@ public class ConnectionFutureAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope onEnter(@Advice.Argument(1) final RedisURI redisUri) { - final AgentSpan span = + public static void onEnter( + @Advice.Argument(1) final RedisURI redisUri, @Advice.Local("ddSpan") AgentSpan span) { + span = startSpan( LettuceClientDecorator.REDIS_CLIENT.toString(), LettuceClientDecorator.OPERATION_NAME); DECORATE.afterStart(span); span.setResourceName(DECORATE.resourceNameForConnection(redisUri)); DECORATE.onConnection(span, redisUri); - return activateSpan(span); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Enter final AgentScope scope, + @Advice.Local("ddSpan") final AgentSpan span, @Advice.Thrown final Throwable throwable, @Advice.Argument(1) final RedisURI redisUri, @Advice.Return(readOnly = false) ConnectionFuture connectionFuture) { - final AgentSpan span = scope.span(); + if (span == null) { + return; + } if (throwable != null) { DECORATE.onError(span, throwable); DECORATE.beforeFinish(span); - scope.close(); span.finish(); return; } @@ -44,7 +43,6 @@ public static void stopSpan( new ConnectionContextBiConsumer( redisUri, InstrumentationContext.get(StatefulConnection.class, RedisURI.class)) .andThen(new LettuceAsyncBiConsumer<>(span))); - scope.close(); // span finished by LettuceAsyncBiConsumer } } diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceAsyncCommandsAdvice.java b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceAsyncCommandsAdvice.java index 8c201abd048..d22f66f8706 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceAsyncCommandsAdvice.java +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceAsyncCommandsAdvice.java @@ -3,7 +3,9 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.lettuce5.LettuceClientDecorator.DECORATE; +import static datadog.trace.instrumentation.lettuce5.LettuceInstrumentationUtil.disableAsyncPropagation; import static datadog.trace.instrumentation.lettuce5.LettuceInstrumentationUtil.expectsResponse; +import static datadog.trace.instrumentation.lettuce5.LettuceInstrumentationUtil.restoreAsyncPropagation; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -20,7 +22,9 @@ public class LettuceAsyncCommandsAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope onEnter( @Advice.Argument(0) final RedisCommand command, - @Advice.This final AbstractRedisAsyncCommands thiz) { + @Advice.This final AbstractRedisAsyncCommands thiz, + @Advice.Local("commandExpectsResponse") boolean commandExpectsResponse, + @Advice.Local("restoreAsyncPropagation") boolean restoreAsyncPropagation) { final AgentSpan span = startSpan( @@ -32,33 +36,53 @@ public static AgentScope onEnter( .get(thiz.getConnection())); DECORATE.onCommand(span, command); - return activateSpan(span); + final AgentScope scope = activateSpan(span); + commandExpectsResponse = expectsResponse(command); + if (!commandExpectsResponse) { + restoreAsyncPropagation = disableAsyncPropagation(); + } + + return scope; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Argument(0) final RedisCommand command, @Advice.Enter final AgentScope scope, + @Advice.Local("commandExpectsResponse") final boolean commandExpectsResponse, + @Advice.Local("restoreAsyncPropagation") final boolean restoreAsyncPropagation, @Advice.Thrown final Throwable throwable, @Advice.Return AsyncCommand asyncCommand) { - final AgentSpan span = scope.span(); - if (throwable != null) { - DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); - scope.close(); - span.finish(); - return; - } + try { + if (scope == null) { + return; + } + final AgentSpan span = scope.span(); + if (throwable != null) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + span.finish(); + return; + } - // close spans on error or normal completion - if (expectsResponse(command)) { - asyncCommand.whenComplete(new LettuceAsyncBiConsumer<>(span)); - } else { - DECORATE.beforeFinish(span); - span.finish(); + // close spans on error or normal completion + if (commandExpectsResponse) { + final boolean restoreCompletionCallbackPropagation = disableAsyncPropagation(); + try { + asyncCommand.whenComplete(new LettuceAsyncBiConsumer<>(span)); + } finally { + restoreAsyncPropagation(restoreCompletionCallbackPropagation); + } + } else { + DECORATE.beforeFinish(span); + span.finish(); + } + // span may be finished by LettuceAsyncBiConsumer + } finally { + restoreAsyncPropagation(restoreAsyncPropagation); + if (scope != null) { + scope.close(); + } } - scope.close(); - // span may be finished by LettuceAsyncBiConsumer } } diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceInstrumentationUtil.java b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceInstrumentationUtil.java index ec4cc28fa78..92aeabe6fd3 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceInstrumentationUtil.java +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceInstrumentationUtil.java @@ -1,5 +1,8 @@ package datadog.trace.instrumentation.lettuce5; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled; + import io.lettuce.core.protocol.RedisCommand; import java.util.Arrays; import java.util.HashSet; @@ -75,4 +78,18 @@ public static String getCommandName(final RedisCommand command) { } return commandName; } + + public static boolean disableAsyncPropagation() { + if (isAsyncPropagationEnabled()) { + setAsyncPropagationEnabled(false); + return true; + } + return false; + } + + public static void restoreAsyncPropagation(final boolean restore) { + if (restore) { + setAsyncPropagationEnabled(true); + } + } } diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/test/groovy/Lettuce5ClientTestBase.groovy b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/test/groovy/Lettuce5ClientTestBase.groovy index b38c142cede..9749dbdfb0f 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/test/groovy/Lettuce5ClientTestBase.groovy +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/test/groovy/Lettuce5ClientTestBase.groovy @@ -42,13 +42,6 @@ abstract class Lettuce5ClientTestBase extends VersionedNamingTestBase { RedisAsyncCommands asyncCommands RedisCommands syncCommands - @Override - boolean useStrictTraceWrites() { - // latest seems leaking continuations that terminates later hence the strict trace will discard our spans. - !isLatestDepTest - } - - def setup() { redisServer.start() println "Using redis: $redisServer.redisURI" From bfcbd9b9491c169632681886a7d652215c9b2b8d Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 10 Jun 2026 00:24:27 +0200 Subject: [PATCH 2/3] suggestions --- .../lettuce5/ConnectionFutureAdvice.java | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java index c0481b09a89..77c2851275d 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java @@ -2,6 +2,8 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.lettuce5.LettuceClientDecorator.DECORATE; +import static datadog.trace.instrumentation.lettuce5.LettuceInstrumentationUtil.disableAsyncPropagation; +import static datadog.trace.instrumentation.lettuce5.LettuceInstrumentationUtil.restoreAsyncPropagation; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -38,11 +40,17 @@ public static void stopSpan( span.finish(); return; } - connectionFuture = - connectionFuture.whenComplete( - new ConnectionContextBiConsumer( - redisUri, InstrumentationContext.get(StatefulConnection.class, RedisURI.class)) - .andThen(new LettuceAsyncBiConsumer<>(span))); + final boolean restoreCompletionCallbackPropagation = disableAsyncPropagation(); + try { + connectionFuture = + connectionFuture.whenComplete( + new ConnectionContextBiConsumer( + redisUri, + InstrumentationContext.get(StatefulConnection.class, RedisURI.class)) + .andThen(new LettuceAsyncBiConsumer<>(span))); + } finally { + restoreAsyncPropagation(restoreCompletionCallbackPropagation); + } // span finished by LettuceAsyncBiConsumer } } From 09b2c415adeedf16c50f09fb1bdf0549a75ecb16 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 10 Jun 2026 14:05:14 +0200 Subject: [PATCH 3/3] simplify the first version --- .../lettuce/lettuce-5.0/build.gradle | 1 + .../lettuce5/AsyncCommandInstrumentation.java | 24 +------ .../lettuce5/ConnectionFutureAdvice.java | 27 +++----- .../lettuce5/LettuceAsyncCommandsAdvice.java | 62 ++++++------------- .../lettuce5/LettuceInstrumentationUtil.java | 17 ----- 5 files changed, 31 insertions(+), 100 deletions(-) diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/build.gradle b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/build.gradle index becb0a29ad7..b306aa29d18 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/build.gradle +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/build.gradle @@ -21,6 +21,7 @@ dependencies { testImplementation group: 'io.lettuce', name: 'lettuce-core', version: '5.0.0.RELEASE' testImplementation project(':dd-java-agent:instrumentation:reactor-core-3.1') testImplementation project(':dd-java-agent:instrumentation:reactive-streams-1.0') + testImplementation project(':dd-java-agent:instrumentation:netty:netty-promise-4.0') latestDepTestImplementation group: 'io.lettuce', name: 'lettuce-core', version: '+' diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/AsyncCommandInstrumentation.java b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/AsyncCommandInstrumentation.java index 414408e3867..3e69c09ddcf 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/AsyncCommandInstrumentation.java +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/AsyncCommandInstrumentation.java @@ -2,8 +2,6 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled; import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture; import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.endTaskScope; import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.startTaskScope; @@ -55,10 +53,8 @@ public void methodAdvice(MethodTransformer transformer) { .and(takesArgument(0, named("io.lettuce.core.protocol.RedisCommand")))), getClass().getName() + "$Capture"); transformer.applyAdvice( - isMethod().and(namedOneOf("complete", "completeExceptionally", "encode")), + isMethod().and(namedOneOf("complete", "completeExceptionally", "onComplete", "encode")), getClass().getName() + "$Activate"); - transformer.applyAdvice( - isMethod().and(named("onComplete")), getClass().getName() + "$SuppressAsyncPropagation"); transformer.applyAdvice( isMethod().and(named("cancel")).and(takesArguments(boolean.class)), getClass().getName() + "$Cancel"); @@ -95,22 +91,4 @@ public static void before(@Advice.This AsyncCommand asyncCommand) { InstrumentationContext.get(AsyncCommand.class, State.class), asyncCommand); } } - - public static final class SuppressAsyncPropagation { - @Advice.OnMethodEnter(suppress = Throwable.class) - public static boolean before() { - if (isAsyncPropagationEnabled()) { - setAsyncPropagationEnabled(false); - return true; - } - return false; - } - - @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void after(@Advice.Enter final boolean restore) { - if (restore) { - setAsyncPropagationEnabled(true); - } - } - } } diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java index 77c2851275d..d14d49f943b 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/ConnectionFutureAdvice.java @@ -2,8 +2,6 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.lettuce5.LettuceClientDecorator.DECORATE; -import static datadog.trace.instrumentation.lettuce5.LettuceInstrumentationUtil.disableAsyncPropagation; -import static datadog.trace.instrumentation.lettuce5.LettuceInstrumentationUtil.restoreAsyncPropagation; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -14,19 +12,19 @@ public class ConnectionFutureAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static void onEnter( - @Advice.Argument(1) final RedisURI redisUri, @Advice.Local("ddSpan") AgentSpan span) { - span = + public static AgentSpan onEnter(@Advice.Argument(1) final RedisURI redisUri) { + final AgentSpan span = startSpan( LettuceClientDecorator.REDIS_CLIENT.toString(), LettuceClientDecorator.OPERATION_NAME); DECORATE.afterStart(span); span.setResourceName(DECORATE.resourceNameForConnection(redisUri)); DECORATE.onConnection(span, redisUri); + return span; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Local("ddSpan") final AgentSpan span, + @Advice.Enter final AgentSpan span, @Advice.Thrown final Throwable throwable, @Advice.Argument(1) final RedisURI redisUri, @Advice.Return(readOnly = false) @@ -40,17 +38,12 @@ public static void stopSpan( span.finish(); return; } - final boolean restoreCompletionCallbackPropagation = disableAsyncPropagation(); - try { - connectionFuture = - connectionFuture.whenComplete( - new ConnectionContextBiConsumer( - redisUri, - InstrumentationContext.get(StatefulConnection.class, RedisURI.class)) - .andThen(new LettuceAsyncBiConsumer<>(span))); - } finally { - restoreAsyncPropagation(restoreCompletionCallbackPropagation); - } + connectionFuture = + connectionFuture.whenComplete( + new ConnectionContextBiConsumer( + redisUri, InstrumentationContext.get(StatefulConnection.class, RedisURI.class)) + .andThen(new LettuceAsyncBiConsumer<>(span))); + // span finished by LettuceAsyncBiConsumer } } diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceAsyncCommandsAdvice.java b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceAsyncCommandsAdvice.java index d22f66f8706..8c201abd048 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceAsyncCommandsAdvice.java +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceAsyncCommandsAdvice.java @@ -3,9 +3,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.lettuce5.LettuceClientDecorator.DECORATE; -import static datadog.trace.instrumentation.lettuce5.LettuceInstrumentationUtil.disableAsyncPropagation; import static datadog.trace.instrumentation.lettuce5.LettuceInstrumentationUtil.expectsResponse; -import static datadog.trace.instrumentation.lettuce5.LettuceInstrumentationUtil.restoreAsyncPropagation; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -22,9 +20,7 @@ public class LettuceAsyncCommandsAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope onEnter( @Advice.Argument(0) final RedisCommand command, - @Advice.This final AbstractRedisAsyncCommands thiz, - @Advice.Local("commandExpectsResponse") boolean commandExpectsResponse, - @Advice.Local("restoreAsyncPropagation") boolean restoreAsyncPropagation) { + @Advice.This final AbstractRedisAsyncCommands thiz) { final AgentSpan span = startSpan( @@ -36,53 +32,33 @@ public static AgentScope onEnter( .get(thiz.getConnection())); DECORATE.onCommand(span, command); - final AgentScope scope = activateSpan(span); - commandExpectsResponse = expectsResponse(command); - if (!commandExpectsResponse) { - restoreAsyncPropagation = disableAsyncPropagation(); - } - - return scope; + return activateSpan(span); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( + @Advice.Argument(0) final RedisCommand command, @Advice.Enter final AgentScope scope, - @Advice.Local("commandExpectsResponse") final boolean commandExpectsResponse, - @Advice.Local("restoreAsyncPropagation") final boolean restoreAsyncPropagation, @Advice.Thrown final Throwable throwable, @Advice.Return AsyncCommand asyncCommand) { - try { - if (scope == null) { - return; - } - final AgentSpan span = scope.span(); - if (throwable != null) { - DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); - span.finish(); - return; - } + final AgentSpan span = scope.span(); + if (throwable != null) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + scope.close(); + span.finish(); + return; + } - // close spans on error or normal completion - if (commandExpectsResponse) { - final boolean restoreCompletionCallbackPropagation = disableAsyncPropagation(); - try { - asyncCommand.whenComplete(new LettuceAsyncBiConsumer<>(span)); - } finally { - restoreAsyncPropagation(restoreCompletionCallbackPropagation); - } - } else { - DECORATE.beforeFinish(span); - span.finish(); - } - // span may be finished by LettuceAsyncBiConsumer - } finally { - restoreAsyncPropagation(restoreAsyncPropagation); - if (scope != null) { - scope.close(); - } + // close spans on error or normal completion + if (expectsResponse(command)) { + asyncCommand.whenComplete(new LettuceAsyncBiConsumer<>(span)); + } else { + DECORATE.beforeFinish(span); + span.finish(); } + scope.close(); + // span may be finished by LettuceAsyncBiConsumer } } diff --git a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceInstrumentationUtil.java b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceInstrumentationUtil.java index 92aeabe6fd3..ec4cc28fa78 100644 --- a/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceInstrumentationUtil.java +++ b/dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/LettuceInstrumentationUtil.java @@ -1,8 +1,5 @@ package datadog.trace.instrumentation.lettuce5; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled; - import io.lettuce.core.protocol.RedisCommand; import java.util.Arrays; import java.util.HashSet; @@ -78,18 +75,4 @@ public static String getCommandName(final RedisCommand command) { } return commandName; } - - public static boolean disableAsyncPropagation() { - if (isAsyncPropagationEnabled()) { - setAsyncPropagationEnabled(false); - return true; - } - return false; - } - - public static void restoreAsyncPropagation(final boolean restore) { - if (restore) { - setAsyncPropagationEnabled(true); - } - } }