Use jvmstat for JDKs 9+ programmatically

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

@@ -44,6 +44,7 @@
 import datadog.trace.util.AgentTaskScheduler;
 import datadog.trace.util.AgentThreadFactory.AgentThread;
 import datadog.trace.util.throwable.FatalAgentMisconfigurationError;
+import de.thetaphi.forbiddenapis.SuppressForbidden;
 import java.lang.instrument.Instrumentation;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
@@ -291,6 +292,8 @@ public static void start(
     codeOriginEnabled = isFeatureEnabled(AgentFeature.CODE_ORIGIN);
     agentlessLogSubmissionEnabled = isFeatureEnabled(AgentFeature.AGENTLESS_LOG_SUBMISSION);
 
+    patchJPSAccess(inst);
+
     if (profilingEnabled) {
       if (!isOracleJDK8()) {
         // Profiling agent startup code is written in a way to allow `startProfilingAgent` be called
@@ -420,6 +423,20 @@ private static void injectAgentArgsConfig(String agentArgs) {
     }
   }
 
+  @SuppressForbidden
+  public static void patchJPSAccess(Instrumentation inst) {
+    if (Platform.isJavaVersionAtLeast(9)) {
+      // Unclear if supported for J9, may need to revisit
+      try {
+        Class.forName("datadog.trace.util.JPMSJPSAccess")
+            .getMethod("patchModuleAccess")
+            .invoke(inst);
+      } catch (Exception e) {
+        log.warn("Failed to patch module access for jvmstat");
+      }
+    }
+  }
+
   public static void shutdown(final boolean sync) {
     StaticEventLogger.end("Agent");
     StaticEventLogger.stop();

dd-smoke-tests/profiling-integration-tests/src/test/java/datadog/smoketest/JFRBasedProfilingIntegrationTest.java

@@ -507,6 +507,8 @@ void testShutdown(final TestInfo testInfo) throws Exception {
             assertTrue(
                 targetProcess.waitFor(
                     duration + PROFILING_UPLOAD_TIMEOUT_SECONDS + 1, TimeUnit.SECONDS));
+            assertTrue(
+                checkLogLines(logFilePath, it -> it.contains("Successfully invoked jvmstat")));
           } finally {
             if (targetProcess != null) {
               targetProcess.destroyForcibly();

internal-api/build.gradle

@@ -175,6 +175,7 @@ excludedClassesCoverage += [
   "datadog.trace.util.ComparableVersion.LongItem",
   "datadog.trace.util.ComparableVersion.StringItem",
   "datadog.trace.util.ConcurrentEnumMap",
+  "datadog.trace.util.JPSUtils",
   "datadog.trace.util.MethodHandles",
   "datadog.trace.util.PidHelper",
   "datadog.trace.util.PidHelper.Fallback",

internal-api/internal-api-9/build.gradle

@@ -41,6 +41,10 @@ forbiddenApisMain {
   failOnMissingClasses = false
 }
 
+forbiddenApisMain_java11 {
+  failOnMissingClasses = false
+}
+
 idea {
   module {
     jdkName = '11'

internal-api/internal-api-9/src/main_java11/java/datadog/trace/util/JPMSJPSAccess.java

@@ -0,0 +1,26 @@
+package datadog.trace.util;
+
+import java.lang.instrument.Instrumentation;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+
+public class JPMSJPSAccess {
+  public static void patchModuleAccess(Instrumentation inst) {
+    Module unnamedModule = JPMSJPSAccess.class.getClassLoader().getUnnamedModule();
+    Module jvmstatModule = ModuleLayer.boot().findModule("jdk.internal.jvmstat").orElse(null);
+
+    if (jvmstatModule != null) {
+      Map<String, Set<Module>> extraOpens = Map.of("sun.jvmstat.monitor", Set.of(unnamedModule));
+
+      // Redefine the module
+      inst.redefineModule(
+          jvmstatModule,
+          Collections.emptySet(),
+          extraOpens,
+          extraOpens,
+          Collections.emptySet(),
+          Collections.emptyMap());
+    }
+  }
+}

internal-api/src/main/java/datadog/trace/util/JPSUtils.java

@@ -0,0 +1,33 @@
+package datadog.trace.util;
+
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.lang.reflect.Method;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public final class JPSUtils {
+  private static final Logger log = LoggerFactory.getLogger(JPSUtils.class);
+
+  @SuppressForbidden
+  public static Set<String> getVMPids() {
+    Set<String> vmPids = new HashSet<>();
+    try {
+      Class<?> monitoredHostClass = Class.forName("sun.jvmstat.monitor.MonitoredHost");
+      Method getMonitoredHostMethod =
+          monitoredHostClass.getDeclaredMethod("getMonitoredHost", String.class);
+      Object vmHost = getMonitoredHostMethod.invoke(null, "localhost");
+      for (Integer vmPid :
+          (List<Integer>) monitoredHostClass.getDeclaredMethod("activeVms").invoke(vmHost)) {
+        vmPids.add(vmPid.toString());
+      }
+      log.debug("Successfully invoked jvmstat");
+    } catch (Exception e) {
+      log.debug("Failed to invoke jvmstat with exception ", e);
+      return null;
+    }
+    return vmPids;
+  }
+}

internal-api/src/main/java/datadog/trace/util/PidHelper.java

@@ -65,6 +65,11 @@ private static String findPid() {
   }
 
   public static Set<String> getJavaPids() {
+    // Attempt to use jvmstat directly, fall through to jps process fork strategy
+    Set<String> directlyObtainedPids = JPSUtils.getVMPids();
+    if (directlyObtainedPids != null) {
+      return directlyObtainedPids;
+    }
     // there is no supported Java API to achieve this
     // one could use sun.jvmstat.monitor.MonitoredHost but it is an internal API and can go away at
     // any time -