Ensure tracer resources cleaned up in exit advice (#11905)

Ensure tracer resources cleaned up in exit advice are still cleaned up even if the patched method throws an exception

Review feedback

Co-authored-by: Andrea Marziali <marziali.andrea@gmail.com>

Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>

Author: mcculls

dd-java-agent/instrumentation/elasticsearch/elasticsearch-transport/elasticsearch-transport-common/src/main/java/datadog/trace/instrumentation/elasticsearch/ShadowExistingScopeAdvice.java

@@ -17,7 +17,7 @@ public static AgentScope enter() {
     return activateSpan(noopSpan());
   }
 
-  @Advice.OnMethodExit(suppress = Throwable.class)
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void exit(@Advice.Enter final AgentScope scope) {
     scope.close();
   }

dd-java-agent/instrumentation/gax-1.4/src/main/java/datadog/trace/instrumentation/gax/CallbackChainRetryingFutureInstrumentation.java

@@ -65,7 +65,7 @@ public static Runnable capturePrevious(
       return previousListener;
     }
 
-    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void cancelSuperseded(
         @Advice.Enter final Runnable previousListener,
         @Advice.FieldValue("attemptFutureCompletionListener") final Runnable newListener) {

dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyByteBodyInstrumentation.java

@@ -207,7 +207,7 @@ static void after(
   }
 
   static class NIOInputStreamRecycleAdvice {
-    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     static void after(@Advice.This final NIOInputStream thiz) {
       InstrumentationContext.get(NIOInputStream.class, StoredByteBody.class).put(thiz, null);
     }

dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyCharBodyInstrumentation.java

@@ -221,7 +221,7 @@ static void after(
   }
 
   static class NIOReaderRecycleAdvice {
-    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     static void after(@Advice.This final NIOReader thiz) {
       InstrumentationContext.get(NIOReader.class, StoredCharBody.class).put(thiz, null);
     }

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/completablefuture/AsyncTaskInstrumentation.java

@@ -52,14 +52,14 @@ public static AgentScope before(@Advice.This ForkJoinTask<?> zis) {
       return startTaskScope(InstrumentationContext.get(ForkJoinTask.class, State.class), zis);
     }
 
-    @Advice.OnMethodExit
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
     public static void after(@Advice.Enter AgentScope scope) {
       endTaskScope(scope);
     }
   }
 
   public static class Cancel {
-    @Advice.OnMethodExit
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
     public static <T> void cancel(@Advice.This ForkJoinTask<T> task) {
       State state = InstrumentationContext.get(ForkJoinTask.class, State.class).get(task);
       if (null != state) {

dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfoInstrumentation.java

@@ -37,6 +37,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.consumer.internals.ConsumerCoordinator;
+import org.apache.kafka.common.errors.WakeupException;
 
 /**
  * This instrumentation saves additional information from the KafkaConsumer, such as consumer group
@@ -258,11 +259,12 @@ public static AgentScope onEnter(@Advice.This KafkaConsumer consumer) {
       return null;
     }
 
-    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void captureGroup(
         @Advice.Enter final AgentScope scope,
         @Advice.This KafkaConsumer consumer,
-        @Advice.Return ConsumerRecords records) {
+        @Advice.Return ConsumerRecords records,
+        @Advice.Thrown Throwable throwable) {
       int recordsCount = 0;
       if (records != null) {
         KafkaConsumerInfo kafkaConsumerInfo =
@@ -281,6 +283,9 @@ public static void captureGroup(
       }
       AgentSpan span = scope.span();
       span.setTag(KAFKA_RECORDS_COUNT, recordsCount);
+      if (!(throwable instanceof WakeupException)) {
+        span.addThrowable(throwable);
+      }
       span.finish();
       scope.close();
     }

dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/RecordsAdvice.java

@@ -17,6 +17,7 @@
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.internals.ConsumerDelegate;
+import org.apache.kafka.common.errors.WakeupException;
 
 /**
  * this method transfers the consumer group from the KafkaConsumer class key to the ConsumerRecords
@@ -45,11 +46,12 @@ public static AgentScope onEnter(@Advice.This ConsumerDelegate consumer) {
     return null;
   }
 
-  @Advice.OnMethodExit(suppress = Throwable.class)
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void captureGroup(
       @Advice.Enter final AgentScope scope,
       @Advice.This ConsumerDelegate consumer,
-      @Advice.Return ConsumerRecords records) {
+      @Advice.Return ConsumerRecords records,
+      @Advice.Thrown Throwable throwable) {
     int recordsCount = 0;
     if (records != null) {
       // new - we are getting the KafkaConsumerInfo from the ConsumerDelegate instead of
@@ -70,6 +72,9 @@ public static void captureGroup(
     }
     AgentSpan span = scope.span();
     span.setTag(KAFKA_RECORDS_COUNT, recordsCount);
+    if (!(throwable instanceof WakeupException)) {
+      span.addThrowable(throwable);
+    }
     span.finish();
     scope.close();
   }

dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/CommandHandlerInstrumentation.java

@@ -63,7 +63,7 @@ public static AgentScope before(@Advice.Argument(2) RedisCommand command) {
       return null;
     }
 
-    @Advice.OnMethodExit
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
     public static void after(@Advice.Enter AgentScope scope) {
       endTaskScope(scope);
     }

dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/rx/RedisSubscriptionCommandCompleteAdvice.java

@@ -12,7 +12,7 @@
 
 public class RedisSubscriptionCommandCompleteAdvice {
 
-  @Advice.OnMethodExit(suppress = Throwable.class)
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void afterComplete(
       @Advice.Origin("#m") String method,
       @Advice.This RedisCommand command,

dd-java-agent/instrumentation/lettuce/lettuce-5.0/src/main/java/datadog/trace/instrumentation/lettuce5/rx/RedisSubscriptionCommandErrorAdvice.java

@@ -9,7 +9,7 @@
 import net.bytebuddy.asm.Advice;
 
 public class RedisSubscriptionCommandErrorAdvice {
-  @Advice.OnMethodExit(suppress = Throwable.class)
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void afterError(
       @Advice.This RedisCommand command, @Advice.Argument(value = 0) Throwable throwable) {