Add tracing support for native method calls via Java FFM API (#10718)

wip

wip2

Trace native methods calls using FFM APIs

Add more tests

try to add linux tests

cleanup

Add config stuff

Add measured methods

Improve the implementation

refactorise some code

leftover

Narrow allowed packages to reduce startup time

Restrict allowed packages to  jdk.internal.loader.NativeLibraries

better trie

Better instrumentation name

Merge branch 'master' into andrea.marziali/ffm-instrument

Merge branch 'master' into andrea.marziali/ffm-instrument

Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>

Author: amarziali

dd-java-agent/agent-bootstrap/src/main/java11/datadog/trace/bootstrap/instrumentation/ffm/FFMNativeMethodDecorator.java

@@ -0,0 +1,216 @@
+package datadog.trace.bootstrap.instrumentation.ffm;
+
+import datadog.context.ContextScope;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+import java.util.Set;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public final class FFMNativeMethodDecorator extends BaseDecorator {
+  private static final Logger LOGGER = LoggerFactory.getLogger(FFMNativeMethodDecorator.class);
+  private static final CharSequence TRACE_FFM = UTF8BytesString.create("trace-ffm");
+  private static final CharSequence OPERATION_NAME = UTF8BytesString.create("trace.native");
+
+  private static final MethodHandle START_SPAN_MH =
+      safeFindStatic(
+          "startSpan",
+          MethodType.methodType(ContextScope.class, CharSequence.class, boolean.class));
+  private static final MethodHandle END_SPAN_MH =
+      safeFindStatic(
+          "endSpan",
+          MethodType.methodType(Object.class, Throwable.class, ContextScope.class, Object.class));
+
+  public static final FFMNativeMethodDecorator DECORATE = new FFMNativeMethodDecorator();
+
+  private static MethodHandle safeFindStatic(String name, MethodType methodType) {
+    try {
+      return MethodHandles.lookup().findStatic(FFMNativeMethodDecorator.class, name, methodType);
+    } catch (Throwable t) {
+      LOGGER.debug("Cannot find method {} in NativeMethodHandleWrapper", name, t);
+      return null;
+    }
+  }
+
+  public static MethodHandle wrap(
+      final MethodHandle original, final String libraryName, final String methodName) {
+    if (START_SPAN_MH == null || END_SPAN_MH == null) {
+      return original;
+    }
+    try {
+      MethodType originalType = original.type();
+      boolean isVoid = originalType.returnType() == void.class;
+
+      // We need the ContextScope to be visible to the finally block.
+      // Easiest way is to artificially prepend it to the target signature.
+      // The added parameter is ignored by the original handle.
+      // originalWithScope: (ContextScope, args...) -> R
+      MethodHandle originalWithScope = MethodHandles.dropArguments(original, 0, ContextScope.class);
+
+      /*
+       * Build the cleanup handle used by MethodHandles.tryFinally.
+       *
+       * tryFinally has a strict calling convention:
+       *   - void target  -> cleanup(Throwable, ContextScope, args...)
+       *   - non-void     -> cleanup(Throwable, R, ContextScope, args...)
+       *
+       * END_SPAN_MH is (Throwable, ContextScope, Object) -> Object,
+       * so we need to reshape it to match what tryFinally expects.
+       */
+      MethodHandle cleanup;
+
+      if (isVoid) {
+        // No return value: bind `null` as the result argument.
+        MethodHandle endWithNull = MethodHandles.insertArguments(END_SPAN_MH, 2, (Object) null);
+
+        // Make it accept the original arguments even though they are unused.
+        MethodHandle endDropped =
+            MethodHandles.dropArguments(endWithNull, 2, originalType.parameterList());
+
+        // tryFinally requires void return for void targets.
+        cleanup = endDropped.asType(endDropped.type().changeReturnType(void.class));
+
+      } else {
+        /*
+         * Non-void case:
+         * tryFinally will call cleanup as:
+         *   (Throwable, returnValue, ContextScope, args...)
+         *
+         * END_SPAN_MH expects:
+         *   (Throwable, ContextScope, result)
+         *
+         * So we first permute parameters to swap returnValue and ContextScope.
+         */
+        MethodHandle endPermuted =
+            MethodHandles.permuteArguments(
+                END_SPAN_MH,
+                MethodType.methodType(
+                    Object.class, Throwable.class, Object.class, ContextScope.class),
+                0,
+                2,
+                1);
+
+        // Accept original arguments (unused) after the required ones.
+        MethodHandle endDropped =
+            MethodHandles.dropArguments(endPermuted, 3, originalType.parameterList());
+
+        // Adapt return and result parameter types to match the original signature.
+        MethodType cleanupType =
+            endDropped
+                .type()
+                .changeParameterType(1, originalType.returnType())
+                .changeReturnType(originalType.returnType());
+
+        cleanup = endDropped.asType(cleanupType);
+      }
+
+      // Wrap the original in try/finally semantics.
+      // Resulting handle:
+      //   (ContextScope, args...) -> R
+      MethodHandle withFinally = MethodHandles.tryFinally(originalWithScope, cleanup);
+
+      // Precompute span metadata so we don't redo the lookup per invocation.
+      final CharSequence resourceName = resourceNameFor(libraryName, methodName);
+      final boolean methodMeasured = isMethodMeasured(libraryName, methodName);
+
+      // Bind both arguments to startSpan.
+      // After binding: () -> ContextScope
+      MethodHandle boundStart =
+          MethodHandles.insertArguments(START_SPAN_MH, 0, resourceName, methodMeasured);
+
+      // Make it look like it takes the same arguments as the original,
+      // even though they are ignored.
+      // (args...) -> ContextScope
+      MethodHandle startCombiner =
+          MethodHandles.dropArguments(boundStart, 0, originalType.parameterList());
+
+      /*
+       * foldArguments wires it all together:
+       *
+       *   scope = startCombiner(args...)
+       *   return withFinally(scope, args...)
+       *
+       * Final shape matches the original:
+       *   (args...) -> R
+       */
+      return MethodHandles.foldArguments(withFinally, startCombiner);
+
+    } catch (Throwable t) {
+      LOGGER.debug(
+          "Cannot wrap method handle for library {} and method {}", libraryName, methodName, t);
+      return original;
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static ContextScope startSpan(CharSequence resourceName, boolean methodMeasured) {
+    AgentSpan span = AgentTracer.startSpan(TRACE_FFM.toString(), OPERATION_NAME);
+    DECORATE.afterStart(span);
+    span.setResourceName(resourceName);
+    if (methodMeasured) {
+      span.setMeasured(true);
+    }
+    return AgentTracer.activateSpan(span);
+  }
+
+  @SuppressWarnings("unused")
+  public static Object endSpan(Throwable t, ContextScope scope, Object result) {
+    try {
+      if (scope != null) {
+        final AgentSpan span = AgentSpan.fromContext(scope.context());
+        scope.close();
+
+        if (span != null) {
+          if (t != null) {
+            DECORATE.onError(span, t);
+            span.addThrowable(t);
+          }
+          span.finish();
+        }
+      }
+    } catch (Throwable ignored) {
+
+    }
+    return result;
+  }
+
+  public static boolean isMethodTraced(final String library, final String method) {
+    return matches(InstrumenterConfig.get().getTraceNativeMethods().get(library), method);
+  }
+
+  public static boolean isMethodMeasured(final String library, final String method) {
+    return matches(InstrumenterConfig.get().getMeasureNativeMethods().get(library), method);
+  }
+
+  public static CharSequence resourceNameFor(final String library, final String method) {
+    return UTF8BytesString.create(library + "." + method);
+  }
+
+  private static boolean matches(final Set<String> allows, final String method) {
+    if (allows == null) {
+      return false;
+    }
+    return allows.contains(method) || allows.contains("*");
+  }
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {TRACE_FFM.toString()};
+  }
+
+  @Override
+  protected CharSequence spanType() {
+    return null;
+  }
+
+  @Override
+  protected CharSequence component() {
+    return TRACE_FFM;
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java11/datadog/trace/bootstrap/instrumentation/ffm/NativeLibraryHelper.java

@@ -0,0 +1,35 @@
+package datadog.trace.bootstrap.instrumentation.ffm;
+
+import datadog.trace.api.Pair;
+import java.io.File;
+import java.util.Locale;
+import java.util.concurrent.ConcurrentHashMap;
+
+public final class NativeLibraryHelper {
+  // this map is unlimited. However, the number of entries depends on the configured methods we want
+  // to trace.
+  private static final ConcurrentHashMap<Long, Pair<String, String>> SYMBOLS_MAP =
+      new ConcurrentHashMap<>();
+
+  private NativeLibraryHelper() {}
+
+  public static void onSymbolLookup(
+      final String libraryName, final String symbol, final long address) {
+    if (libraryName != null && !libraryName.isEmpty()) {
+      if (FFMNativeMethodDecorator.isMethodTraced(libraryName, symbol)) {
+        SYMBOLS_MAP.put(address, Pair.of(libraryName, symbol));
+      }
+    }
+  }
+
+  public static Pair<String, String> reverseResolveLibraryAndSymbol(long address) {
+    return SYMBOLS_MAP.get(address);
+  }
+
+  public static String extractLibraryName(String fullPath) {
+    String libraryName = new File(fullPath).getName().toLowerCase(Locale.ROOT);
+    int dot = libraryName.lastIndexOf('.');
+    libraryName = (dot > 0) ? libraryName.substring(0, dot) : libraryName;
+    return libraryName;
+  }
+}

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java

@@ -106,7 +106,7 @@ public ClassVisitor wrap(
       final MethodList<?> methods,
       final int writerFlags,
       final int readerFlags) {
-    return new ClassVisitor(Opcodes.ASM8, classVisitor) {
+    return new ClassVisitor(Opcodes.ASM9, classVisitor) {
 
       private final boolean frames =
           implementationContext.getClassFileVersion().isAtLeast(ClassFileVersion.JAVA_V6);

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextRequestRewriter.java

@@ -84,7 +84,7 @@ public ClassVisitor wrap(
       final MethodList<?> methods,
       final int writerFlags,
       final int readerFlags) {
-    return new ClassVisitor(Opcodes.ASM8, classVisitor) {
+    return new ClassVisitor(Opcodes.ASM9, classVisitor) {
       @Override
       public MethodVisitor visitMethod(
           final int access,

dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie

@@ -44,6 +44,10 @@
 1 io.opentelemetry.javaagent.*
 1 java.*
 0 java.lang.ClassLoader
+0 jdk.internal.foreign.abi.AbstractLinker
+0 jdk.internal.loader.NativeLibrary
+0 jdk.internal.loader.RawNativeLibraries$*
+0 jdk.internal.loader.NativeLibraries$*
 # allow exception profiling instrumentation
 0 java.lang.Exception
 0 java.lang.Error

dd-java-agent/instrumentation/java/java-lang/java-lang-22.0/build.gradle

@@ -0,0 +1,38 @@
+plugins {
+  id 'idea'
+}
+
+muzzle {
+  pass {
+    coreJdk('25')
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+apply from: "$rootDir/gradle/slf4j-simple.gradle"
+
+tracerJava {
+  addSourceSetFor(JavaVersion.VERSION_25)
+}
+
+testJvmConstraints {
+  minJavaVersion = JavaVersion.VERSION_25
+}
+
+idea {
+  module {
+    jdkName = '25'
+  }
+}
+
+
+tasks.named("compileMain_java25Java", JavaCompile) {
+  configureCompiler(it, 25, JavaVersion.VERSION_1_8)
+}
+
+tasks.named("compileTestGroovy", GroovyCompile) {
+  configureCompiler(it, 25)
+}
+dependencies {
+  implementation project(':dd-java-agent:instrumentation:datadog:tracing:trace-annotation')
+}