Skip to content

Commit 6fec4d9

Browse files
daniel-mohedanodaniel-mohedano
committed
feat: propagation of custom tags to all test event levels
1 parent a5ac864 commit 6fec4d9

15 files changed

Lines changed: 194 additions & 26 deletions

File tree

‎dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/AbstractTestModule.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ public AbstractTestModule(
6464
}
6565

6666
span = spanBuilder.start();
67-
tagsPropagator = new SpanTagsPropagator(span);
67+
tagsPropagator = new SpanTagsPropagator(span, config.getCiVisibilityPropagatedTagKeys());
6868

6969
span.setSpanType(InternalSpanTypes.TEST_MODULE_END);
7070
span.setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_TEST_MODULE);

‎dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/AbstractTestSession.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ public AbstractTestSession(
9999
}
100100

101101
span = spanBuilder.start();
102-
tagPropagator = new SpanTagsPropagator(span);
102+
tagPropagator = new SpanTagsPropagator(span, config.getCiVisibilityPropagatedTagKeys());
103103

104104
span.setSpanType(InternalSpanTypes.TEST_SESSION_END);
105105
span.setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_TEST_SESSION);

‎dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/SpanTagsPropagator.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,23 @@ public class SpanTagsPropagator {
1818
public static final Consumer<AgentSpan> NOOP_PROPAGATOR = span -> {};
1919

2020
private final AgentSpan parentSpan;
21+
private final Collection<String> propagatedTagKeys;
2122
private final Object tagPropagationLock = new Object();
2223

2324
public SpanTagsPropagator(AgentSpan parentSpan) {
25+
this(parentSpan, Collections.emptyList());
26+
}
27+
28+
public SpanTagsPropagator(AgentSpan parentSpan, Collection<String> propagatedTagKeys) {
2429
this.parentSpan = parentSpan;
30+
this.propagatedTagKeys =
31+
propagatedTagKeys != null ? propagatedTagKeys : Collections.emptyList();
2532
}
2633

2734
public void propagateCiVisibilityTags(AgentSpan childSpan) {
2835
mergeTestFrameworks(getFrameworks(childSpan));
2936
propagateStatus(childSpan);
37+
propagateCustomTags(childSpan);
3038
}
3139

3240
public void propagateStatus(AgentSpan childSpan) {
@@ -49,6 +57,34 @@ public void propagateTags(AgentSpan childSpan, TagMergeSpec<?>... specs) {
4957
}
5058
}
5159

60+
public void propagateCustomTags(AgentSpan childSpan) {
61+
if (propagatedTagKeys.isEmpty()) {
62+
return;
63+
}
64+
synchronized (tagPropagationLock) {
65+
for (String key : propagatedTagKeys) {
66+
Object value = childSpan.getTag(key);
67+
if (value != null) {
68+
parentSpan.setTag(key, String.valueOf(value));
69+
}
70+
}
71+
}
72+
}
73+
74+
public void propagateCustomTags(Map<String, String> tags) {
75+
if (propagatedTagKeys.isEmpty() || tags == null || tags.isEmpty()) {
76+
return;
77+
}
78+
synchronized (tagPropagationLock) {
79+
for (String key : propagatedTagKeys) {
80+
String value = tags.get(key);
81+
if (value != null) {
82+
parentSpan.setTag(key, value);
83+
}
84+
}
85+
}
86+
}
87+
5288
private void unsafeMergeTestFrameworks(Collection<TestFramework> childFrameworks) {
5389
Collection<TestFramework> parentFrameworks = getFrameworks(parentSpan);
5490
Collection<TestFramework> merged = merge(parentFrameworks, childFrameworks);

‎dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestSuiteImpl.java‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ public TestSuiteImpl(
112112
}
113113

114114
span = spanBuilder.start();
115-
tagsPropagator = new SpanTagsPropagator(span);
115+
tagsPropagator = new SpanTagsPropagator(span, config.getCiVisibilityPropagatedTagKeys());
116116

117117
span.setSpanType(InternalSpanTypes.TEST_SUITE_END);
118118
span.setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_TEST_SUITE);
@@ -275,6 +275,11 @@ public TestImpl testStart(
275275
executionResults,
276276
configurationErrors,
277277
capabilities,
278-
tagsPropagator::propagateStatus);
278+
this::propagateTags);
279+
}
280+
281+
private void propagateTags(AgentSpan childSpan) {
282+
tagsPropagator.propagateStatus(childSpan);
283+
tagsPropagator.propagateCustomTags(childSpan);
279284
}
280285
}

