From 8949d6c05f9de4c1287fe98cec0a6d6d564a8c11 Mon Sep 17 00:00:00 2001 From: Trask Stalnaker Date: Thu, 13 Feb 2020 14:17:47 -0800 Subject: [PATCH] Make sure CallDepthThreadLocalMap is always reset --- .../CouchbaseBucketInstrumentation.java | 2 +- .../CouchbaseClusterInstrumentation.java | 2 +- .../PreparedStatementInstrumentation.java | 8 +++--- .../jdbc/StatementInstrumentation.java | 8 +++--- .../NettyChannelPipelineInstrumentation.java | 2 +- .../NettyChannelPipelineInstrumentation.java | 2 +- .../MemcachedClientInstrumentation.java | 27 ++++++++++++------- 7 files changed, 30 insertions(+), 21 deletions(-) diff --git a/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java b/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java index 4a805b6aef..89e10472e3 100644 --- a/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java +++ b/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java @@ -64,7 +64,7 @@ public class CouchbaseBucketInstrumentation extends Instrumenter.Default { return CallDepthThreadLocalMap.incrementCallDepth(CouchbaseCluster.class); } - @Advice.OnMethodExit + @Advice.OnMethodExit(onThrowable = Throwable.class) public static void subscribeResult( @Advice.Enter final int callDepth, @Advice.Origin final Method method, diff --git a/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java b/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java index 0af4109ec6..b9f58044d6 100644 --- a/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java +++ b/dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java @@ -64,7 +64,7 @@ public class CouchbaseClusterInstrumentation extends Instrumenter.Default { return CallDepthThreadLocalMap.incrementCallDepth(CouchbaseCluster.class); } - @Advice.OnMethodExit + @Advice.OnMethodExit(onThrowable = Throwable.class) public static void subscribeResult( @Advice.Enter final int callDepth, @Advice.Origin final Method method, diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/PreparedStatementInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/PreparedStatementInstrumentation.java index 8d2c5aa675..0e4f413b33 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/PreparedStatementInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/PreparedStatementInstrumentation.java @@ -61,13 +61,13 @@ public final class PreparedStatementInstrumentation extends Instrumenter.Default @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope onEnter(@Advice.This final PreparedStatement statement) { - final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(PreparedStatement.class); - if (callDepth > 0) { + final Connection connection = connectionFromStatement(statement); + if (connection == null) { return null; } - final Connection connection = connectionFromStatement(statement); - if (connection == null) { + final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(PreparedStatement.class); + if (callDepth > 0) { return null; } diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java index 345529f58c..25e4782bd0 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java @@ -62,13 +62,13 @@ public final class StatementInstrumentation extends Instrumenter.Default { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope onEnter( @Advice.Argument(0) final String sql, @Advice.This final Statement statement) { - final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Statement.class); - if (callDepth > 0) { + final Connection connection = connectionFromStatement(statement); + if (connection == null) { return null; } - final Connection connection = connectionFromStatement(statement); - if (connection == null) { + final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Statement.class); + if (callDepth > 0) { return null; } diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelPipelineInstrumentation.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelPipelineInstrumentation.java index 7beb747880..22dcc87d78 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelPipelineInstrumentation.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelPipelineInstrumentation.java @@ -101,7 +101,7 @@ public class NettyChannelPipelineInstrumentation extends Instrumenter.Default { return CallDepthThreadLocalMap.incrementCallDepth(ChannelPipeline.class); } - @Advice.OnMethodExit(suppress = Throwable.class) + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void addHandler( @Advice.Enter final int depth, @Advice.This final ChannelPipeline pipeline, diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelPipelineInstrumentation.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelPipelineInstrumentation.java index f96732cadd..b8903651cf 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelPipelineInstrumentation.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelPipelineInstrumentation.java @@ -110,7 +110,7 @@ public class NettyChannelPipelineInstrumentation extends Instrumenter.Default { return CallDepthThreadLocalMap.incrementCallDepth(handler.getClass()); } - @Advice.OnMethodExit(suppress = Throwable.class) + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void addHandler( @Advice.Enter final int depth, @Advice.This final ChannelPipeline pipeline, diff --git a/dd-java-agent/instrumentation/spymemcached-2.12/src/main/java/datadog/trace/instrumentation/spymemcached/MemcachedClientInstrumentation.java b/dd-java-agent/instrumentation/spymemcached-2.12/src/main/java/datadog/trace/instrumentation/spymemcached/MemcachedClientInstrumentation.java index fbb571eb48..087b70e16a 100644 --- a/dd-java-agent/instrumentation/spymemcached-2.12/src/main/java/datadog/trace/instrumentation/spymemcached/MemcachedClientInstrumentation.java +++ b/dd-java-agent/instrumentation/spymemcached-2.12/src/main/java/datadog/trace/instrumentation/spymemcached/MemcachedClientInstrumentation.java @@ -81,14 +81,17 @@ public final class MemcachedClientInstrumentation extends Instrumenter.Default { return CallDepthThreadLocalMap.incrementCallDepth(MemcachedClient.class) <= 0; } - @Advice.OnMethodExit(suppress = Throwable.class) + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( @Advice.Enter final boolean shouldInjectListener, @Advice.This final MemcachedClient client, @Advice.Origin("#m") final String methodName, @Advice.Return final OperationFuture future) { - if (shouldInjectListener && future != null) { - CallDepthThreadLocalMap.reset(MemcachedClient.class); + if (!shouldInjectListener) { + return; + } + CallDepthThreadLocalMap.reset(MemcachedClient.class); + if (future != null) { final OperationCompletionListener listener = new OperationCompletionListener(client.getConnection(), methodName); future.addListener(listener); @@ -103,14 +106,17 @@ public final class MemcachedClientInstrumentation extends Instrumenter.Default { return CallDepthThreadLocalMap.incrementCallDepth(MemcachedClient.class) <= 0; } - @Advice.OnMethodExit(suppress = Throwable.class) + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( @Advice.Enter final boolean shouldInjectListener, @Advice.This final MemcachedClient client, @Advice.Origin("#m") final String methodName, @Advice.Return final GetFuture future) { - if (shouldInjectListener && future != null) { - CallDepthThreadLocalMap.reset(MemcachedClient.class); + if (!shouldInjectListener) { + return; + } + CallDepthThreadLocalMap.reset(MemcachedClient.class); + if (future != null) { final GetCompletionListener listener = new GetCompletionListener(client.getConnection(), methodName); future.addListener(listener); @@ -125,14 +131,17 @@ public final class MemcachedClientInstrumentation extends Instrumenter.Default { return CallDepthThreadLocalMap.incrementCallDepth(MemcachedClient.class) <= 0; } - @Advice.OnMethodExit(suppress = Throwable.class) + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( @Advice.Enter final boolean shouldInjectListener, @Advice.This final MemcachedClient client, @Advice.Origin("#m") final String methodName, @Advice.Return final BulkFuture future) { - if (shouldInjectListener && future != null) { - CallDepthThreadLocalMap.reset(MemcachedClient.class); + if (!shouldInjectListener) { + return; + } + CallDepthThreadLocalMap.reset(MemcachedClient.class); + if (future != null) { final BulkGetCompletionListener listener = new BulkGetCompletionListener(client.getConnection(), methodName); future.addListener(listener);