1- package datadog .trace .civisibility .utils ;
1+ package datadog .trace .civisibility .domain ;
22
33import datadog .trace .api .civisibility .execution .TestStatus ;
44import datadog .trace .bootstrap .instrumentation .api .AgentSpan ;
1414import java .util .function .BinaryOperator ;
1515import java .util .function .Consumer ;
1616
17- public class SpanUtils {
18- public static final Consumer <AgentSpan > DO_NOT_PROPAGATE_CI_VISIBILITY_TAGS = span -> {};
17+ public class SpanTagsPropagator {
18+ public static final Consumer <AgentSpan > NOOP_PROPAGATOR = span -> {};
1919
20- public static Consumer <AgentSpan > propagateCiVisibilityTagsTo (AgentSpan parentSpan ) {
21- return childSpan -> propagateCiVisibilityTags (parentSpan , childSpan );
20+ private final AgentSpan parentSpan ;
21+ private final Object tagPropagationLock = new Object ();
22+
23+ public SpanTagsPropagator (AgentSpan parentSpan ) {
24+ this .parentSpan = parentSpan ;
25+ }
26+
27+ public void propagateCiVisibilityTags (AgentSpan childSpan ) {
28+ mergeTestFrameworks (getFrameworks (childSpan ));
29+ propagateStatus (childSpan );
30+ }
31+
32+ public void propagateStatus (AgentSpan childSpan ) {
33+ synchronized (tagPropagationLock ) {
34+ unsafePropagateStatus (childSpan );
35+ }
36+ }
37+
38+ public void mergeTestFrameworks (Collection <TestFramework > testFrameworks ) {
39+ synchronized (tagPropagationLock ) {
40+ unsafeMergeTestFrameworks (testFrameworks );
41+ }
2242 }
2343
24- public static void propagateCiVisibilityTags (AgentSpan parentSpan , AgentSpan childSpan ) {
25- mergeTestFrameworks (parentSpan , getFrameworks (childSpan ));
26- propagateStatus (parentSpan , childSpan );
44+ public void propagateTags (AgentSpan childSpan , TagMergeSpec <?>... specs ) {
45+ synchronized (tagPropagationLock ) {
46+ for (TagMergeSpec <?> spec : specs ) {
47+ unsafePropagateTag (childSpan , spec );
48+ }
49+ }
2750 }
2851
29- public static void mergeTestFrameworks ( AgentSpan span , Collection <TestFramework > testFrameworks ) {
30- Collection <TestFramework > spanFrameworks = getFrameworks (span );
31- Collection <TestFramework > merged = merge (spanFrameworks , testFrameworks );
32- setFrameworks (span , merged );
52+ private void unsafeMergeTestFrameworks ( Collection <TestFramework > childFrameworks ) {
53+ Collection <TestFramework > parentFrameworks = getFrameworks (parentSpan );
54+ Collection <TestFramework > merged = merge (parentFrameworks , childFrameworks );
55+ setFrameworks (merged );
3356 }
3457
35- private static Collection <TestFramework > getFrameworks (AgentSpan span ) {
58+ static Collection <TestFramework > getFrameworks (AgentSpan span ) {
3659 Object nameTag = span .getTag (Tags .TEST_FRAMEWORK );
3760 Object versionTag = span .getTag (Tags .TEST_FRAMEWORK_VERSION );
38- if (nameTag == null && versionTag == null ) {
61+ if (nameTag == null ) {
3962 return Collections .emptyList ();
4063 }
4164
@@ -45,9 +68,11 @@ private static Collection<TestFramework> getFrameworks(AgentSpan span) {
4568
4669 } else if (nameTag instanceof Collection ) {
4770 Iterator <String > names = ((Collection <String >) nameTag ).iterator ();
48- Iterator <String > versions = ((Collection <String >) versionTag ).iterator ();
71+ Iterator <String > versions =
72+ versionTag != null ? ((Collection <String >) versionTag ).iterator () : null ;
4973 while (names .hasNext ()) {
50- frameworks .add (new TestFramework (names .next (), versions .next ()));
74+ String version = (versions != null && versions .hasNext ()) ? versions .next () : null ;
75+ frameworks .add (new TestFramework (names .next (), version ));
5176 }
5277
5378 } else {
@@ -79,7 +104,7 @@ private static Collection<TestFramework> merge(
79104 return merged ;
80105 }
81106
82- private static void setFrameworks (AgentSpan span , Collection <TestFramework > frameworks ) {
107+ private void setFrameworks (Collection <TestFramework > frameworks ) {
83108 if (frameworks .isEmpty ()) {
84109 return ;
85110 }
@@ -90,7 +115,7 @@ private static void setFrameworks(AgentSpan span, Collection<TestFramework> fram
90115 if (framework .getVersion () != null ) {
91116 tags .put (Tags .TEST_FRAMEWORK_VERSION , framework .getVersion ());
92117 }
93- span .setAllTags (tags );
118+ parentSpan .setAllTags (tags );
94119 return ;
95120 }
96121 Collection <String > names = new ArrayList <>(frameworks .size ());
@@ -102,10 +127,10 @@ private static void setFrameworks(AgentSpan span, Collection<TestFramework> fram
102127 Map <String , Collection <String >> tags = new HashMap <>();
103128 tags .put (Tags .TEST_FRAMEWORK , names );
104129 tags .put (Tags .TEST_FRAMEWORK_VERSION , versions );
105- span .setAllTags (tags );
130+ parentSpan .setAllTags (tags );
106131 }
107132
108- private static void propagateStatus ( AgentSpan parentSpan , AgentSpan childSpan ) {
133+ private void unsafePropagateStatus ( AgentSpan childSpan ) {
109134 TestStatus childStatus = (TestStatus ) childSpan .getTag (Tags .TEST_STATUS );
110135 if (childStatus == null ) {
111136 return ;
@@ -131,25 +156,32 @@ private static void propagateStatus(AgentSpan parentSpan, AgentSpan childSpan) {
131156 }
132157 }
133158
134- public static void propagateTags (AgentSpan parentSpan , AgentSpan childSpan , String ... tagNames ) {
135- for (String tagName : tagNames ) {
136- parentSpan .setTag (tagName , childSpan .getTag (tagName ));
159+ public static class TagMergeSpec <T > {
160+ private final String tagKey ;
161+ private final BinaryOperator <T > mergeFunction ;
162+
163+ TagMergeSpec (String tagKey , BinaryOperator <T > mergeFunction ) {
164+ this .tagKey = tagKey ;
165+ this .mergeFunction = mergeFunction ;
166+ }
167+
168+ public static <T > TagMergeSpec <T > of (String key , BinaryOperator <T > mergeFunction ) {
169+ return new TagMergeSpec <>(key , mergeFunction );
137170 }
138- }
139171
140- public static <T > void propagateTag (AgentSpan parentSpan , AgentSpan childSpan , String tagName ) {
141- propagateTag (parentSpan , childSpan , tagName , (p , c ) -> c );
172+ public static TagMergeSpec <Object > of (String tagKey ) {
173+ return new TagMergeSpec <>(tagKey , (parent , child ) -> child );
174+ }
142175 }
143176
144- public static <T > void propagateTag (
145- AgentSpan parentSpan , AgentSpan childSpan , String tagName , BinaryOperator <T > mergeStrategy ) {
146- T childTag = (T ) childSpan .getTag (tagName );
177+ private <T > void unsafePropagateTag (AgentSpan childSpan , TagMergeSpec <T > spec ) {
178+ T childTag = (T ) childSpan .getTag (spec .tagKey );
147179 if (childTag != null ) {
148- T parentTag = (T ) parentSpan .getTag (tagName );
180+ T parentTag = (T ) parentSpan .getTag (spec . tagKey );
149181 if (parentTag == null ) {
150- parentSpan .setTag (tagName , childTag );
182+ parentSpan .setTag (spec . tagKey , childTag );
151183 } else {
152- parentSpan .setTag (tagName , mergeStrategy .apply (parentTag , childTag ));
184+ parentSpan .setTag (spec . tagKey , spec . mergeFunction .apply (parentTag , childTag ));
153185 }
154186 }
155187 }
0 commit comments