Skip to content

Commit cde4fbf

Browse files
daniel-mohedanodaniel-mohedano
authored
feat: avoid final field modifications in junit instrumentation (#11752)
feat: avoid final field modifications in junit instrumentation Merge branch 'master' into daniel.mohedano/jep-500-junit Co-authored-by: daniel.mohedano <daniel.mohedano@datadoghq.com>
1 parent a4087b6 commit cde4fbf

4 files changed

Lines changed: 133 additions & 31 deletions

File tree

‎dd-java-agent/instrumentation/junit/junit-5/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/JUnitPlatformUtils.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,18 @@ private JUnitPlatformUtils() {}
9999
private static final MethodHandles METHOD_HANDLES =
100100
new MethodHandles(ClassLoaderUtils.getDefaultClassLoader());
101101

102+
/**
103+
* Loads a class by name from the default class loader, returning {@code null} if it is absent.
104+
*/
105+
@Nullable
106+
public static Class<?> loadClass(String className) {
107+
try {
108+
return ClassLoaderUtils.getDefaultClassLoader().loadClass(className);
109+
} catch (Throwable t) {
110+
return null;
111+
}
112+
}
113+
102114
/*
103115
* We have to support older versions of JUnit 5 that do not have certain methods that we would
104116
* like to use. We try to get method handles in runtime, and if we fail to do it there's a

‎dd-java-agent/instrumentation/junit/junit-5/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/execution/JUnit5ExecutionInstrumentation.java‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@
3030
import org.junit.platform.engine.TestDescriptor;
3131
import org.junit.platform.engine.support.hierarchical.EngineExecutionContext;
3232
import org.junit.platform.engine.support.hierarchical.HierarchicalTestExecutorService;
33-
import org.junit.platform.engine.support.hierarchical.Node;
3433
import org.junit.platform.engine.support.hierarchical.ThrowableCollector;
3534

3635
@AutoService(InstrumenterModule.class)
@@ -161,12 +160,13 @@ public static Boolean execute(@Advice.This HierarchicalTestExecutorService.TestT
161160
EngineExecutionContext parentContext = taskHandle.getParentContext();
162161
TestDescriptorHandle descriptorHandle = new TestDescriptorHandle(testDescriptor);
163162

163+
HierarchicalTestExecutorService.TestTask currentTask = testTask;
164164
int retryAttemptIdx = 0;
165165
while (true) {
166166
factory.setSuppressFailures(executionPolicy.suppressFailures());
167167

168168
CallDepthThreadLocalMap.incrementCallDepth(HierarchicalTestExecutorService.TestTask.class);
169-
testTask.execute();
169+
currentTask.execute();
170170
CallDepthThreadLocalMap.decrementCallDepth(HierarchicalTestExecutorService.TestTask.class);
171171

172172
factory.setSuppressFailures(false); // restore default behavior
@@ -185,13 +185,12 @@ public static Boolean execute(@Advice.This HierarchicalTestExecutorService.TestT
185185
JUnitPlatformUtils.RETRY_DESCRIPTOR_ID_SUFFIX, String.valueOf(++retryAttemptIdx));
186186

187187
TestDescriptor retryDescriptor = descriptorHandle.withIdSuffix(suffix);
188-
taskHandle.setTestDescriptor(retryDescriptor);
189-
taskHandle.setNode((Node<?>) retryDescriptor);
190188
taskHandle.getListener().dynamicTestRegistered(retryDescriptor);
191189
TestEventsHandlerHolder.setExecutionTracker(retryDescriptor, executionPolicy);
192190

193-
// restore parent context, since the reference is overwritten with null after execution
194-
taskHandle.setParentContext(parentContext);
191+
// build a fresh task for the retry and reuse the original parent context, since execution
192+
// overwrites it with null
193+
currentTask = taskHandle.createRetryTask(retryDescriptor, parentContext);
195194
}
196195
return Boolean.TRUE; // skip original method execution
197196
}

‎dd-java-agent/instrumentation/junit/junit-5/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/execution/TestDescriptorHandle.java‎

Lines changed: 65 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
package datadog.trace.instrumentation.junit5.execution;
22

33
import datadog.trace.agent.tooling.muzzle.Reference;
4+
import datadog.trace.instrumentation.junit5.JUnitPlatformUtils;
45
import datadog.trace.util.MethodHandles;
56
import datadog.trace.util.UnsafeUtils;
67
import java.lang.invoke.MethodHandle;
78
import java.util.Collection;
89
import java.util.Collections;
910
import java.util.Map;
11+
import java.util.function.UnaryOperator;
1012
import org.junit.platform.commons.util.ClassLoaderUtils;
1113
import org.junit.platform.engine.TestDescriptor;
1214
import org.junit.platform.engine.UniqueId;
@@ -17,8 +19,28 @@ public class TestDescriptorHandle {
1719
private static final MethodHandles METHOD_HANDLES =
1820
new MethodHandles(ClassLoaderUtils.getDefaultClassLoader());
1921

20-
private static final MethodHandle UNIQUE_ID_SETTER =
21-
METHOD_HANDLES.privateFieldSetter(AbstractTestDescriptor.class, "uniqueId");
22+
private static final String JUPITER_TEST_DESCRIPTOR =
23+
"org.junit.jupiter.engine.descriptor.JupiterTestDescriptor";
24+
private static final Class<?> JUPITER_TEST_DESCRIPTOR_CLASS =
25+
JUnitPlatformUtils.loadClass(JUPITER_TEST_DESCRIPTOR);
26+
27+
/** {@code JupiterTestDescriptor#copyIncludingDescendants(UnaryOperator<UniqueId>)} (5.13+) */
28+
private static final MethodHandle COPY_INCLUDING_DESCENDANTS =
29+
METHOD_HANDLES.method(
30+
JUPITER_TEST_DESCRIPTOR, "copyIncludingDescendants", UnaryOperator.class);
31+
32+
// Legacy fallback used when copyIncludingDescendants is unavailable.
33+
// Overwrites the final unique ID field by reflection. Lazily created to avoid JEP 500 warnings.
34+
private static volatile MethodHandle uniqueIdSetter;
35+
36+
private static MethodHandle uniqueIdSetter() {
37+
MethodHandle handle = uniqueIdSetter;
38+
if (handle == null) {
39+
handle = METHOD_HANDLES.privateFieldSetter(AbstractTestDescriptor.class, "uniqueId");
40+
uniqueIdSetter = handle;
41+
}
42+
return handle;
43+
}
2244

2345
public static final class MuzzleHelper {
2446
public static Collection<? extends Reference> compileReferences() {
@@ -33,24 +55,57 @@ public static Collection<? extends Reference> compileReferences() {
3355

3456
public TestDescriptorHandle(TestDescriptor testDescriptor) {
3557
/*
36-
* We're cloning the descriptor to preserve its original state:
58+
* We're copying the descriptor to preserve its original state:
3759
* JUnit will modify some of its fields during and after test execution
3860
* (one example is parameterized test descriptor,
3961
* whose invocation context is overwritten with null).
40-
* Cloning has to be done before each test retry to
41-
* compensate for the state modifications.
62+
* The snapshot is taken before the first execution so that every retry
63+
* can be derived from the pristine state.
4264
*/
43-
this.testDescriptor = UnsafeUtils.tryShallowClone(testDescriptor);
65+
this.testDescriptor = copy(testDescriptor, UnaryOperator.identity());
4466
}
4567

4668
public TestDescriptor withIdSuffix(Map<String, Object> suffices) {
47-
UniqueId updatedId = testDescriptor.getUniqueId();
48-
for (Map.Entry<String, Object> e : suffices.entrySet()) {
49-
updatedId = updatedId.append(e.getKey(), String.valueOf(e.getValue()));
69+
return copy(
70+
testDescriptor,
71+
id -> {
72+
UniqueId updatedId = id;
73+
for (Map.Entry<String, Object> e : suffices.entrySet()) {
74+
updatedId = updatedId.append(e.getKey(), String.valueOf(e.getValue()));
75+
}
76+
return updatedId;
77+
});
78+
}
79+
80+
private static TestDescriptor copy(
81+
TestDescriptor testDescriptor, UnaryOperator<UniqueId> idTransform) {
82+
if (COPY_INCLUDING_DESCENDANTS != null
83+
&& JUPITER_TEST_DESCRIPTOR_CLASS != null
84+
&& JUPITER_TEST_DESCRIPTOR_CLASS.isInstance(testDescriptor)) {
85+
TestDescriptor copy =
86+
METHOD_HANDLES.invoke(COPY_INCLUDING_DESCENDANTS, testDescriptor, idTransform);
87+
if (copy != null) {
88+
// copyIncludingDescendants returns a detached copy so we link it back to its suite
89+
if (copy instanceof AbstractTestDescriptor) {
90+
((AbstractTestDescriptor) copy).setParent(testDescriptor.getParent().orElse(null));
91+
}
92+
return copy;
93+
}
5094
}
95+
return legacyCopy(testDescriptor, idTransform);
96+
}
5197

98+
/**
99+
* Fallback for engines without {@code copyIncludingDescendants}: shallow-clone the descriptor and
100+
* overwrite the cloned unique ID field by reflection. Not JEP 500 compliant.
101+
*/
102+
private static TestDescriptor legacyCopy(
103+
TestDescriptor testDescriptor, UnaryOperator<UniqueId> idTransform) {
52104
TestDescriptor descriptorClone = UnsafeUtils.tryShallowClone(testDescriptor);
53-
METHOD_HANDLES.invoke(UNIQUE_ID_SETTER, descriptorClone, updatedId);
105+
UniqueId updatedId = idTransform.apply(testDescriptor.getUniqueId());
106+
if (descriptorClone != testDescriptor && !updatedId.equals(testDescriptor.getUniqueId())) {
107+
METHOD_HANDLES.invoke(uniqueIdSetter(), descriptorClone, updatedId);
108+
}
54109
return descriptorClone;
55110
}
56111
}

‎dd-java-agent/instrumentation/junit/junit-5/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/execution/TestTaskHandle.java‎

Lines changed: 51 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import datadog.trace.agent.tooling.muzzle.Reference;
44
import datadog.trace.agent.tooling.muzzle.ReferenceProvider;
5+
import datadog.trace.instrumentation.junit5.JUnitPlatformUtils;
56
import datadog.trace.util.MethodHandles;
67
import java.lang.invoke.MethodHandle;
78
import java.util.Arrays;
@@ -13,7 +14,6 @@
1314
import org.junit.platform.engine.TestDescriptor;
1415
import org.junit.platform.engine.support.hierarchical.EngineExecutionContext;
1516
import org.junit.platform.engine.support.hierarchical.HierarchicalTestExecutorService;
16-
import org.junit.platform.engine.support.hierarchical.Node;
1717
import org.junit.platform.engine.support.hierarchical.ThrowableCollector;
1818

1919
public class TestTaskHandle {
@@ -28,17 +28,44 @@ public class TestTaskHandle {
2828

2929
private static final MethodHandle TEST_DESCRIPTOR_GETTER =
3030
METHOD_HANDLES.privateFieldGetter(TEST_TASK_CLASS, "testDescriptor");
31-
private static final MethodHandle TEST_DESCRIPTOR_SETTER =
32-
METHOD_HANDLES.privateFieldSetter(TEST_TASK_CLASS, "testDescriptor");
33-
34-
private static final MethodHandle NODE_SETTER =
35-
METHOD_HANDLES.privateFieldSetter(TEST_TASK_CLASS, "node");
3631

3732
private static final MethodHandle PARENT_CONTEXT_GETTER =
3833
METHOD_HANDLES.privateFieldGetter(TEST_TASK_CLASS, "parentContext");
3934
private static final MethodHandle PARENT_CONTEXT_SETTER =
4035
METHOD_HANDLES.privateFieldSetter(TEST_TASK_CLASS, "parentContext");
4136

37+
/** NodeTestTask's {@code (NodeTestTaskContext, TestDescriptor)} constructor (1.3.1+) */
38+
private static final Class<?> TEST_TASK_CONTEXT_CLASS_REF =
39+
JUnitPlatformUtils.loadClass(TEST_TASK_CONTEXT_CLASS);
40+
41+
private static final MethodHandle TEST_TASK_CONSTRUCTOR =
42+
TEST_TASK_CONTEXT_CLASS_REF != null
43+
? METHOD_HANDLES.constructor(
44+
TEST_TASK_CLASS, TEST_TASK_CONTEXT_CLASS_REF, TestDescriptor.class)
45+
: null;
46+
47+
// Legacy fallback setters, lazily created to avoid JEP 500 warnings, only used on 1.3.0
48+
private static volatile MethodHandle testDescriptorSetter;
49+
private static volatile MethodHandle nodeSetter;
50+
51+
private static MethodHandle testDescriptorSetter() {
52+
MethodHandle handle = testDescriptorSetter;
53+
if (handle == null) {
54+
handle = METHOD_HANDLES.privateFieldSetter(TEST_TASK_CLASS, "testDescriptor");
55+
testDescriptorSetter = handle;
56+
}
57+
return handle;
58+
}
59+
60+
private static MethodHandle nodeSetter() {
61+
MethodHandle handle = nodeSetter;
62+
if (handle == null) {
63+
handle = METHOD_HANDLES.privateFieldSetter(TEST_TASK_CLASS, "node");
64+
nodeSetter = handle;
65+
}
66+
return handle;
67+
}
68+
4269
private static final MethodHandle THROWABLE_COLLECTOR_FACTORY_GETTER =
4370
METHOD_HANDLES.privateFieldGetter(TEST_TASK_CLASS, "throwableCollectorFactory");
4471
private static final MethodHandle TASK_CONTEXT_THROWABLE_COLLECTOR_FACTORY_GETTER =
@@ -118,20 +145,29 @@ public TestDescriptor getTestDescriptor() {
118145
return METHOD_HANDLES.invoke(TEST_DESCRIPTOR_GETTER, testTask);
119146
}
120147

121-
public void setTestDescriptor(TestDescriptor testDescriptor) {
122-
METHOD_HANDLES.invoke(TEST_DESCRIPTOR_SETTER, testTask, testDescriptor);
123-
}
124-
125-
public void setNode(Node<?> node) {
126-
METHOD_HANDLES.invoke(NODE_SETTER, testTask, node);
127-
}
128-
129148
public EngineExecutionContext getParentContext() {
130149
return METHOD_HANDLES.invoke(PARENT_CONTEXT_GETTER, testTask);
131150
}
132151

133-
public void setParentContext(EngineExecutionContext parentContext) {
152+
/**
153+
* Returns a task that will execute the given retry descriptor. If possible, a brand-new
154+
* NodeTestTask is constructed; otherwise we fall back to overwriting the current task's fields
155+
* (non-compliant with JEP500).
156+
*/
157+
public HierarchicalTestExecutorService.TestTask createRetryTask(
158+
TestDescriptor descriptor, EngineExecutionContext parentContext) {
159+
if (TEST_TASK_CONSTRUCTOR != null && testTaskContext != null) {
160+
Object retryTask = METHOD_HANDLES.invoke(TEST_TASK_CONSTRUCTOR, testTaskContext, descriptor);
161+
if (retryTask != null) {
162+
METHOD_HANDLES.invoke(PARENT_CONTEXT_SETTER, retryTask, parentContext);
163+
return (HierarchicalTestExecutorService.TestTask) retryTask;
164+
}
165+
}
166+
// fallback (< 1.3.1): reuse the current task by overwriting its final fields.
167+
METHOD_HANDLES.invoke(testDescriptorSetter(), testTask, descriptor);
168+
METHOD_HANDLES.invoke(nodeSetter(), testTask, descriptor);
134169
METHOD_HANDLES.invoke(PARENT_CONTEXT_SETTER, testTask, parentContext);
170+
return testTask;
135171
}
136172

137173
public EngineExecutionListener getListener() {

0 commit comments

Comments
 (0)