Skip to content

Commit fa3ea25

Browse files
amarzialiamarziali
committed
wip2
1 parent 4cea17c commit fa3ea25

6 files changed

Lines changed: 289 additions & 7 deletions

File tree

‎dd-java-agent/instrumentation/java/java-lang/java-lang-22.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk22/FFMApiModule.java‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,38 +1,44 @@
11
package datadog.trace.instrumentation.java.lang.jdk22;
22

3+
import static java.util.Arrays.asList;
4+
35
import com.google.auto.service.AutoService;
46
import datadog.trace.agent.tooling.Instrumenter;
57
import datadog.trace.agent.tooling.InstrumenterModule;
68
import datadog.trace.api.InstrumenterConfig;
9+
import java.util.HashMap;
710
import java.util.List;
811
import java.util.Map;
9-
import java.util.Set;
1012

1113
@AutoService(InstrumenterModule.class)
1214
public class FFMApiModule extends InstrumenterModule.Tracing {
13-
private final Map<String, Set<String>> tracedNativeMethods;
1415
public FFMApiModule() {
1516
super("java-lang-22");
16-
tracedNativeMethods = InstrumenterConfig.get().getTraceNativeMethods();
1717
}
1818

1919
@Override
2020
public Map<String, String> contextStore() {
21-
return super.contextStore();
21+
final Map<String, String> ret = new HashMap<>();
22+
ret.put("java.lang.foreign.SymbolLookup", "java.lang.String");
23+
ret.put("java.lang.foreign.MemorySegment", "java.lang.CharSequence");
24+
return ret;
2225
}
2326

2427
@Override
2528
public boolean isEnabled() {
26-
return super.isEnabled() && !tracedNativeMethods.isEmpty();
29+
return super.isEnabled() && !InstrumenterConfig.get().getTraceNativeMethods().isEmpty();
2730
}
2831

2932
@Override
3033
public List<Instrumenter> typeInstrumentations() {
31-
return super.typeInstrumentations();
34+
return asList(new LinkerInstrumentation(), new SymbolLookupInstrumentation());
3235
}
3336

3437
@Override
3538
public String[] helperClassNames() {
36-
return super.helperClassNames();
39+
return new String[] {
40+
// this could be moved to the boostrap eventually
41+
"datadog.trace.instrumentation.trace_annotation.TraceDecorator",
42+
};
3743
}
3844
}

‎dd-java-agent/instrumentation/java/java-lang/java-lang-22.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk22/LinkerInstrumentation.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
package datadog.trace.instrumentation.java.lang.jdk22;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
6+
7+
import datadog.trace.agent.tooling.Instrumenter;
8+
import net.bytebuddy.description.type.TypeDescription;
9+
import net.bytebuddy.matcher.ElementMatcher;
10+
11+
public class LinkerInstrumentation
12+
implements Instrumenter.ForTypeHierarchy,
13+
Instrumenter.ForBootstrap,
14+
Instrumenter.HasMethodAdvice {
15+
@Override
16+
public String hierarchyMarkerType() {
17+
return null; // bootstrap type
18+
}
19+
20+
@Override
21+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
22+
return implementsInterface(named("java.lang.foreign.Linker"));
23+
}
24+
25+
@Override
26+
public void methodAdvice(MethodTransformer transformer) {
27+
transformer.applyAdvice(
28+
isMethod().and(named("downcallHandle")),
29+
"datadog.trace.instrumentation.java.lang.jdk22.DownCallWrapAdvice");
30+
}
31+
}

‎dd-java-agent/instrumentation/java/java-lang/java-lang-22.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk22/SymbolLookupInstrumentation.java‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package datadog.trace.instrumentation.java.lang.jdk22;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
6+
import static net.bytebuddy.matcher.ElementMatchers.isStatic;
7+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
8+
9+
import datadog.trace.agent.tooling.Instrumenter;
10+
import datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers;
11+
import net.bytebuddy.description.type.TypeDescription;
12+
import net.bytebuddy.matcher.ElementMatcher;
13+
14+
public class SymbolLookupInstrumentation
15+
implements Instrumenter.ForTypeHierarchy,
16+
Instrumenter.ForBootstrap,
17+
Instrumenter.HasMethodAdvice {
18+
private static final String SYMBOL_LOOKUP = "java.lang.foreign.SymbolLookup";
19+
20+
@Override
21+
public String hierarchyMarkerType() {
22+
return null; // bootstrap type
23+
}
24+
25+
@Override
26+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
27+
// instrument both interface and sub-implementations
28+
return implementsInterface(named(SYMBOL_LOOKUP)).or(named(SYMBOL_LOOKUP));
29+
}
30+
31+
@Override
32+
public void methodAdvice(MethodTransformer transformer) {
33+
transformer.applyAdvice(
34+
isMethod().and(isStatic()).and(named("defaultLookup")),
35+
"datadog.trace.instrumentation.java.lang.jdk22.SymbolLookupAdvices$CaptureDefaultLookup");
36+
transformer.applyAdvice(
37+
isMethod()
38+
.and(isStatic())
39+
.and(named("libraryLookup").and(takesArgument(0, named("java.lang.String")))),
40+
"datadog.trace.instrumentation.java.lang.jdk22.SymbolLookupAdvices$CaptureLibraryName");
41+
transformer.applyAdvice(
42+
isMethod()
43+
.and(isStatic())
44+
.and(named("libraryLookup").and(takesArgument(0, named("java.nio.Path")))),
45+
"datadog.trace.instrumentation.java.lang.jdk22.SymbolLookupAdvices$CaptureLibraryPath");
46+
transformer.applyAdvice(
47+
isMethod().and(named("find").and(takesArgument(0, named("java.lang.String")))),
48+
"datadog.trace.instrumentation.java.lang.jdk22.SymbolLookupAdvices$CaptureMemorySegment");
49+
}
50+
}