‎dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/buildsystem/BuildSystemModuleImpl.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -290,6 +290,7 @@ private SignalResponse onModuleExecutionResultReceived(ModuleExecutionResult res
290290
testsSkipped.add(result.getTestsSkippedTotal());
291291

292292
tagsPropagator.mergeTestFrameworks(result.getTestFrameworks());
293+
tagsPropagator.propagateCustomTags(result.getPropagatedTags());
293294

294295
return AckResponse.INSTANCE;
295296
}

‎dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/buildsystem/ProxyTestModule.java‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
import datadog.trace.civisibility.test.ExecutionResults;
3232
import datadog.trace.civisibility.test.ExecutionStrategy;
3333
import java.util.Collection;
34+
import java.util.Map;
35+
import java.util.Set;
3436
import java.util.TreeSet;
3537
import java.util.concurrent.ConcurrentHashMap;
3638
import javax.annotation.Nonnull;
@@ -61,6 +63,8 @@ public class ProxyTestModule implements TestFrameworkModule {
6163
private final LinesResolver linesResolver;
6264
private final CoverageStore.Factory coverageStoreFactory;
6365
private final Collection<TestFramework> testFrameworks = ConcurrentHashMap.newKeySet();
66+
private final Map<String, String> propagatedTags = new ConcurrentHashMap<>();
67+
private final Set<String> propagatedTagKeys;
6468
private final Collection<LibraryCapability> capabilities;
6569

6670
public ProxyTestModule(
@@ -91,6 +95,7 @@ public ProxyTestModule(
9195
this.linesResolver = linesResolver;
9296
this.coverageStoreFactory = coverageStoreFactory;
9397
this.capabilities = capabilities;
98+
this.propagatedTagKeys = config.getCiVisibilityPropagatedTagKeys();
9499
}
95100

96101
@Override
@@ -180,7 +185,8 @@ private void sendModuleExecutionResult() {
180185
testManagementEnabled,
181186
hasFailedTestReplayTests,
182187
testsSkippedTotal,
183-
new TreeSet<>(testFrameworks)));
188+
new TreeSet<>(testFrameworks),
189+
propagatedTags));
184190

185191
} catch (Exception e) {
186192
log.error("Error while reporting module execution result", e);
@@ -215,13 +221,24 @@ public TestSuiteImpl testSuiteStart(
215221
executionResults,
216222
executionStrategy.getExecutionSettings().getConfigurationErrors(),
217223
capabilities,
218-
this::propagateTestFrameworkData);
224+
this::propagateData);
219225
}
220226

221-
private void propagateTestFrameworkData(AgentSpan childSpan) {
227+
private void propagateData(AgentSpan childSpan) {
222228
testFrameworks.add(
223229
new TestFramework(
224230
(String) childSpan.getTag(Tags.TEST_FRAMEWORK),
225231
(String) childSpan.getTag(Tags.TEST_FRAMEWORK_VERSION)));
232+
233+
if (propagatedTagKeys.isEmpty()) {
234+
return;
235+
}
236+
237+
for (String key : propagatedTagKeys) {
238+
Object value = childSpan.getTag(key);
239+
if (value != null) {
240+
propagatedTags.put(key, String.valueOf(value));
241+
}
242+
}
226243
}
227244
}

‎dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/ipc/ModuleExecutionResult.java‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
import datadog.trace.civisibility.ipc.serialization.Serializer;
55
import java.nio.ByteBuffer;
66
import java.util.Collection;
7+
import java.util.Collections;
8+
import java.util.Map;
79
import java.util.Objects;
810

