Merge branch 'master' into jb/repo_cleanup

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelper.java

@@ -72,9 +72,11 @@ public static boolean shouldPropagate(
     return Boolean.TRUE.equals(contextStore.get(executor));
   }
 
-  public static void capture(ContextStore<Runnable, State> contextStore, Runnable task) {
+  public static void capture(
+      ContextStore<Runnable, State> contextStore, ThreadPoolExecutor executor, Runnable task) {
     if (task != null && !exclude(RUNNABLE, task)) {
       AdviceUtils.capture(contextStore, task, true);
+      QueueTimerHelper.startQueuingTimer(contextStore, executor.getClass(), task);
     }
   }
 

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java

@@ -63,7 +63,6 @@ public class CapturedContextInstrumentor extends Instrumentor {
   private int exitContextVar = -1;
   private int timestampStartVar = -1;
   private int throwableListVar = -1;
-  private final List<FinallyBlock> finallyBlocks = new ArrayList<>();
 
   public CapturedContextInstrumentor(
       ProbeDefinition definition,
@@ -97,14 +96,6 @@ public InstrumentationResult.Status instrument() {
     return InstrumentationResult.Status.INSTALLED;
   }
 
-  private void installFinallyBlocks() {
-    for (FinallyBlock finallyBlock : finallyBlocks) {
-      methodNode.tryCatchBlocks.add(
-          new TryCatchBlockNode(
-              finallyBlock.startLabel, finallyBlock.endLabel, finallyBlock.handlerLabel, null));
-    }
-  }
-
   private boolean addLineCaptures(LineMap lineMap) {
     Where.SourceLine[] targetLines = definition.getWhere().getSourceLines();
     if (targetLines == null) {
@@ -1045,16 +1036,4 @@ private void addCapturedValueOf(InsnList insnList, Limits limits) {
         INT_TYPE);
     // stack: [captured_value]
   }
-
-  private static class FinallyBlock {
-    final LabelNode startLabel;
-    final LabelNode endLabel;
-    final LabelNode handlerLabel;
-
-    public FinallyBlock(LabelNode startLabel, LabelNode endLabel, LabelNode handlerLabel) {
-      this.startLabel = startLabel;
-      this.endLabel = endLabel;
-      this.handlerLabel = handlerLabel;
-    }
-  }
 }

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/Instrumentor.java

@@ -10,8 +10,11 @@
 import com.datadog.debugger.probe.ProbeDefinition;
 import com.datadog.debugger.probe.Where;
 import datadog.trace.bootstrap.debugger.ProbeId;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.AbstractInsnNode;
@@ -24,6 +27,7 @@
 import org.objectweb.asm.tree.LocalVariableNode;
 import org.objectweb.asm.tree.MethodInsnNode;
 import org.objectweb.asm.tree.MethodNode;
+import org.objectweb.asm.tree.TryCatchBlockNode;
 import org.objectweb.asm.tree.TypeInsnNode;
 
 /** Common class for generating instrumentation */
@@ -45,6 +49,7 @@ public abstract class Instrumentor {
   protected int argOffset;
   protected final LocalVariableNode[] localVarsBySlot;
   protected LabelNode returnHandlerLabel;
+  protected final List<CapturedContextInstrumentor.FinallyBlock> finallyBlocks = new ArrayList<>();
 
   public Instrumentor(
       ProbeDefinition definition,
@@ -270,4 +275,38 @@ protected ProbeDefinition.Tag[] addProbeIdWithTags(String probeId, ProbeDefiniti
     newTags[newTags.length - 1] = new ProbeDefinition.Tag(PROBEID_TAG_NAME, probeId);
     return newTags;
   }
+
+  protected InsnList clone(InsnList insnList) {
+    InsnList result = new InsnList();
+    Map<LabelNode, LabelNode> labels = new HashMap<>();
+    for (AbstractInsnNode node : insnList) {
+      if (node instanceof LabelNode) {
+        labels.put((LabelNode) node, new LabelNode());
+      }
+    }
+    for (AbstractInsnNode node : insnList) {
+      result.add(node.clone(labels));
+    }
+    return result;
+  }
+
+  protected void installFinallyBlocks() {
+    for (FinallyBlock finallyBlock : finallyBlocks) {
+      methodNode.tryCatchBlocks.add(
+          new TryCatchBlockNode(
+              finallyBlock.startLabel, finallyBlock.endLabel, finallyBlock.handlerLabel, null));
+    }
+  }
+
+  protected static class FinallyBlock {
+    final LabelNode startLabel;
+    final LabelNode endLabel;
+    final LabelNode handlerLabel;
+
+    public FinallyBlock(LabelNode startLabel, LabelNode endLabel, LabelNode handlerLabel) {
+      this.startLabel = startLabel;
+      this.endLabel = endLabel;
+      this.handlerLabel = handlerLabel;
+    }
+  }
 }

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/MetricInstrumentor.java

@@ -80,6 +80,8 @@ public class MetricInstrumentor extends Instrumentor {
   private static final InsnList EMPTY_INSN_LIST = new InsnList();
 
   private final MetricProbe metricProbe;
+  private int durationStartVar = -1;
+  private LabelNode durationStartLabel;
 
   public MetricInstrumentor(
       MetricProbe metricProbe,
@@ -109,6 +111,8 @@ public InstrumentationResult.Status instrument() {
       case EXIT:
         {
           processInstructions();
+          addFinallyHandler(returnHandlerLabel);
+          installFinallyBlocks();
           break;
         }
       default:
@@ -118,6 +122,15 @@ public InstrumentationResult.Status instrument() {
     return InstrumentationResult.Status.INSTALLED;
   }
 
+  private void addFinallyHandler(LabelNode endLabel) {
+    // only for @duration case
+    if (durationStartVar == -1) {
+      return;
+    }
+    InsnList insnList = callMetric(metricProbe, null);
+    createFinallyHandler(durationStartLabel, endLabel, wrapTryCatch(insnList));
+  }
+
   private InsnList wrapTryCatch(InsnList insnList) {
     if (insnList == null || insnList == EMPTY_INSN_LIST) {
       return EMPTY_INSN_LIST;
@@ -190,6 +203,25 @@ protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
     return insnList;
   }
 
+  private void createFinallyHandler(LabelNode startLabel, LabelNode endLabel, InsnList insnList) {
+    LabelNode handlerLabel = new LabelNode();
+    InsnList handler = new InsnList();
+    handler.add(handlerLabel);
+    // declare a local var to store the current Throwable of the 'finally' block
+    int throwableTmpVar = newVar(Type.getType(Throwable.class));
+    // stack [exception]
+    handler.add(new VarInsnNode(Opcodes.ASTORE, throwableTmpVar));
+    // stack []
+    handler.add(insnList);
+    // stack []
+    // restore the current Throwable to the stack
+    handler.add(new VarInsnNode(Opcodes.ALOAD, throwableTmpVar));
+    // stack [exception]
+    handler.add(new InsnNode(Opcodes.ATHROW));
+    methodNode.instructions.add(handler);
+    finallyBlocks.add(new FinallyBlock(startLabel, endLabel, handlerLabel));
+  }
+
   private InsnList callCount(MetricProbe metricProbe, ReturnContext returnContext) {
     if (metricProbe.getValue() == null) {
       InsnList insnList = new InsnList();
@@ -354,7 +386,7 @@ public VisitorResult(ASMHelper.Type type, InsnList insnList) {
     }
   }
 
-  private static class MetricValueVisitor implements Visitor<VisitorResult> {
+  private class MetricValueVisitor implements Visitor<VisitorResult> {
     private final MetricInstrumentor instrumentor;
     private final InsnList nullBranch;
     private final ReturnContext returnContext;
@@ -780,21 +812,26 @@ private ASMHelper.Type tryRetrieveSynthetic(String name, InsnList insnList) {
           return null;
         }
         // call System.nanoTime at the beginning of the method
-        int var = instrumentor.newVar(LONG_TYPE);
-        InsnList nanoTimeList = new InsnList();
-        invokeStatic(nanoTimeList, Type.getType(System.class), "nanoTime", LONG_TYPE);
-        nanoTimeList.add(new VarInsnNode(Opcodes.LSTORE, var));
-        instrumentor.methodNode.instructions.insert(instrumentor.methodEnterLabel, nanoTimeList);
+        if (durationStartVar == -1) {
+          durationStartVar = instrumentor.newVar(LONG_TYPE);
+          InsnList nanoTimeList = new InsnList();
+          invokeStatic(nanoTimeList, Type.getType(System.class), "nanoTime", LONG_TYPE);
+          nanoTimeList.add(new VarInsnNode(Opcodes.LSTORE, durationStartVar));
+          durationStartLabel = new LabelNode();
+          nanoTimeList.add(durationStartLabel);
+          instrumentor.methodNode.instructions.insert(instrumentor.methodEnterLabel, nanoTimeList);
+        }
         // diff nanoTime before calling metric
         invokeStatic(insnList, Type.getType(System.class), "nanoTime", LONG_TYPE);
         // stack [long]
-        insnList.add(new VarInsnNode(Opcodes.LLOAD, var));
+        insnList.add(new VarInsnNode(Opcodes.LLOAD, durationStartVar));
         // stack [long, long]
         insnList.add(new InsnNode(Opcodes.LSUB));
         // stack [long]
         insnList.add(new InsnNode(Opcodes.L2D));
+        // stack [double]
         insnList.add(new LdcInsnNode(1_000_000D));
-        // stack [long, double]
+        // stack [double, double]
         insnList.add(new InsnNode(Opcodes.DDIV));
         // stack [double]
         return new ASMHelper.Type(DOUBLE_TYPE);

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/MetricProbesInstrumentationTest.java

@@ -8,6 +8,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
@@ -425,6 +426,29 @@ public void methodSyntheticDurationGaugeMetric() throws IOException, URISyntaxEx
     assertArrayEquals(new String[] {METRIC_PROBEID_TAG}, listener.lastTags);
   }
 
+  @Test
+  public void methodSyntheticDurationExceptionGaugeMetric() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "CapturedSnapshot05";
+    String METRIC_NAME = "syn_gauge";
+    MetricProbe metricProbe =
+        createMetricBuilder(METRIC_ID, METRIC_NAME, GAUGE)
+            .where(CLASS_NAME, "main", "(String)")
+            .valueScript(new ValueScript(DSL.ref("@duration"), "@duration"))
+            .evaluateAt(MethodLocation.EXIT)
+            .build();
+    MetricForwarderListener listener = installMetricProbes(metricProbe);
+    Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    try {
+      Reflect.on(testClass).call("main", "triggerUncaughtException").get();
+      fail("should not reach this code");
+    } catch (Exception ex) {
+      assertEquals("oops", ex.getCause().getCause().getMessage());
+    }
+    assertTrue(listener.doubleGauges.containsKey(METRIC_NAME));
+    assertTrue(listener.doubleGauges.get(METRIC_NAME).doubleValue() > 0);
+    assertArrayEquals(new String[] {METRIC_PROBEID_TAG}, listener.lastTags);
+  }
+
   @Test
   public void lineArgumentRefValueCountMetric() throws IOException, URISyntaxException {
     final String CLASS_NAME = "CapturedSnapshot03";

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/JavaForkJoinPoolInstrumentation.java

@@ -60,7 +60,7 @@ public static <T> void externalPush(
     @Advice.OnMethodExit(onThrowable = Throwable.class)
     public static <T> void cleanup(
         @Advice.Argument(0) ForkJoinTask<T> task, @Advice.Thrown Throwable thrown) {
-      if (null != thrown) {
+      if (null != thrown && !exclude(FORK_JOIN_TASK, task)) {
         cancelTask(InstrumentationContext.get(ForkJoinTask.class, State.class), task);
       }
     }

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/JavaTimerInstrumentation.java

@@ -3,6 +3,8 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.cancelTask;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.QueueTimerHelper.startQueuingTimer;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -58,15 +60,17 @@ public static void before(@Advice.Argument(0) TimerTask task, @Advice.Argument(2
       if (period != 0) {
         return;
       }
-      ContextStore<Runnable, State> contextStore =
-          InstrumentationContext.get(Runnable.class, State.class);
-      capture(contextStore, task, true);
-      startQueuingTimer(contextStore, Timer.class, task);
+      if (!exclude(RUNNABLE, task)) {
+        ContextStore<Runnable, State> contextStore =
+            InstrumentationContext.get(Runnable.class, State.class);
+        capture(contextStore, task, true);
+        startQueuingTimer(contextStore, Timer.class, task);
+      }
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class)
     public static void after(@Advice.Argument(0) TimerTask task, @Advice.Thrown Throwable thrown) {
-      if (null != thrown) {
+      if (null != thrown && !exclude(RUNNABLE, task)) {
         cancelTask(InstrumentationContext.get(Runnable.class, State.class), task);
       }
     }

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ThreadPoolExecutorInstrumentation.java

@@ -16,11 +16,9 @@
 import datadog.trace.agent.tooling.ExcludeFilterProvider;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.api.Platform;
-import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
-import datadog.trace.bootstrap.instrumentation.java.concurrent.QueueTimerHelper;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.TPEHelper;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.Wrapper;
@@ -150,10 +148,7 @@ public static void capture(
         if (TPEHelper.useWrapping(task)) {
           task = Wrapper.wrap(task);
         } else {
-          ContextStore<Runnable, State> contextStore =
-              InstrumentationContext.get(Runnable.class, State.class);
-          TPEHelper.capture(contextStore, task);
-          QueueTimerHelper.startQueuingTimer(contextStore, tpe.getClass(), task);
+          TPEHelper.capture(InstrumentationContext.get(Runnable.class, State.class), tpe, task);
         }
       }
     }