Use jvmstat for JDKs 9+ programmatically (#8641)

Author: Matt

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

@@ -3,6 +3,7 @@
 import static datadog.trace.api.ConfigDefaults.DEFAULT_STARTUP_LOGS_ENABLED;
 import static datadog.trace.api.Platform.isJavaVersionAtLeast;
 import static datadog.trace.api.Platform.isOracleJDK8;
+import static datadog.trace.api.telemetry.LogCollector.SEND_TELEMETRY;
 import static datadog.trace.bootstrap.Library.WILDFLY;
 import static datadog.trace.bootstrap.Library.detectLibraries;
 import static datadog.trace.util.AgentThreadFactory.AgentThread.JMX_STARTUP;
@@ -44,6 +45,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 +293,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 +424,23 @@ 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", Instrumentation.class)
+            .invoke(null, inst);
+      } catch (Exception e) {
+        log.debug(
+            SEND_TELEMETRY,
+            "Failed to patch module access for jvmstat and Java version "
+                + Platform.getRuntimeVersion());
+      }
+    }
+  }
+
   public static void shutdown(final boolean sync) {
     StaticEventLogger.end("Agent");
     StaticEventLogger.stop();

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/src/main/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 = ClassLoader.getSystemClassLoader().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/internal-api-9/src/test/groovy/datadog/trace/util/PidHelperTest.groovy

@@ -1,11 +1,21 @@
 package datadog.trace.util
 
 import datadog.trace.test.util.DDSpecification
+import net.bytebuddy.agent.ByteBuddyAgent
 
 class PidHelperTest extends DDSpecification {
 
   def "PID is available everywhere we test"() {
     expect:
     !PidHelper.getPid().isEmpty()
   }
+
+  def "JPS via jvmstat is used when possible"() {
+    when:
+    def inst = ByteBuddyAgent.install()
+    JPMSJPSAccess.patchModuleAccess(inst)
+
+    then:
+    JPSUtils.VMPids != null
+  }
 }

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

@@ -0,0 +1,25 @@
+package datadog.trace.util;
+
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.lang.reflect.Method;
+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() {
+    try {
+      Class<?> monitoredHostClass = Class.forName("sun.jvmstat.monitor.MonitoredHost");
+      Method getMonitoredHostMethod =
+          monitoredHostClass.getDeclaredMethod("getMonitoredHost", String.class);
+      Object vmHost = getMonitoredHostMethod.invoke(null, "localhost");
+      return (Set<String>) monitoredHostClass.getDeclaredMethod("activeVms").invoke(vmHost);
+    } catch (Exception e) {
+      log.debug("Failed to invoke jvmstat with exception ", e);
+      return null;
+    }
+  }
+}

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 -