‎dd-java-agent/instrumentation/java/java-lang/java-lang-22.0/src/main/java25/datadog/trace/instrumentation/java/lang/jdk22/DownCallWrapAdvice.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package datadog.trace.instrumentation.java.lang.jdk22;
2+
3+
4+
import datadog.trace.bootstrap.InstrumentationContext;
5+
import net.bytebuddy.asm.Advice;
6+
import java.lang.foreign.MemorySegment;
7+
import java.lang.invoke.MethodHandle;
8+
import java.lang.invoke.MethodHandles;
9+
10+
public class DownCallWrapAdvice {
11+
@Advice.OnMethodExit(suppress = Throwable.class)
12+
public static void onExit(@Advice.Argument(0) final MemorySegment memorySegment, @Advice.Return(readOnly = false)MethodHandle handle) {
13+
if (memorySegment == null || !Boolean.TRUE.equals(InstrumentationContext.get(MemorySegment.class, Boolean.class).get(memorySegment))) {
14+
return;
15+
}
16+
MethodHandles.
17+
}
18+
}

‎dd-java-agent/instrumentation/java/java-lang/java-lang-22.0/src/main/java25/datadog/trace/instrumentation/java/lang/jdk22/FFMNativeMethodDecorator.java‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
package datadog.trace.instrumentation.java.lang.jdk22;
2+
3+
import datadog.context.Context;
4+
import datadog.context.ContextScope;
5+
import datadog.trace.api.DDSpanTypes;
6+
import datadog.trace.api.cache.DDCache;
7+
import datadog.trace.api.cache.DDCaches;
8+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
9+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
10+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
11+
import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
12+
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
13+
import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator;
14+
import org.slf4j.Logger;
15+
import org.slf4j.LoggerFactory;
16+
17+
import java.lang.invoke.MethodHandle;
18+
import java.lang.invoke.MethodHandles;
19+
import java.lang.invoke.MethodType;
20+
import java.util.Deque;
21+
22+
public final class FFMNativeMethodDecorator extends BaseDecorator {
23+
24+
private static final Logger LOGGER = LoggerFactory.getLogger(FFMNativeMethodDecorator.class);
25+
private static final CharSequence TRACE_FFM = UTF8BytesString.create("trace-ffm");
26+
private static final CharSequence OPERATION_NAME = UTF8BytesString.create("trace.native");
27+
28+
private static final MethodHandle START_SPAN_MH = safeFindStatic("startSpan", MethodType.methodType(void.class, String.class));
29+
private static final MethodHandle END_SPAN_MH = safeFindStatic("endSpan", MethodType.methodType(void.class, String.class)))
30+
31+
public static final FFMNativeMethodDecorator DECORATE = new FFMNativeMethodDecorator();
32+
33+
private static MethodHandle safeFindStatic(String name, MethodType methodType) {
34+
try {
35+
return MethodHandles.lookup().findStatic(FFMNativeMethodDecorator.class, name, methodType);
36+
} catch (Throwable t) {
37+
LOGGER.debug("Cannot find method {} in NativeMethodHandleWrapper", name, t);
38+
return null;
39+
}
40+
}
41+
42+
public static MethodHandle wrap(MethodHandle original, String operationName) {
43+
if (START_SPAN_MH == null || END_SPAN_MH == null) {
44+
return original;
45+
}
46+
MethodType originalType = original.type();
47+
boolean isVoid = originalType.returnType() == void.class;
48+
49+
MethodHandle startSpanMH = START_SPAN_MH.bindTo(operationName);
50+
/*
51+
Return a methodHandle chain that
52+
1. first calls startspans
53+
2. than calls the original
54+
3. As a tryfinally calls the endSpan providing the return value of (1) as argument
55+
4. Eventually drops the return value if the original return was void in order to have a method handle wrapped that's transparent
56+
*/
57+
return null;
58+
59+
}
60+
61+
public static ContextScope startSpan(CharSequence resourceName) {
62+
AgentSpan span = AgentTracer.startSpan(TRACE_FFM.toString(), OPERATION_NAME);
63+
DECORATE.afterStart(span);
64+
span.setResourceName(resourceName);
65+
return AgentTracer.activateSpan(span);
66+
}
67+
68+
public static Object endSpan(Throwable t, ContextScope scope, Object result) {
69+
try {
70+
if (scope != null) {
71+
final AgentSpan span = AgentSpan.fromContext(scope.context());
72+
scope.close();
73+
74+
if (span != null) {
75+
if (t != null) {
76+
DECORATE.onError(span, t);
77+
span.addThrowable(t);
78+
}
79+
80+
span.finish();
81+
}
82+
}
83+
} catch (Throwable ignored) {
84+
85+
}
86+
return result;
87+
}
88+
89+
90+
@Override
91+
protected String[] instrumentationNames() {
92+
return new String[] {TRACE_FFM.toString()};
93+
}
94+
95+
@Override
96+
protected CharSequence spanType() {
97+
return null;
98+
}
99+
100+
@Override
101+
protected CharSequence component() {
102+
return TRACE_FFM;
103+
}
104+
105+
106+
}

