Skip to content

Commit b166862

Browse files
amarzialiamarziali
authored
Use pure FFM to write on service discovery memfd from java 22 (DataDog#10295)
* Use pure FFI to write on service discovery memfd from java 25 * Fix build * refactor
1 parent 92c84e2 commit b166862

7 files changed

Lines changed: 333 additions & 38 deletions

File tree

‎dd-java-agent/agent-tooling/build.gradle‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,11 @@ minimumBranchCoverage = 0.6
99
excludedClassesCoverage += ['datadog.trace.agent.tooling.*']
1010

1111
sourceSets {
12+
register("main_java25") {
13+
java {
14+
srcDirs = [file('src/main/java25')]
15+
}
16+
}
1217
register("test_java11") {
1318
java {
1419
srcDirs = [file('src/test/java11')]
@@ -19,11 +24,17 @@ sourceSets {
1924
srcDirs = [file('src/test/java21')]
2025
}
2126
}
27+
named("main_java25") {
28+
compileClasspath += sourceSets.main.output
29+
runtimeClasspath += sourceSets.main.output
30+
}
2231
named("test") {
2332
compileClasspath += sourceSets.test_java11.output
2433
runtimeClasspath += sourceSets.test_java11.output
2534
compileClasspath += sourceSets.test_java21.output
2635
runtimeClasspath += sourceSets.test_java21.output
36+
compileClasspath += sourceSets.main_java25.output
37+
runtimeClasspath += sourceSets.main_java25.output
2738
}
2839
}
2940

@@ -39,6 +50,7 @@ configurations {
3950
}
4051
}
4152

53+
4254
dependencies {
4355
api(project(':dd-java-agent:agent-bootstrap')) {
4456
exclude group: 'com.datadoghq', module: 'agent-logging'
@@ -53,13 +65,15 @@ dependencies {
5365
implementation group: 'net.java.dev.jna', name: 'jna-platform', version: '5.8.0'
5466

5567
api project(':dd-trace-core')
56-
5768
implementation project(':dd-java-agent:agent-crashtracking')
5869

70+
main_java25Implementation project(':dd-trace-core')
71+
5972
testImplementation project(':dd-java-agent:testing')
6073
testImplementation libs.bytebuddy
6174
testImplementation group: 'com.google.guava', name: 'guava-testlib', version: '20.0'
6275

76+
6377
jmhImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '2.3.5.RELEASE'
6478
}
6579

@@ -68,6 +82,10 @@ jmh {
6882
includeTests = true
6983
}
7084

85+
tasks.named("jar") {
86+
from(sourceSets.main_java25.output)
87+
}
88+
7189
tasks.named("compileJava") { dependsOn 'generateClassNameTries' }
7290
tasks.named("sourcesJar") { dependsOn 'generateClassNameTries' }
7391

@@ -88,6 +106,14 @@ tasks.named("compileTestGroovy") {
88106
" otherwise anonymous class has one `loadClass` accessor's signature has `java.lang.Module`"
89107
)
90108
}
109+
tasks.named("compileMain_java25Java") {
110+
configureCompiler(
111+
it,
112+
25,
113+
JavaVersion.VERSION_1_8,
114+
"Java 25 sourceset for Foreign Function & Memory API, compiled with Java 25 but targeting Java 8 bytecode."
115+
)
116+
}
91117
tasks.named("compileTest_java11Java") {
92118
configureCompiler(it, 11, JavaVersion.VERSION_11)
93119
}

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/TracerInstaller.java‎

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,9 @@
99
import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
1010
import datadog.trace.core.CoreTracer;
1111
import datadog.trace.core.servicediscovery.ForeignMemoryWriter;
12+
import datadog.trace.core.servicediscovery.ForeignMemoryWriterFactory;
1213
import datadog.trace.core.servicediscovery.ServiceDiscovery;
1314
import datadog.trace.core.servicediscovery.ServiceDiscoveryFactory;
14-
import de.thetaphi.forbiddenapis.SuppressForbidden;
1515
import org.slf4j.Logger;
1616
import org.slf4j.LoggerFactory;
1717

@@ -57,20 +57,12 @@ private static ServiceDiscoveryFactory serviceDiscoveryFactory() {
5757
return TracerInstaller::initServiceDiscovery;
5858
}
5959

60-
@SuppressForbidden // intentional use of Class.forName
6160
private static ServiceDiscovery initServiceDiscovery() {
62-
try {
63-
// use reflection to load MemFDUnixWriter so it doesn't get picked up when we
64-
// transitively look for all tracer class dependencies to install in GraalVM via
65-
// VMRuntimeInstrumentation
66-
Class<?> memFdClass =
67-
Class.forName("datadog.trace.agent.tooling.servicediscovery.MemFDUnixWriter");
68-
ForeignMemoryWriter memFd = (ForeignMemoryWriter) memFdClass.getConstructor().newInstance();
69-
return new ServiceDiscovery(memFd);
70-
} catch (Throwable e) {
71-
log.debug("service discovery not supported", e);
72-
return null;
61+
final ForeignMemoryWriter writer = new ForeignMemoryWriterFactory().get();
62+
if (writer != null) {
63+
return new ServiceDiscovery(writer);
7364
}
65+
return null;
7466
}
7567

7668
public static void installGlobalTracer(final CoreTracer tracer) {

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/servicediscovery/MemFDUnixWriter.java‎

Lines changed: 18 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,26 +2,23 @@
22

33
import static datadog.trace.api.telemetry.LogCollector.SEND_TELEMETRY;
44

5-
import com.sun.jna.Library;
6-
import com.sun.jna.Memory;
75
import com.sun.jna.Native;
8-
import com.sun.jna.NativeLong;
9-
import com.sun.jna.Pointer;
106
import datadog.environment.OperatingSystem;
7+
import datadog.environment.SystemProperties;
118
import datadog.trace.core.servicediscovery.ForeignMemoryWriter;
129
import org.slf4j.Logger;
1310
import org.slf4j.LoggerFactory;
1411

15-
public class MemFDUnixWriter implements ForeignMemoryWriter {
12+
abstract class MemFDUnixWriter implements ForeignMemoryWriter {
1613
private static final Logger log = LoggerFactory.getLogger(MemFDUnixWriter.class);
1714

18-
private interface LibC extends Library {
19-
int syscall(int number, Object... args);
15+
protected abstract long syscall(long number, String name, int flags);
2016

21-
NativeLong write(int fd, Pointer buf, NativeLong count);
17+
protected abstract long write(int fd, byte[] payload);
2218

23-
int fcntl(int fd, int cmd, int arg);
24-
}
19+
protected abstract int fcntl(int fd, int cmd, int arg);
20+
21+
protected abstract int getLastError();
2522

2623
// https://elixir.bootlin.com/linux/v6.17.1/source/include/uapi/linux/memfd.h#L8-L9
2724
private static final int MFD_CLOEXEC = 0x0001;
@@ -36,36 +33,33 @@ private interface LibC extends Library {
3633
private static final int F_SEAL_GROW = 0x0004;
3734

3835
@Override
39-
public void write(String fileName, byte[] payload) {
40-
final LibC libc = Native.load("c", LibC.class);
41-
36+
public final void write(String fileName, byte[] payload) {
4237
OperatingSystem.Architecture arch = OperatingSystem.architecture();
4338
int memfdSyscall = getMemfdSyscall(arch);
4439
if (memfdSyscall <= 0) {
45-
log.debug(SEND_TELEMETRY, "service discovery not supported for arch={}", arch);
40+
log.debug(
41+
SEND_TELEMETRY,
42+
"service discovery not supported for arch={}",
43+
SystemProperties.get("os.arch"));
4644
return;
4745
}
48-
int memFd = libc.syscall(memfdSyscall, fileName, MFD_CLOEXEC | MFD_ALLOW_SEALING);
46+
int memFd = (int) syscall(memfdSyscall, fileName, MFD_CLOEXEC | MFD_ALLOW_SEALING);
4947
if (memFd < 0) {
50-
log.warn("{} memfd create failed, errno={}", fileName, Native.getLastError());
48+
log.warn("{} memfd create failed, errno={}", fileName, getLastError());
5149
return;
5250
}
5351

5452
log.debug("{} memfd created (fd={})", fileName, memFd);
5553

56-
Memory buf = new Memory(payload.length);
57-
buf.write(0, payload, 0, payload.length);
58-
59-
NativeLong written = libc.write(memFd, buf, new NativeLong(payload.length));
60-
if (written.longValue() != payload.length) {
54+
long written = write(memFd, payload);
55+
if (written != payload.length) {
6156
log.warn("write to {} memfd failed errno={}", fileName, Native.getLastError());
6257
return;
6358
}
64-
log.debug("wrote {} bytes to memfd {}", written.longValue(), memFd);
65-
int returnCode = libc.fcntl(memFd, F_ADD_SEALS, F_SEAL_SHRINK | F_SEAL_GROW | F_SEAL_SEAL);
59+
log.debug("wrote {} bytes to memfd {}", written, memFd);
60+
int returnCode = fcntl(memFd, F_ADD_SEALS, F_SEAL_SHRINK | F_SEAL_GROW | F_SEAL_SEAL);
6661
if (returnCode == -1) {
6762
log.warn("failed to add seal to {} memfd errno={}", fileName, Native.getLastError());
68-
return;
6963
}
7064
// memfd is not closed to keep it readable for the lifetime of the process.
7165
}

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/servicediscovery/MemFDUnixWriterJNA.java‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package datadog.trace.agent.tooling.servicediscovery;
2+
3+
import com.sun.jna.Library;
4+
import com.sun.jna.Memory;
5+
import com.sun.jna.Native;
6+
import com.sun.jna.NativeLong;
7+
import com.sun.jna.Pointer;
8+
9+
public final class MemFDUnixWriterJNA extends MemFDUnixWriter {
10+
private final LibC libc = Native.load("c", LibC.class);
11+
12+
private interface LibC extends Library {
13+
long syscall(long number, Object... args);
14+
15+
NativeLong write(int fd, Pointer buf, NativeLong count);
16+
17+
int fcntl(int fd, int cmd, int arg);
18+
}
19+
20+
@Override
21+
protected long syscall(long number, String name, int flags) {
22+
return libc.syscall(number, name, flags);
23+
}
24+
25+
@Override
26+
protected long write(int fd, byte[] payload) {
27+
Memory buf = new Memory(payload.length);
28+
buf.write(0, payload, 0, payload.length);
29+
return libc.write(fd, buf, new NativeLong(payload.length)).longValue();
30+
}
31+
32+
@Override
33+
protected int fcntl(int fd, int cmd, int arg) {
34+
return libc.fcntl(fd, cmd, arg);
35+
}
36+
37+
@Override
38+
protected int getLastError() {
39+
return Native.getLastError();
40+
}
41+
}

‎dd-java-agent/agent-tooling/src/main/java25/datadog/trace/agent/tooling/servicediscovery/MemFDUnixWriterFFM.java‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
package datadog.trace.agent.tooling.servicediscovery;
2+
3+
import java.lang.foreign.Arena;
4+
import java.lang.foreign.FunctionDescriptor;
5+
import java.lang.foreign.Linker;
6+
import java.lang.foreign.MemoryLayout;
7+
import java.lang.foreign.MemorySegment;
8+
import java.lang.foreign.StructLayout;
9+
import java.lang.foreign.SymbolLookup;
10+
import java.lang.foreign.ValueLayout;
11+
import java.lang.invoke.MethodHandle;
12+
import org.slf4j.Logger;
13+
import org.slf4j.LoggerFactory;
14+
15+
public class MemFDUnixWriterFFM extends MemFDUnixWriter {
16+
private static final Logger log = LoggerFactory.getLogger(MemFDUnixWriterFFM.class);
17+
18+
// Captured call state layout for errno
19+
private static final StructLayout CAPTURE_STATE_LAYOUT = Linker.Option.captureStateLayout();
20+
private static final long ERRNO_OFFSET =
21+
CAPTURE_STATE_LAYOUT.byteOffset(MemoryLayout.PathElement.groupElement("errno"));
22+
23+
// Function handles - initialized once
24+
private final MethodHandle syscallMH;
25+
private final MethodHandle writeMH;
26+
private final MethodHandle fcntlMH;
27+
28+
private final MemorySegment captureState;
29+
30+
public MemFDUnixWriterFFM() {
31+
final Linker linker = Linker.nativeLinker();
32+
final SymbolLookup LIBC = linker.defaultLookup();
33+
34+
// Allocate memory for capturing errno (need to be alive until the class instance is collected)
35+
this.captureState = Arena.ofAuto().allocate(CAPTURE_STATE_LAYOUT);
36+
37+
// long syscall(long number, ...)
38+
// Note: variadic functions require special handling, we'll use a fixed signature
39+
syscallMH =
40+
linker.downcallHandle(
41+
LIBC.find("syscall").orElseThrow(),
42+
FunctionDescriptor.of(
43+
ValueLayout.JAVA_LONG, // return type: long
44+
ValueLayout.JAVA_LONG, // syscall number
45+
ValueLayout.ADDRESS, // const char* name
46+
ValueLayout.JAVA_INT // int flags
47+
),
48+
Linker.Option.captureCallState("errno"));
49+
50+
// ssize_t write(int fd, const void *buf, size_t count)
51+
writeMH =
52+
linker.downcallHandle(
53+
LIBC.find("write").orElseThrow(),
54+
FunctionDescriptor.of(
55+
ValueLayout.JAVA_LONG, // return type: ssize_t
56+
ValueLayout.JAVA_INT, // int fd
57+
ValueLayout.ADDRESS, // const void* buf
58+
ValueLayout.JAVA_LONG // size_t count
59+
),
60+
Linker.Option.captureCallState("errno"));
61+
62+
// int fcntl(int fd, int cmd, ... /* arg */)
63+
fcntlMH =
64+
linker.downcallHandle(
65+
LIBC.find("fcntl").orElseThrow(),
66+
FunctionDescriptor.of(
67+
ValueLayout.JAVA_INT, // return type: int
68+
ValueLayout.JAVA_INT, // int fd
69+
ValueLayout.JAVA_INT, // int cmd
70+
ValueLayout.JAVA_INT // int arg
71+
),
72+
Linker.Option.captureCallState("errno"));
73+
}
74+
75+
@Override
76+
protected long syscall(long number, String name, int flags) {
77+
try (Arena arena = Arena.ofConfined()) {
78+
// Allocate native string for file name
79+
MemorySegment fileNameSegment = arena.allocateFrom(name);
80+
// Call memfd_create via syscall, passing captureState as first arg
81+
return (long) syscallMH.invoke(captureState, (long) number, fileNameSegment, flags);
82+
} catch (Throwable t) {
83+
log.error("Unable to make a syscall through FFM", t);
84+
return -1;
85+
}
86+
}
87+
88+
@Override
89+
protected long write(int fd, byte[] payload) {
90+
try (Arena arena = Arena.ofConfined()) {
91+
// Allocate native memory for payload
92+
MemorySegment buffer = arena.allocate(payload.length);
93+
MemorySegment.copy(payload, 0, buffer, ValueLayout.JAVA_BYTE, 0, payload.length);
94+
95+
// Write payload to memfd, passing captureState as first arg
96+
return (long) writeMH.invoke(captureState, fd, buffer, (long) payload.length);
97+
} catch (Throwable t) {
98+
log.error("Unable to make a write call through FFM", t);
99+
return -1;
100+
}
101+
}
102+
103+
@Override
104+
protected int fcntl(int fd, int cmd, int arg) {
105+
try {
106+
return (int) fcntlMH.invoke(captureState, fd, cmd, arg);
107+
} catch (Throwable t) {
108+
log.error("Unable to make a fcntl call through FFM", t);
109+
return -1;
110+
}
111+
}
112+
113+
@Override
114+
protected int getLastError() {
115+
try {
116+
// Read errno from the captured state memory segment
117+
return captureState.get(ValueLayout.JAVA_INT, ERRNO_OFFSET);
118+
} catch (Throwable t) {
119+
log.error("Unable to read errno from captured state", t);
120+
return -1;
121+
}
122+
}
123+
}

0 commit comments

Comments
 (0)