911
public class ModuleExecutionResult extends ModuleSignal {
@@ -23,6 +25,7 @@ public class ModuleExecutionResult extends ModuleSignal {
2325
private final boolean hasFailedTestReplayTests;
2426
private final long testsSkippedTotal;
2527
private final Collection<TestFramework> testFrameworks;
28+
private final Map<String, String> propagatedTags;
2629

2730
public ModuleExecutionResult(
2831
DDTraceId sessionId,
@@ -34,7 +37,8 @@ public ModuleExecutionResult(
3437
boolean testManagementEnabled,
3538
boolean hasFailedTestReplayTests,
3639
long testsSkippedTotal,
37-
Collection<TestFramework> testFrameworks) {
40+
Collection<TestFramework> testFrameworks,
41+
Map<String, String> propagatedTags) {
3842
super(sessionId, moduleId);
3943
this.coverageEnabled = coverageEnabled;
4044
this.testSkippingEnabled = testSkippingEnabled;
@@ -44,6 +48,7 @@ public ModuleExecutionResult(
4448
this.hasFailedTestReplayTests = hasFailedTestReplayTests;
4549
this.testsSkippedTotal = testsSkippedTotal;
4650
this.testFrameworks = testFrameworks;
51+
this.propagatedTags = propagatedTags != null ? propagatedTags : Collections.emptyMap();
4752
}
4853

4954
public boolean isCoverageEnabled() {
@@ -78,6 +83,10 @@ public Collection<TestFramework> getTestFrameworks() {
7883
return testFrameworks;
7984
}
8085

86+
public Map<String, String> getPropagatedTags() {
87+
return propagatedTags;
88+
}
89+
8190
@Override
8291
public boolean equals(Object o) {
8392
if (this == o) {
@@ -94,7 +103,8 @@ public boolean equals(Object o) {
94103
&& testSkippingEnabled == that.testSkippingEnabled
95104
&& hasFailedTestReplayTests == that.hasFailedTestReplayTests
96105
&& testsSkippedTotal == that.testsSkippedTotal
97-
&& Objects.equals(testFrameworks, that.testFrameworks);
106+
&& Objects.equals(testFrameworks, that.testFrameworks)
107+
&& Objects.equals(propagatedTags, that.propagatedTags);
98108
}
99109

100110
@Override
@@ -106,7 +116,8 @@ public int hashCode() {
106116
testSkippingEnabled,
107117
hasFailedTestReplayTests,
108118
testsSkippedTotal,
109-
testFrameworks);
119+
testFrameworks,
120+
propagatedTags);
110121
}
111122

112123
@Override
@@ -161,6 +172,7 @@ public ByteBuffer serialize() {
161172

162173
s.write(testsSkippedTotal);
163174
s.write(testFrameworks, TestFramework::serialize);
175+
s.write(propagatedTags);
164176

165177
return s.flush();
166178
}
@@ -180,6 +192,7 @@ public static ModuleExecutionResult deserialize(ByteBuffer buffer) {
180192
long testsSkippedTotal = Serializer.readLong(buffer);
181193
Collection<TestFramework> testFrameworks =
182194
Serializer.readList(buffer, TestFramework::deserialize);
195+
Map<String, String> propagatedTags = Serializer.readStringMap(buffer);
183196

184197
return new ModuleExecutionResult(
185198
sessionId,
@@ -191,6 +204,7 @@ public static ModuleExecutionResult deserialize(ByteBuffer buffer) {
191204
testManagementEnabled,
192205
hasFailedTestReplayTests,
193206
testsSkippedTotal,
194-
testFrameworks);
207+
testFrameworks,
208+
propagatedTags);
195209
}
196210
}

‎dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/domain/SpanTagsPropagatorTest.groovy‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package datadog.trace.civisibility.domain
33
import static datadog.trace.civisibility.domain.SpanTagsPropagator.TagMergeSpec
44

55
import datadog.trace.api.civisibility.execution.TestStatus
6+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
67
import datadog.trace.bootstrap.instrumentation.api.Tags
78
import datadog.trace.civisibility.ipc.TestFramework
89
import datadog.trace.core.DDSpan
@@ -126,6 +127,61 @@ class SpanTagsPropagatorTest extends Specification {
126127
TagMergeSpec.of("tag", Boolean::logicalOr) | false | false | true | false
127128
}
128129

130+
// Mocks AgentSpan (interface) rather than DDSpan because propagateCustomTags writes through
131+
// the final DDSpan#setTag(String, String) overload, which Spock cannot intercept on a class mock.
132+
def "test custom tag propagation from span: child=#childValue, parent=#parentValue, key=#key, allowlist=#allowlist"() {
133+
given:
134+
def parentSpan = Mock(AgentSpan)
135+
parentSpan.getTag(key) >> parentValue
136+
137+
def childSpan = Mock(AgentSpan)
138+
childSpan.getTag(key) >> childValue
139+
140+
def propagator = new SpanTagsPropagator(parentSpan, allowlist)
141+
142+
when:
143+
propagator.propagateCustomTags(childSpan)
144+
145+
then:
146+
if (expectedValue != null) {
147+
1 * parentSpan.setTag(key, expectedValue)
148+
} else {
149+
0 * parentSpan.setTag(key, _)
150+
}
151+
152+
where:
153+
allowlist | key | childValue | parentValue | expectedValue
154+
["bazel.shard_index"] | "bazel.shard_index" | "0" | null | "0"
155+
["bazel.shard_index"] | "bazel.shard_index" | "1" | "0" | "1" // child overrides parent
156+
["bazel.shard_index"] | "bazel.shard_index" | null | "0" | null // missing on child, no-op
157+
["bazel.shard_index"] | "bazel.total_shards" | "2" | null | null // not in allowlist
158+
[] | "bazel.shard_index" | "0" | null | null // empty allowlist
159+
null | "bazel.shard_index" | "0" | null | null // null allowlist
160+
["bazel.shard_index"] | "bazel.shard_index" | 0L | null | "0" // non-string child stringified
161+
["bazel.shard_index"] | "bazel.shard_index" | true | null | "true" // boolean stringified
162+
}
163+
164+
def "test custom tag propagation from map: allowlist=#allowlist, tags=#tags"() {
165+
given:
166+
def parentSpan = Mock(AgentSpan)
167+
def propagator = new SpanTagsPropagator(parentSpan, allowlist)
168+
169+
when:
170+
propagator.propagateCustomTags(tags)
171+
172+
then:
173+
expectedSets * parentSpan.setTag(_, _)
174+
175+
where:
176+
allowlist | tags | expectedSets
177+
["bazel.shard_index", "bazel.total_shards"] | ["bazel.shard_index": "0", "bazel.total_shards": "2"] | 2
178+
["bazel.shard_index"] | ["bazel.shard_index": "0"] | 1
179+
["bazel.shard_index"] | ["bazel.shard_index": "0", "bazel.total_shards": "2"] | 1 // only allowlisted keys are copied
180+
["bazel.shard_index"] | [:] | 0 // empty tags
181+
[] | ["bazel.shard_index": "0"] | 0 // empty allowlist
182+
null | ["bazel.shard_index": "0"] | 0 // null allowlist
183+
}
184+
129185
def "test synchronized propagation"() {
130186
given:
131187
def parentSpan = Mock(DDSpan)

‎dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/domain/headless/HeadlessTestSessionTest.groovy‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ class HeadlessTestSessionTest extends SpanWriterTest {
2424
def module = session.testModuleStart("module-name", null)
2525

2626
when:
27+
module.setTag("custom.propagated_tag", "value")
2728
module.end(null)
2829
session.end(null)
2930

@@ -34,10 +35,14 @@ class HeadlessTestSessionTest extends SpanWriterTest {
3435
spanType DDSpanTypes.TEST_SESSION_END
3536
tags(false) {
3637
"$Tags.TEST_TEST_MANAGEMENT_ENABLED" true
38+
"custom.propagated_tag" "value"
3739
}
3840
}
3941
span(1) {
4042
spanType DDSpanTypes.TEST_MODULE_END
43+
tags(false) {
44+
"custom.propagated_tag" "value"
45+
}
4146
}
4247
}
4348
})
@@ -47,13 +52,16 @@ class HeadlessTestSessionTest extends SpanWriterTest {
4752
def executionSettings = Stub(ExecutionSettings)
4853
executionSettings.getTestManagementSettings() >> new TestManagementSettings(true, 10)
4954

50-
def executionStrategy = new ExecutionStrategy(Stub(Config), executionSettings, Stub(SourcePathResolver), Stub(LinesResolver))
55+
def config = Stub(Config)
56+
config.getCiVisibilityPropagatedTagKeys() >> ["custom.propagated_tag"]
57+
58+
def executionStrategy = new ExecutionStrategy(config, executionSettings, Stub(SourcePathResolver), Stub(LinesResolver))
5159

5260
new HeadlessTestSession(
5361
"project-name",
5462
null,
5563
Provider.UNSUPPORTED,
56-
Stub(Config),
64+
config,
5765
Stub(CiVisibilityMetricCollector),
5866
Stub(TestDecorator),
5967
Stub(SourcePathResolver),

0 commit comments

Comments
 (0)