‎dd-java-agent/instrumentation/java/java-lang/java-lang-22.0/src/main/java25/datadog/trace/instrumentation/java/lang/jdk22/SymbolLookupAdvices.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
package datadog.trace.instrumentation.java.lang.jdk22;
2+
3+
import datadog.trace.api.InstrumenterConfig;
4+
import datadog.trace.bootstrap.InstrumentationContext;
5+
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
6+
import java.lang.foreign.MemorySegment;
7+
import java.lang.foreign.SymbolLookup;
8+
import java.nio.file.Path;
9+
import java.util.Locale;
10+
import java.util.Optional;
11+
import java.util.Set;
12+
import net.bytebuddy.asm.Advice;
13+
14+
public class SymbolLookupAdvices {
15+
16+
public static class DefaultLookup {
17+
public static class CaptureDefault {
18+
@Advice.OnMethodExit(suppress = Throwable.class)
19+
public static void onExit(@Advice.Return final SymbolLookup symbolLookup) {
20+
if (symbolLookup != null) {
21+
InstrumentationContext.get(SymbolLookup.class, String.class).put(symbolLookup, "default");
22+
}
23+
}
24+
}
25+
26+
public static class CaptureLibraryName {
27+
@Advice.OnMethodExit(suppress = Throwable.class)
28+
public static void onExit(
29+
@Advice.Return final SymbolLookup symbolLookup,
30+
@Advice.Argument(0) final String libraryName) {
31+
if (symbolLookup != null && libraryName != null) {
32+
InstrumentationContext.get(SymbolLookup.class, String.class)
33+
.put(symbolLookup, libraryName.toLowerCase(Locale.ROOT));
34+
}
35+
}
36+
}
37+
}
38+
39+
public static class CaptureLibraryPath {
40+
@Advice.OnMethodExit(suppress = Throwable.class)
41+
public static void onExit(
42+
@Advice.Return final SymbolLookup symbolLookup,
43+
@Advice.Argument(0) final Path libraryPath) {
44+
if (symbolLookup != null && libraryPath != null) {
45+
InstrumentationContext.get(SymbolLookup.class, String.class)
46+
.put(symbolLookup, libraryPath.getFileName().toString().toLowerCase(Locale.ROOT));
47+
}
48+
}
49+
}
50+
51+
public static class CaptureMemorySegment {
52+
@Advice.OnMethodExit(suppress = Throwable.class)
53+
public static void onExit(
54+
@Advice.This final SymbolLookup self,
55+
@Advice.Argument(0) final String name,
56+
@Advice.Return final Optional<MemorySegment> maybeSegment) {
57+
if (name != null && maybeSegment != null && maybeSegment.isPresent()) {
58+
final String libName =
59+
InstrumentationContext.get(SymbolLookup.class, String.class).get(self);
60+
if (libName != null) {
61+
final Set<String> tracedMethods =
62+
InstrumenterConfig.get().getTraceNativeMethods().get(libName);
63+
if (tracedMethods != null && tracedMethods.contains(name)) {
64+
InstrumentationContext.get(MemorySegment.class, CharSequence.class)
65+
.put(maybeSegment.get(), UTF8BytesString.create(libName + "." + name));
66+
}
67+
}
68+
}
69+
}
70+
}
71+
}

0 commit comments

Comments
 (0)