Fix Timeout for tests

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/JsonSnapshotSerializer.java

@@ -9,11 +9,16 @@
 import datadog.trace.api.ProcessTags;
 import datadog.trace.bootstrap.debugger.CapturedContext;
 import datadog.trace.bootstrap.debugger.DebuggerContext;
+import java.time.Duration;
+import java.time.temporal.ChronoUnit;
 
 /** Serializes snapshots in Json using Moshi */
 public class JsonSnapshotSerializer implements DebuggerContext.ValueSerializer {
   private static final JsonAdapter<IntakeRequest> ADAPTER =
-      MoshiHelper.createMoshiSnapshot().adapter(IntakeRequest.class);
+      MoshiHelper.createMoshiSnapshot(
+              Duration.of(
+                  Config.get().getDynamicInstrumentationCaptureTimeout(), ChronoUnit.MILLIS))
+          .adapter(IntakeRequest.class);
   private static final JsonAdapter<CapturedContext.CapturedValue> VALUE_ADAPTER =
       new MoshiSnapshotHelper.CapturedValueAdapter();
 

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/util/MoshiHelper.java

@@ -11,6 +11,7 @@
 import com.squareup.moshi.Types;
 import datadog.trace.bootstrap.debugger.el.DebuggerScript;
 import java.lang.reflect.ParameterizedType;
+import java.time.Duration;
 import java.util.Map;
 
 /** Helper for creating Moshi instances with the right adapters depending on the context */
@@ -32,9 +33,9 @@ public static Moshi.Builder createMoshiConfigBuilder() {
         .add(ProbeDefinition.Tag[].class, new ProbeDefinition.TagAdapter());
   }
 
-  public static Moshi createMoshiSnapshot() {
+  public static Moshi createMoshiSnapshot(Duration captureTimeOut) {
     return new Moshi.Builder()
-        .add(new MoshiSnapshotHelper.SnapshotJsonFactory())
+        .add(new MoshiSnapshotHelper.SnapshotJsonFactory(captureTimeOut))
         .add(
             DebuggerScript.class,
             new ProbeCondition.ProbeConditionJsonAdapter()) // ProbeDetails in Snapshot

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/util/MoshiSnapshotHelper.java

@@ -6,6 +6,7 @@
 import com.squareup.moshi.JsonWriter;
 import com.squareup.moshi.Moshi;
 import com.squareup.moshi.Types;
+import datadog.trace.api.Config;
 import datadog.trace.bootstrap.debugger.CapturedContext;
 import datadog.trace.bootstrap.debugger.CapturedStackFrame;
 import datadog.trace.bootstrap.debugger.Limits;
@@ -62,14 +63,20 @@ public class MoshiSnapshotHelper {
   public static final String LOCATION = "location";
   public static final String MESSAGE = "message";
   public static final String STACKTRACE = "stacktrace";
-  private static final Duration TIME_OUT = Duration.ofMillis(200);
 
   public static class SnapshotJsonFactory implements JsonAdapter.Factory {
+    private final Duration captureTimeOut;
+
+    public SnapshotJsonFactory(Duration captureTimeOut) {
+      this.captureTimeOut = captureTimeOut;
+    }
+
     @Override
     public JsonAdapter<?> create(Type type, Set<? extends Annotation> set, Moshi moshi) {
       if (Types.equals(type, Snapshot.Captures.class)) {
         return new CapturesAdapter(
             moshi,
+            captureTimeOut,
             new CapturedContextAdapter(
                 moshi, new CapturedValueAdapter(), new CapturedThrowableAdapter(moshi)));
       }
@@ -88,12 +95,15 @@ public JsonAdapter<?> create(Type type, Set<? extends Annotation> set, Moshi mos
   }
 
   public static class CapturesAdapter extends JsonAdapter<Snapshot.Captures> {
+    protected final Duration captureTimeOut;
     protected final JsonAdapter<CapturedContext> capturedContextAdapter;
     protected final JsonAdapter<Map<Integer, CapturedContext>> linesAdapter;
     protected final JsonAdapter<List<CapturedContext.CapturedThrowable>> caughtExceptionsAdapter;
 
-    public CapturesAdapter(Moshi moshi, JsonAdapter<CapturedContext> capturedContextAdapter) {
+    public CapturesAdapter(
+        Moshi moshi, Duration captureTimeout, JsonAdapter<CapturedContext> capturedContextAdapter) {
       this.capturedContextAdapter = capturedContextAdapter;
+      this.captureTimeOut = captureTimeout;
       linesAdapter =
           moshi.adapter(
               Types.newParameterizedType(Map.class, Integer.class, CapturedContext.class));
@@ -109,7 +119,7 @@ public void toJson(JsonWriter jsonWriter, Snapshot.Captures captures) throws IOE
         return;
       }
       jsonWriter.setTag(
-          TimeoutChecker.class, new TimeoutChecker(System.currentTimeMillis(), TIME_OUT));
+          TimeoutChecker.class, new TimeoutChecker(System.currentTimeMillis(), captureTimeOut));
       jsonWriter.beginObject();
       jsonWriter.name(ENTRY);
       capturedContextAdapter.toJson(jsonWriter, captures.getEntry());
@@ -149,7 +159,9 @@ public void toJson(JsonWriter jsonWriter, CapturedContext capturedContext) throw
       }
       TimeoutChecker timeoutChecker = jsonWriter.tag(TimeoutChecker.class);
       if (timeoutChecker == null) {
-        timeoutChecker = new TimeoutChecker(TIME_OUT);
+        Duration timeout =
+            Duration.of(Config.get().getDynamicInstrumentationCaptureTimeout(), ChronoUnit.MILLIS);
+        timeoutChecker = new TimeoutChecker(timeout);
       }
       // need to 'freeze' the context before serializing it
       capturedContext.freeze(timeoutChecker);

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

@@ -789,7 +789,10 @@ public void length255() throws IOException {
   @Test
   public void capturesAdapterNull() {
     MoshiSnapshotHelper.CapturesAdapter capturesAdapter =
-        new MoshiSnapshotHelper.CapturesAdapter(MoshiHelper.createMoshiSnapshot(), null);
+        new MoshiSnapshotHelper.CapturesAdapter(
+            MoshiSnapshotTestHelper.createMoshiSnapshot(),
+            Duration.of(5, ChronoUnit.SECONDS),
+            null);
     Assertions.assertEquals("null", capturesAdapter.toJson(null));
   }
 
@@ -921,7 +924,9 @@ public void fieldCount20() throws IOException {
   public void timeOut() throws IOException {
     DebuggerContext.initValueSerializer(
         new TimeoutSnapshotSerializer(Duration.of(150, ChronoUnit.MILLIS)));
-    JsonAdapter<Snapshot> adapter = createSnapshotAdapter();
+    JsonAdapter<Snapshot> adapter =
+        MoshiHelper.createMoshiSnapshot(Duration.of(100, ChronoUnit.MILLIS))
+            .adapter(Snapshot.class);
     Snapshot snapshot = createSnapshot();
     CapturedContext context = new CapturedContext();
     CapturedContext.CapturedValue arg1 = CapturedContext.CapturedValue.of("arg1", "int", 42);

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/util/MoshiSnapshotTestHelper.java

@@ -39,6 +39,8 @@
 import java.io.IOException;
 import java.lang.annotation.Annotation;
 import java.lang.reflect.Type;
+import java.time.Duration;
+import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -90,11 +92,18 @@ public static String getValue(CapturedContext.CapturedValue capturedValue) {
   }
 
   public static class SnapshotJsonFactory implements JsonAdapter.Factory {
+    private final Duration captureTimeOut;
+
+    public SnapshotJsonFactory(Duration captureTimeout) {
+      this.captureTimeOut = captureTimeout;
+    }
+
     @Override
     public JsonAdapter<?> create(Type type, Set<? extends Annotation> set, Moshi moshi) {
       if (Types.equals(type, Snapshot.Captures.class)) {
         return new MoshiSnapshotTestHelper.CapturesAdapter(
             moshi,
+            captureTimeOut,
             new CapturedContextAdapter(
                 moshi, new CapturedValueAdapter(), new CapturedThrowableAdapter(moshi)));
       }
@@ -113,8 +122,9 @@ public JsonAdapter<?> create(Type type, Set<? extends Annotation> set, Moshi mos
   }
 
   public static Moshi createMoshiSnapshot() {
+    // By default, for tests we have a capture timeout of 5 seconds
     return new Moshi.Builder()
-        .add(new MoshiSnapshotTestHelper.SnapshotJsonFactory())
+        .add(new MoshiSnapshotTestHelper.SnapshotJsonFactory(Duration.of(5, ChronoUnit.SECONDS)))
         .add(
             DebuggerScript.class,
             new ProbeCondition.ProbeConditionJsonAdapter()) // ProbeDetails in Snapshot
@@ -152,8 +162,9 @@ private static String primitiveArrayToString(Object obj) {
 
   private static class CapturesAdapter extends MoshiSnapshotHelper.CapturesAdapter {
 
-    public CapturesAdapter(Moshi moshi, JsonAdapter<CapturedContext> capturedContextAdapter) {
-      super(moshi, capturedContextAdapter);
+    public CapturesAdapter(
+        Moshi moshi, Duration captureTimeout, JsonAdapter<CapturedContext> capturedContextAdapter) {
+      super(moshi, captureTimeout, capturedContextAdapter);
     }
 
     @Override