Add CompletableFuture in WellKnownClasses

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

@@ -1,6 +1,12 @@
 package datadog.trace.bootstrap.debugger.util;
 
+import static datadog.trace.api.telemetry.LogCollector.EXCLUDE_TELEMETRY;
+import static java.lang.invoke.MethodType.methodType;
+
+import datadog.environment.JavaVirtualMachine;
 import datadog.trace.bootstrap.debugger.CapturedContext;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.Arrays;
@@ -14,6 +20,7 @@
 import java.util.OptionalInt;
 import java.util.OptionalLong;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 import java.util.function.ToLongFunction;
 import org.slf4j.Logger;
@@ -144,12 +151,18 @@ public class WellKnownClasses {
       OPTIONALDOUBLE_SPECIAL_FIELDS = new HashMap<>();
   private static final Map<String, Function<Object, CapturedContext.CapturedValue>>
       OPTIONALLONG_SPECIAL_FIELDS = new HashMap<>();
+  private static final Map<String, Function<Object, CapturedContext.CapturedValue>>
+      COMPLETABLEFUTURE_SPECIAL_FIELDS = new HashMap<>();
 
   static {
     OPTIONAL_SPECIAL_FIELDS.put("value", OptionalFields::value);
     OPTIONALINT_SPECIAL_FIELDS.put("value", OptionalFields::valueInt);
     OPTIONALDOUBLE_SPECIAL_FIELDS.put("value", OptionalFields::valueDouble);
     OPTIONALLONG_SPECIAL_FIELDS.put("value", OptionalFields::valueLong);
+    if (JavaVirtualMachine.isJavaVersionAtLeast(19)) {
+      // Future::resultNow method is available since JDK 19
+      COMPLETABLEFUTURE_SPECIAL_FIELDS.put("result", CompletableFutureFields::result);
+    }
   }
 
   static {
@@ -158,6 +171,10 @@ public class WellKnownClasses {
     SPECIAL_TYPE_ACCESS.put(OptionalInt.class, OPTIONALINT_SPECIAL_FIELDS);
     SPECIAL_TYPE_ACCESS.put(OptionalDouble.class, OPTIONALDOUBLE_SPECIAL_FIELDS);
     SPECIAL_TYPE_ACCESS.put(OptionalLong.class, OPTIONALLONG_SPECIAL_FIELDS);
+    if (JavaVirtualMachine.isJavaVersionAtLeast(19)) {
+      // Future::resultNow method is available since JDK 19
+      SPECIAL_TYPE_ACCESS.put(CompletableFuture.class, COMPLETABLEFUTURE_SPECIAL_FIELDS);
+    }
   }
 
   private static final Map<String, Function<Object, CapturedContext.CapturedValue>>
@@ -405,4 +422,36 @@ public static CapturedContext.CapturedValue valueLong(Object o) {
           "value", Long.TYPE.getTypeName(), ((OptionalLong) o).orElse(0L));
     }
   }
+
+  private static class CompletableFutureFields {
+    private static final MethodHandle RESULT_NOW;
+
+    static {
+      MethodHandle methodHandle = null;
+      try {
+        MethodHandles.Lookup lookup = MethodHandles.lookup();
+        methodHandle =
+            lookup.findVirtual(CompletableFuture.class, "resultNow", methodType(Object.class));
+      } catch (Exception e) {
+        LOGGER.debug(EXCLUDE_TELEMETRY, "Looking up CompletableFuture::resultNow failed: ", e);
+      }
+      RESULT_NOW = methodHandle;
+    }
+
+    public static CapturedContext.CapturedValue result(Object o) {
+      if (RESULT_NOW == null) {
+        throw new UnsupportedOperationException("CompletableFuture::resultNow not available");
+      }
+      try {
+        CompletableFuture<?> future = (CompletableFuture<?>) o;
+        // need to check with isDone() to avoid getting exception if null.
+        // Known benign rare race condition result != null => result == null
+        // between isDone() and resultNow()
+        Object result = future.isDone() ? RESULT_NOW.invokeExact((future)) : null;
+        return CapturedContext.CapturedValue.of("result", Object.class.getTypeName(), result);
+      } catch (Throwable t) {
+        throw new RuntimeException(t);
+      }
+    }
+  }
 }

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

@@ -30,6 +30,7 @@
 import com.datadog.debugger.util.MoshiSnapshotHelper;
 import com.datadog.debugger.util.MoshiSnapshotTestHelper;
 import com.squareup.moshi.JsonAdapter;
+import datadog.environment.JavaVirtualMachine;
 import datadog.trace.bootstrap.debugger.CapturedContext;
 import datadog.trace.bootstrap.debugger.CapturedStackFrame;
 import datadog.trace.bootstrap.debugger.DebuggerContext;
@@ -63,6 +64,7 @@
 import java.util.OptionalLong;
 import java.util.Random;
 import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.LockSupport;
 import org.junit.jupiter.api.Assertions;
@@ -306,6 +308,7 @@ static class WellKnownClasses {
     StackTraceElement element = new StackTraceElement("Foo", "bar", "foo.java", 42);
     File file = new File("/tmp/foo");
     Path path = file.toPath();
+    CompletableFuture<String> future = CompletableFuture.completedFuture("FutureCompleted!");
   }
 
   @Test
@@ -385,6 +388,12 @@ public void wellKnownClasses() throws IOException {
     assertPrimitiveValue(objLocalFields, "file", File.class.getTypeName(), "/tmp/foo");
     // path
     assertPrimitiveValue(objLocalFields, "path", "sun.nio.fs.UnixPath", "/tmp/foo");
+    if (JavaVirtualMachine.isJavaVersionAtLeast(19)) {
+      Map<String, Object> future = (Map<String, Object>) objLocalFields.get("future");
+      assertComplexClass(future, CompletableFuture.class.getTypeName());
+      Map<String, Object> futureFields = (Map<String, Object>) future.get(FIELDS);
+      assertPrimitiveValue(futureFields, "result", String.class.getTypeName(), "FutureCompleted!");
+    }
   }
 
   @Test