Add evalution timeout for expressions

dd-java-agent/agent-debugger/debugger-bootstrap/build.gradle

@@ -14,4 +14,5 @@ dependencies {
   implementation libs.slf4j
   implementation libs.instrument.java
   implementation project(':internal-api')
+  testImplementation libs.bundles.mockito
 }

dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/DebuggerContext.java

@@ -6,7 +6,6 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.lang.reflect.Method;
 import java.time.Duration;
-import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -324,8 +323,8 @@ public static void evalContext(
       // snapshot
       if (needFreeze) {
         Duration timeout =
-            Duration.of(Config.get().getDynamicInstrumentationCaptureTimeout(), ChronoUnit.MILLIS);
-        context.freeze(new TimeoutChecker(timeout));
+            Duration.ofMillis(Config.get().getDynamicInstrumentationCaptureTimeout());
+        context.freeze(TimeoutChecker.create(Config.get(), timeout));
       }
     } catch (Exception ex) {
       LOGGER.debug("Error in evalContext: ", ex);
@@ -359,8 +358,8 @@ public static void evalContext(
       // snapshot
       if (needFreeze) {
         Duration timeout =
-            Duration.of(Config.get().getDynamicInstrumentationCaptureTimeout(), ChronoUnit.MILLIS);
-        context.freeze(new TimeoutChecker(timeout));
+            Duration.ofMillis(Config.get().getDynamicInstrumentationCaptureTimeout());
+        context.freeze(TimeoutChecker.create(Config.get(), timeout));
       }
     } catch (Exception ex) {
       LOGGER.debug("Error in evalContext: ", ex);

dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/el/DebuggerScript.java

@@ -1,10 +1,12 @@
 package datadog.trace.bootstrap.debugger.el;
 
+import datadog.trace.bootstrap.debugger.util.TimeoutChecker;
+
 /**
  * A debugger EL script interface used for communication between the instrumented code and the
  * debugger EL.<br>
  * Because it must be reachable from the instrumented code it must be placed in bootstrap.
  */
 public interface DebuggerScript<R> {
-  R execute(ValueReferenceResolver valueRefResolver);
+  R execute(ValueReferenceResolver valueRefResolver, TimeoutChecker timeoutChecker);
 }

dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/util/CpuTimeoutChecker.java

@@ -0,0 +1,26 @@
+package datadog.trace.bootstrap.debugger.util;
+
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadMXBean;
+import java.time.Duration;
+
+public class CpuTimeoutChecker implements TimeoutChecker {
+  private final Duration timeOut;
+  private final ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean();
+  private final long startCpuTime;
+
+  public CpuTimeoutChecker(Duration timeout) {
+    this.timeOut = timeout;
+    startCpuTime = threadMXBean.getCurrentThreadCpuTime();
+  }
+
+  @Override
+  public boolean isTimedOut() {
+    return threadMXBean.getCurrentThreadCpuTime() - startCpuTime >= timeOut.toNanos();
+  }
+
+  @Override
+  public Duration getTimeOut() {
+    return timeOut;
+  }
+}

dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/util/TimeoutChecker.java

@@ -1,33 +1,21 @@
 package datadog.trace.bootstrap.debugger.util;
 
+import datadog.trace.api.Config;
 import java.time.Duration;
-import java.time.temporal.ChronoUnit;
 
-public class TimeoutChecker {
-  public static final Duration DEFAULT_TIME_OUT = Duration.of(100, ChronoUnit.MILLIS);
+public interface TimeoutChecker {
 
-  private final long start;
-  private final Duration timeOut;
+  String CPU = "CPU";
+  String WALL = "WALL";
 
-  public TimeoutChecker(Duration timeOut) {
-    this.start = System.currentTimeMillis();
-    this.timeOut = timeOut;
-  }
-
-  public TimeoutChecker(long start, Duration timeOut) {
-    this.start = start;
-    this.timeOut = timeOut;
-  }
+  boolean isTimedOut();
 
-  public boolean isTimedOut(long currentTimeMillis) {
-    return (currentTimeMillis - start) > timeOut.toMillis();
-  }
-
-  public long getStart() {
-    return start;
-  }
+  Duration getTimeOut();
 
-  public Duration getTimeOut() {
-    return timeOut;
+  static TimeoutChecker create(Config config, Duration timeout) {
+    if (config.getDynamicInstrumentationTimeoutCheckerMode().equals(CPU)) {
+      return new CpuTimeoutChecker(timeout);
+    }
+    return new WallTimeoutChecker(timeout);
   }
 }

dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/util/WallTimeoutChecker.java

@@ -0,0 +1,37 @@
+package datadog.trace.bootstrap.debugger.util;
+
+import java.time.Duration;
+
+public class WallTimeoutChecker implements TimeoutChecker {
+
+  private final long start;
+  private final Duration timeOut;
+
+  public WallTimeoutChecker(Duration timeOut) {
+    this.start = System.currentTimeMillis();
+    this.timeOut = timeOut;
+  }
+
+  public WallTimeoutChecker(long start, Duration timeOut) {
+    this.start = start;
+    this.timeOut = timeOut;
+  }
+
+  public boolean isTimedOut(long currentTimeMillis) {
+    return (currentTimeMillis - start) > timeOut.toMillis();
+  }
+
+  @Override
+  public boolean isTimedOut() {
+    return isTimedOut(System.currentTimeMillis());
+  }
+
+  public long getStart() {
+    return start;
+  }
+
+  @Override
+  public Duration getTimeOut() {
+    return timeOut;
+  }
+}

dd-java-agent/agent-debugger/debugger-bootstrap/src/test/java/datadog/trace/bootstrap/debugger/util/TimeoutCheckerTest.java

@@ -1,20 +1,55 @@
 package datadog.trace.bootstrap.debugger.util;
 
-import static datadog.trace.bootstrap.debugger.util.TimeoutChecker.DEFAULT_TIME_OUT;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
-import org.junit.jupiter.api.Assertions;
+import datadog.trace.api.Config;
+import java.time.Duration;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.LockSupport;
 import org.junit.jupiter.api.Test;
 
 public class TimeoutCheckerTest {
   @Test
-  public void timedOut() {
+  public void wallTimedOut() {
     long start = System.currentTimeMillis();
+    Duration timeout = Duration.ofMillis(100);
     // assume that start & timestamp captured in instance below are very close
-    TimeoutChecker timeoutChecker = new TimeoutChecker(DEFAULT_TIME_OUT);
-    Assertions.assertFalse(timeoutChecker.isTimedOut(start + 1));
-    Assertions.assertTrue(timeoutChecker.isTimedOut(start + DEFAULT_TIME_OUT.toMillis() * 2));
-    timeoutChecker = new TimeoutChecker(start, DEFAULT_TIME_OUT);
-    Assertions.assertFalse(timeoutChecker.isTimedOut(start + 1));
-    Assertions.assertTrue(timeoutChecker.isTimedOut(start + DEFAULT_TIME_OUT.toMillis() * 2));
+    WallTimeoutChecker timeoutChecker = new WallTimeoutChecker(timeout);
+    assertFalse(timeoutChecker.isTimedOut(start + 1));
+    assertTrue(timeoutChecker.isTimedOut(start + timeout.toMillis() * 2));
+    timeoutChecker = new WallTimeoutChecker(start, timeout);
+    assertFalse(timeoutChecker.isTimedOut(start + 1));
+    assertTrue(timeoutChecker.isTimedOut(start + timeout.toMillis() * 2));
+  }
+
+  @Test
+  public void cpuTimedOut() {
+    CpuTimeoutChecker cpuTimeoutChecker = new CpuTimeoutChecker(Duration.ofMillis(1));
+    LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(10)); // not consume cpu
+    assertFalse(cpuTimeoutChecker.isTimedOut());
+    burnCpu(5_000_000);
+    assertTrue(cpuTimeoutChecker.isTimedOut());
+  }
+
+  static volatile int counter;
+
+  private static void burnCpu(int iterations) {
+    for (int i = 0; i < iterations; i++) {
+      counter++;
+    }
+  }
+
+  @Test
+  public void configTimeout() {
+    Config config = mock(Config.class);
+    when(config.getDynamicInstrumentationTimeoutCheckerMode()).thenReturn(TimeoutChecker.CPU);
+    assertInstanceOf(CpuTimeoutChecker.class, TimeoutChecker.create(config, Duration.ofMillis(50)));
+    when(config.getDynamicInstrumentationTimeoutCheckerMode()).thenReturn(TimeoutChecker.WALL);
+    assertInstanceOf(
+        WallTimeoutChecker.class, TimeoutChecker.create(config, Duration.ofMillis(50)));
   }
 }

dd-java-agent/agent-debugger/debugger-el/src/main/java/com/datadog/debugger/el/BooleanValueExpressionAdapter.java

@@ -3,7 +3,6 @@
 import com.datadog.debugger.el.expressions.BooleanExpression;
 import com.datadog.debugger.el.expressions.ValueExpression;
 import com.datadog.debugger.el.values.BooleanValue;
-import datadog.trace.bootstrap.debugger.el.ValueReferenceResolver;
 
 public class BooleanValueExpressionAdapter implements ValueExpression<BooleanValue> {
 
@@ -14,8 +13,8 @@ public BooleanValueExpressionAdapter(BooleanExpression booleanExpression) {
   }
 
   @Override
-  public BooleanValue evaluate(ValueReferenceResolver valueRefResolver) {
-    Boolean result = booleanExpression.evaluate(valueRefResolver);
+  public BooleanValue evaluate(EvalContext evalContext) {
+    Boolean result = booleanExpression.evaluate(evalContext);
     if (result == null) {
       throw new EvaluationException(
           "Boolean expression returning null", PrettyPrintVisitor.print(this));

dd-java-agent/agent-debugger/debugger-el/src/main/java/com/datadog/debugger/el/EvalContext.java

@@ -0,0 +1,23 @@
+package com.datadog.debugger.el;
+
+import datadog.trace.bootstrap.debugger.el.ValueReferenceResolver;
+import datadog.trace.bootstrap.debugger.util.TimeoutChecker;
+
+public class EvalContext {
+  private final ValueReferenceResolver valueRefResolver;
+  private final TimeoutChecker timeoutChecker;
+
+  public EvalContext(
+      final ValueReferenceResolver valueRefResolver, final TimeoutChecker timeoutChecker) {
+    this.valueRefResolver = valueRefResolver;
+    this.timeoutChecker = timeoutChecker;
+  }
+
+  public ValueReferenceResolver getValueRefResolver() {
+    return valueRefResolver;
+  }
+
+  public TimeoutChecker getTimeoutChecker() {
+    return timeoutChecker;
+  }
+}

dd-java-agent/agent-debugger/debugger-el/src/main/java/com/datadog/debugger/el/Expression.java

@@ -1,11 +1,9 @@
 package com.datadog.debugger.el;
 
-import datadog.trace.bootstrap.debugger.el.ValueReferenceResolver;
-
 /** Represents any evaluable expression */
 @FunctionalInterface
 public interface Expression<ReturnType> {
-  ReturnType evaluate(ValueReferenceResolver valueRefResolver);
+  ReturnType evaluate(EvalContext evalContext);
 
   default <R> R accept(Visitor<R> visitor) {
     return null;

dd-java-agent/agent-debugger/debugger-el/src/main/java/com/datadog/debugger/el/Literal.java

@@ -1,7 +1,6 @@
 package com.datadog.debugger.el;
 
 import com.datadog.debugger.el.expressions.ValueExpression;
-import datadog.trace.bootstrap.debugger.el.ValueReferenceResolver;
 import datadog.trace.bootstrap.debugger.el.Values;
 import java.util.Objects;
 
@@ -27,7 +26,7 @@ public boolean isUndefined() {
   }
 
   @Override
-  public Value<ConstantType> evaluate(ValueReferenceResolver valueRefResolver) {
+  public Value<ConstantType> evaluate(EvalContext evalContext) {
     return this;
   }
 

dd-java-agent/agent-debugger/debugger-el/src/main/java/com/datadog/debugger/el/ProbeCondition.java

@@ -9,6 +9,7 @@
 import com.squareup.moshi.JsonWriter;
 import datadog.trace.bootstrap.debugger.el.DebuggerScript;
 import datadog.trace.bootstrap.debugger.el.ValueReferenceResolver;
+import datadog.trace.bootstrap.debugger.util.TimeoutChecker;
 import java.io.IOException;
 import javax.annotation.Nonnull;
 
@@ -99,12 +100,12 @@ public static ProbeCondition load(JsonReader reader) throws IOException {
   }
 
   @Override
-  public Boolean execute(ValueReferenceResolver valueRefResolver) {
+  public Boolean execute(ValueReferenceResolver valueRefResolver, TimeoutChecker timeoutChecker) {
     if (when == null) {
       return true;
     }
-    if (when.evaluate(valueRefResolver)) {
-      then.evaluate(valueRefResolver);
+    if (when.evaluate(new EvalContext(valueRefResolver, timeoutChecker))) {
+      then.evaluate(new EvalContext(valueRefResolver, timeoutChecker));
       return true;
     }
     return false;