Save upstream span context in build action

jenkinsci · Mar 22, 2024 · 28f77b9 · 28f77b9
1 parent 42e0572
commit 28f77b9
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 14 deletions.
diff --git a/src/main/java/org/datadog/jenkins/plugins/datadog/listeners/DatadogBuildListener.java b/src/main/java/org/datadog/jenkins/plugins/datadog/listeners/DatadogBuildListener.java
@@ -117,7 +117,13 @@ public void onInitialize(Run run) {
             TraceSpan.TraceSpanContext buildSpanContext = new TraceSpan.TraceSpanContext();
             BuildSpanManager.get().put(buildData.getBuildTag(""), buildSpanContext);
 
-            final BuildSpanAction buildSpanAction = new BuildSpanAction(buildSpanContext);
+            TraceSpan.TraceSpanContext upstreamBuildSpanContext = null;
+            String upstreamBuildTag = buildData.getUpstreamBuildTag("");
+            if (upstreamBuildTag != null) {
+                upstreamBuildSpanContext = BuildSpanManager.get().get(upstreamBuildTag);
+            }
+
+            final BuildSpanAction buildSpanAction = new BuildSpanAction(buildSpanContext, upstreamBuildSpanContext);
             run.addAction(buildSpanAction);
 
             run.addAction(new GitCommitAction());

diff --git a/src/main/java/org/datadog/jenkins/plugins/datadog/model/BuildData.java b/src/main/java/org/datadog/jenkins/plugins/datadog/model/BuildData.java
@@ -93,6 +93,8 @@ public class BuildData implements Serializable {
     private String jobName;
     private Map<String, String> buildConfigurations;
     private String buildTag;
+    @Nullable
+    private String upstreamBuildTag;
     private String jenkinsUrl;
     private String executorNumber;
     private String javaHome;
@@ -295,8 +297,15 @@ private void populateUpstreamPipelineData(Run<?, ?> run, EnvVars envVars) {
 
         String upstreamProject = upstreamCause.getUpstreamProject();
         if (upstreamProject != null) {
-            String upstreamPipelineBuildTag = "jenkins-" + upstreamProject.replace('/', '-') + "-" + upstreamBuild;
-            TraceSpan.TraceSpanContext upstreamPipelineContext = BuildSpanManager.get().get(upstreamPipelineBuildTag);
+            upstreamBuildTag = "jenkins-" + upstreamProject.replace('/', '-') + "-" + upstreamBuild;
+            TraceSpan.TraceSpanContext upstreamPipelineContext = BuildSpanManager.get().get(upstreamBuildTag);
+            if (upstreamPipelineContext == null) {
+                BuildSpanAction buildSpanAction = run.getAction(BuildSpanAction.class);
+                if (buildSpanAction != null) {
+                    upstreamPipelineContext = buildSpanAction.getUpstreamSpanContext();
+                }
+            }
+
             if (upstreamPipelineContext != null) {
                 upstreamPipelineTraceId = upstreamPipelineContext.getTraceId();
             }
@@ -705,6 +714,11 @@ public String getBuildTag(String value) {
         return defaultIfNull(buildTag, value);
     }
 
+    @Nullable
+    public String getUpstreamBuildTag(String value) {
+        return defaultIfNull(upstreamBuildTag, value);
+    }
+
     public String getJenkinsUrl(String value) {
         return defaultIfNull(jenkinsUrl, value);
     }

diff --git a/src/main/java/org/datadog/jenkins/plugins/datadog/traces/BuildSpanAction.java b/src/main/java/org/datadog/jenkins/plugins/datadog/traces/BuildSpanAction.java
@@ -7,6 +7,7 @@
 import com.thoughtworks.xstream.io.HierarchicalStreamWriter;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
+import javax.annotation.Nullable;
 import org.datadog.jenkins.plugins.datadog.model.DatadogPluginAction;
 import org.datadog.jenkins.plugins.datadog.traces.message.TraceSpan;
 import org.datadog.jenkins.plugins.datadog.util.conversion.DatadogActionConverter;
@@ -20,16 +21,19 @@ public class BuildSpanAction extends DatadogPluginAction {
     private static final long serialVersionUID = 1L;
 
     private final TraceSpan.TraceSpanContext buildSpanContext;
+    private final TraceSpan.TraceSpanContext upstreamSpanContext;
     private final AtomicInteger version;
     private volatile String buildUrl;
 
-    public BuildSpanAction(final TraceSpan.TraceSpanContext buildSpanContext){
-       this.buildSpanContext = buildSpanContext;
-       this.version = new AtomicInteger(0);
+    public BuildSpanAction(final TraceSpan.TraceSpanContext buildSpanContext, @Nullable final TraceSpan.TraceSpanContext upstreamSpanContext) {
+        this.buildSpanContext = buildSpanContext;
+        this.upstreamSpanContext = upstreamSpanContext;
+        this.version = new AtomicInteger(0);
     }
 
-    public BuildSpanAction(TraceSpan.TraceSpanContext buildSpanContext, int version, String buildUrl) {
+    public BuildSpanAction(TraceSpan.TraceSpanContext buildSpanContext, TraceSpan.TraceSpanContext upstreamSpanContext, int version, String buildUrl) {
         this.buildSpanContext = buildSpanContext;
+        this.upstreamSpanContext = upstreamSpanContext;
         this.version = new AtomicInteger(version);
         this.buildUrl = buildUrl;
     }
@@ -38,6 +42,10 @@ public TraceSpan.TraceSpanContext getBuildSpanContext() {
         return buildSpanContext;
     }
 
+    public TraceSpan.TraceSpanContext getUpstreamSpanContext() {
+        return upstreamSpanContext;
+    }
+
     public String getBuildUrl() {
         return buildUrl;
     }
@@ -74,7 +82,7 @@ public String toString() {
 
     public static final class ConverterImpl extends DatadogActionConverter<BuildSpanAction> {
         public ConverterImpl(XStream xs) {
-            super(new ConverterV1());
+            super(new ConverterV1(), new ConverterV2());
         }
     }
 
@@ -110,7 +118,50 @@ public BuildSpanAction unmarshal(HierarchicalStreamReader reader, UnmarshallingC
                 reader.moveUp();
             }
 
-            return new BuildSpanAction(spanContext, version, buildUrl);
+            return new BuildSpanAction(spanContext, null, version, buildUrl);
+        }
+    }
+
+    public static final class ConverterV2 extends VersionedConverter<BuildSpanAction> {
+
+        private static final int VERSION = 2;
+
+        public ConverterV2() {
+            super(VERSION);
+        }
+
+        @Override
+        public void marshal(BuildSpanAction action, HierarchicalStreamWriter writer, MarshallingContext context) {
+            writeField("version", action.version.get(), writer, context);
+            writeField("spanContext", action.buildSpanContext, writer, context);
+            if (action.upstreamSpanContext != null) {
+                writeField("upstreamSpanContext", action.upstreamSpanContext, writer, context);
+            }
+            if (action.buildUrl != null) {
+                writeField("buildUrl", action.buildUrl, writer, context);
+            }
+        }
+
+        @Override
+        public BuildSpanAction unmarshal(HierarchicalStreamReader reader, UnmarshallingContext context) {
+            int version = readField(reader, context, int.class);
+            TraceSpan.TraceSpanContext spanContext = readField(reader, context, TraceSpan.TraceSpanContext.class);
+
+            String buildUrl = null;
+            TraceSpan.TraceSpanContext upstreamSpanContext = null;
+            while (reader.hasMoreChildren()) {
+                reader.moveDown();
+                String fieldName = reader.getNodeName();
+                if ("buildUrl".equals(fieldName)) {
+                    buildUrl = (String) context.convertAnother(null, String.class);
+                }
+                if ("upstreamSpanContext".equals(fieldName)) {
+                    upstreamSpanContext= readField(reader, context, TraceSpan.TraceSpanContext.class);
+                }
+                reader.moveUp();
+            }
+
+            return new BuildSpanAction(spanContext, upstreamSpanContext, version, buildUrl);
         }
     }
 }
diff --git a/src/test/java/org/datadog/jenkins/plugins/datadog/model/BuildSpanActionConverterTest.java b/src/test/java/org/datadog/jenkins/plugins/datadog/model/BuildSpanActionConverterTest.java
@@ -15,11 +15,11 @@ public class BuildSpanActionConverterTest extends ActionConverterTest<BuildSpanA
     @Parameterized.Parameters(name = "{0}")
     public static Collection<Object[]> data() {
         return Arrays.asList(new Object[][]{
-                {new BuildSpanAction(new TraceSpan.TraceSpanContext(123, 456, 789), 0, "buildUrl")},
-                {new BuildSpanAction(new TraceSpan.TraceSpanContext(0, 456, 789), 1, "buildUrl")},
-                {new BuildSpanAction(new TraceSpan.TraceSpanContext(123, 0, 789), 2, null)},
-                {new BuildSpanAction(new TraceSpan.TraceSpanContext(123, 456, 0), 3, "buildUrl")},
-                {new BuildSpanAction(new TraceSpan.TraceSpanContext(0, 0, 0), 4, null)},
+                {new BuildSpanAction(new TraceSpan.TraceSpanContext(123, 456, 789), null, 0, "buildUrl")},
+                {new BuildSpanAction(new TraceSpan.TraceSpanContext(0, 456, 789), null, 1, "buildUrl")},
+                {new BuildSpanAction(new TraceSpan.TraceSpanContext(123, 0, 789), null, 2, null)},
+                {new BuildSpanAction(new TraceSpan.TraceSpanContext(123, 456, 0), new TraceSpan.TraceSpanContext(222, 333, 444), 3, "buildUrl")},
+                {new BuildSpanAction(new TraceSpan.TraceSpanContext(0, 0, 0), new TraceSpan.TraceSpanContext(0, 0, 0), 4, null)},
         });
     }