Propagate trace/span IDs via environment variables for Jenkins Pipelines

src/main/java/org/datadog/jenkins/plugins/datadog/DatadogUtilities.java

@@ -47,6 +47,7 @@ of this software and associated documentation files (the "Software"), to deal
 import org.datadog.jenkins.plugins.datadog.traces.BuildSpanAction;
 import org.datadog.jenkins.plugins.datadog.traces.IsPipelineAction;
 import org.datadog.jenkins.plugins.datadog.traces.StepDataAction;
+import org.datadog.jenkins.plugins.datadog.traces.StepTraceDataAction;
 import org.datadog.jenkins.plugins.datadog.util.SuppressFBWarnings;
 import org.datadog.jenkins.plugins.datadog.util.TagsUtil;
 import org.jenkinsci.plugins.pipeline.StageStatus;
@@ -878,6 +879,7 @@ public static void cleanUpTraceActions(final Run<?, ?> run) {
             run.removeActions(PipelineQueueInfoAction.class);
             run.removeActions(StageBreakdownAction.class);
             run.removeActions(IsPipelineAction.class);
+            run.removeActions(StepTraceDataAction.class);
         }
     }
 

src/main/java/org/datadog/jenkins/plugins/datadog/listeners/DatadogBuildListener.java

@@ -27,10 +27,16 @@ of this software and associated documentation files (the "Software"), to deal
 
 import static org.datadog.jenkins.plugins.datadog.DatadogUtilities.cleanUpTraceActions;
 import static org.datadog.jenkins.plugins.datadog.DatadogUtilities.isPipeline;
+import static org.datadog.jenkins.plugins.datadog.traces.TracerConstants.SPAN_ID_ENVVAR_KEY;
+import static org.datadog.jenkins.plugins.datadog.traces.TracerConstants.TRACE_ID_ENVVAR_KEY;
 
 import com.cloudbees.workflow.rest.external.RunExt;
 import com.cloudbees.workflow.rest.external.StageNodeExt;
 import hudson.Extension;
+import hudson.Launcher;
+import hudson.model.AbstractBuild;
+import hudson.model.BuildListener;
+import hudson.model.Environment;
 import hudson.model.Queue;
 import hudson.model.Result;
 import hudson.model.Run;
@@ -44,15 +50,9 @@ of this software and associated documentation files (the "Software"), to deal
 import org.datadog.jenkins.plugins.datadog.events.BuildFinishedEventImpl;
 import org.datadog.jenkins.plugins.datadog.events.BuildStartedEventImpl;
 import org.datadog.jenkins.plugins.datadog.model.BuildData;
-import org.datadog.jenkins.plugins.datadog.model.CIGlobalTagsAction;
-import org.datadog.jenkins.plugins.datadog.model.GitCommitAction;
-import org.datadog.jenkins.plugins.datadog.model.GitRepositoryAction;
-import org.datadog.jenkins.plugins.datadog.model.PipelineNodeInfoAction;
-import org.datadog.jenkins.plugins.datadog.model.PipelineQueueInfoAction;
-import org.datadog.jenkins.plugins.datadog.model.StageBreakdownAction;
 import org.datadog.jenkins.plugins.datadog.traces.BuildSpanAction;
-import org.datadog.jenkins.plugins.datadog.traces.IsPipelineAction;
-import org.datadog.jenkins.plugins.datadog.traces.StepDataAction;
+
+import org.datadog.jenkins.plugins.datadog.traces.message.TraceSpan;
 import org.jenkinsci.plugins.workflow.job.WorkflowRun;
 
 import javax.annotation.Nonnull;
@@ -112,6 +112,42 @@ public void onInitialize(Run run) {
         }
     }
 
+    /**
+     * Called before the SCMCheckout is run in a Jenkins build.
+     * This method is called after onInitialize callback.
+     */
+    @Override
+    public Environment setUpEnvironment(AbstractBuild build, Launcher launcher, BuildListener listener) throws Run.RunnerAbortedException {
+        try {
+            logger.fine("Start DatadogBuildListener#setUpEnvironment");
+
+            final BuildSpanAction buildSpanAction = build.getAction(BuildSpanAction.class);
+            if(buildSpanAction == null || buildSpanAction.getBuildData() == null) {
+                return new Environment() {
+                };
+            }
+
+            final TraceSpan.TraceSpanContext traceSpanContext = buildSpanAction.getBuildSpanContext();
+            final long traceId = traceSpanContext.getTraceId();
+            final long spanId = traceSpanContext.getSpanId();
+
+            final Environment newEnv = new Environment() {
+                @Override
+                public void buildEnvVars(Map<String, String> env) {
+                    env.put(TRACE_ID_ENVVAR_KEY, Long.toString(traceId));
+                    env.put(SPAN_ID_ENVVAR_KEY, Long.toString(spanId));
+                }
+            };
+
+            logger.fine("End DatadogBuildListener#setUpEnvironment");
+            return newEnv;
+        } catch (Exception e) {
+            DatadogUtilities.severe(logger, e, null);
+            return new Environment() {
+            };
+        }
+    }
+
     /**
      * Called when a build is first started.
      *
@@ -147,7 +183,7 @@ public void onStarted(Run run, TaskListener listener) {
             DatadogEvent event = new BuildStartedEventImpl(buildData);
             client.event(event);
 
-            // Send an metric
+            // Send a metric
             // item.getInQueueSince() may raise a NPE if a worker node is spinning up to run the job.
             // This could be expected behavior with ec2 spot instances/ecs containers, meaning no waiting
             // queue times if the plugin is spinning up an instance/container for one/first job.

src/main/java/org/datadog/jenkins/plugins/datadog/model/BuildPipelineNode.java

@@ -4,6 +4,7 @@
 import hudson.model.Run;
 import org.datadog.jenkins.plugins.datadog.DatadogUtilities;
 import org.datadog.jenkins.plugins.datadog.traces.StepDataAction;
+import org.datadog.jenkins.plugins.datadog.traces.StepTraceDataAction;
 import org.datadog.jenkins.plugins.datadog.util.SuppressFBWarnings;
 import org.jenkinsci.plugins.workflow.actions.ArgumentsAction;
 import org.jenkinsci.plugins.workflow.actions.ErrorAction;
@@ -88,6 +89,9 @@ public String getBuildLevel() {
     // Although the error flag was true, this can be null.
     private Throwable errorObj;
 
+    //Tracing
+    private long spanId = -1L;
+
     public BuildPipelineNode(final String id, final String name) {
         this(new BuildPipelineNodeKey(id, name));
     }
@@ -174,6 +178,11 @@ public BuildPipelineNode(final StepAtomNode stepNode) {
             this.nodeLabels = stepData.getNodeLabels();
         }
 
+        final StepTraceData stepTraceData = getStepTraceData(stepNode);
+        if(stepTraceData != null) {
+            this.spanId = stepTraceData.getSpanId();
+        }
+
         final FlowNodeQueueData queueData = getQueueData(stepNode);
         if(queueData != null) {
             this.nanosInQueue = queueData.getNanosInQueue();
@@ -320,6 +329,10 @@ public void setError(boolean error) {
         this.error = error;
     }
 
+    public long getSpanId() {
+        return spanId;
+    }
+
     public List<BuildPipelineNode> getParents(){ return parents; }
 
     public List<BuildPipelineNode> getChildren() {
@@ -365,6 +378,7 @@ public void updateData(final BuildPipelineNode buildNode) {
         this.error = buildNode.error;
         this.errorObj = buildNode.errorObj;
         this.parents.addAll(buildNode.parents);
+        this.spanId = buildNode.spanId;
     }
 
     public void addChild(final BuildPipelineNode child) {
@@ -440,6 +454,23 @@ private StepData getStepData(final FlowNode flowNode) {
         return stepDataAction.synchronizedGet(run, flowNode);
     }
 
+    @SuppressFBWarnings("NP_NULL_ON_SOME_PATH_FROM_RETURN_VALUE")
+    private StepTraceData getStepTraceData(FlowNode flowNode) {
+        final Run<?, ?> run = getRun(flowNode);
+        if(run == null) {
+            logger.fine("Unable to get StepTraceData from flowNode '"+flowNode.getDisplayName()+"'. Run is null");
+            return null;
+        }
+
+        final StepTraceDataAction stepTraceDataAction = run.getAction(StepTraceDataAction.class);
+        if(stepTraceDataAction == null) {
+            logger.fine("Unable to get StepTraceData from flowNode '"+flowNode.getDisplayName()+"'. StepTraceDataAction is null");
+            return null;
+        }
+
+        return stepTraceDataAction.synchronizedGet(run, flowNode);
+    }
+
     private FlowNodeQueueData getQueueData(FlowNode node) {
         final Run<?, ?> run = getRun(node);
         if(run == null) {

src/main/java/org/datadog/jenkins/plugins/datadog/model/StepData.java

@@ -5,6 +5,7 @@
 import hudson.model.Computer;
 import org.datadog.jenkins.plugins.datadog.DatadogUtilities;
 import org.datadog.jenkins.plugins.datadog.audit.DatadogAudit;
+import org.datadog.jenkins.plugins.datadog.traces.IdGenerator;
 import org.jenkinsci.plugins.workflow.steps.StepContext;
 
 import java.io.Serializable;
@@ -60,7 +61,6 @@ public Set<String> getNodeLabels() {
         return nodeLabels;
     }
 
-
     /**
      * Returns the workspace filepath of the remote node which is executing a determined {@code Step}
      * @param stepContext

src/main/java/org/datadog/jenkins/plugins/datadog/model/StepTraceData.java

@@ -0,0 +1,16 @@
+package org.datadog.jenkins.plugins.datadog.model;
+
+import java.io.Serializable;
+
+public class StepTraceData implements Serializable {
+
+    private final long spanId;
+
+    public StepTraceData(final long spanId) {
+        this.spanId = spanId;
+    }
+
+    public long getSpanId() {
+        return spanId;
+    }
+}

src/main/java/org/datadog/jenkins/plugins/datadog/traces/DatadogTraceBuildLogic.java

@@ -74,6 +74,9 @@ public void startBuildTrace(final BuildData buildData, Run run) {
         final StepDataAction stepDataAction = new StepDataAction();
         run.addAction(stepDataAction);
 
+        final StepTraceDataAction stepTraceDataAction = new StepTraceDataAction();
+        run.addAction(stepTraceDataAction);
+
         final StageBreakdownAction stageBreakdownAction = new StageBreakdownAction();
         run.addAction(stageBreakdownAction);
 

src/main/java/org/datadog/jenkins/plugins/datadog/traces/DatadogTracePipelineLogic.java

@@ -238,7 +238,8 @@ private void collectTraces(final Run run, final List<PayloadMessage> spanBuffer,
         final long fixedEndTimeNanos = TimeUnit.MICROSECONDS.toNanos(current.getEndTimeMicros() - TimeUnit.MILLISECONDS.toMicros(propagatedMillisInQueue));
 
         // At this point, the current node is traceable.
-        final TraceSpan span = new TraceSpan(buildOperationName(current), fixedStartTimeNanos + getNanosInQueue(current), parentSpanContext);
+        final TraceSpan.TraceSpanContext spanContext = new TraceSpan.TraceSpanContext(parentSpanContext.getTraceId(), parentSpanContext.getSpanId(), current.getSpanId());
+        final TraceSpan span = new TraceSpan(buildOperationName(current), fixedStartTimeNanos + getNanosInQueue(current), spanContext);
         span.setServiceName(DatadogUtilities.getDatadogGlobalDescriptor().getCiInstanceName());
         span.setResourceName(current.getName());
         span.setType("ci");

src/main/java/org/datadog/jenkins/plugins/datadog/traces/StepTraceDataAction.java

@@ -0,0 +1,30 @@
+package org.datadog.jenkins.plugins.datadog.traces;
+
+import hudson.model.InvisibleAction;
+import hudson.model.Run;
+import org.datadog.jenkins.plugins.datadog.model.StepTraceData;
+import org.jenkinsci.plugins.workflow.graph.FlowNode;
+
+import java.io.Serializable;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+public class StepTraceDataAction extends InvisibleAction implements Serializable {
+
+    private static final long serialVersionUID = 1L;
+
+    private final ConcurrentMap<String, StepTraceData> stepTraceDataByDescriptor = new ConcurrentHashMap<>();
+
+    public StepTraceData synchronizedPut(final Run<?,?> run, final FlowNode flowNode, final StepTraceData stepTraceData) {
+        synchronized (run){
+            return stepTraceDataByDescriptor.put(flowNode.getId(), stepTraceData);
+        }
+    }
+
+    public StepTraceData synchronizedGet(final Run<?,?> run, final FlowNode flowNode) {
+        synchronized (run){
+            return stepTraceDataByDescriptor.get(flowNode.getId());
+        }
+    }
+
+}

src/main/java/org/datadog/jenkins/plugins/datadog/traces/TraceStepEnvironmentContributor.java

@@ -0,0 +1,83 @@
+package org.datadog.jenkins.plugins.datadog.traces;
+
+import static org.datadog.jenkins.plugins.datadog.traces.TracerConstants.SPAN_ID_ENVVAR_KEY;
+import static org.datadog.jenkins.plugins.datadog.traces.TracerConstants.TRACE_ID_ENVVAR_KEY;
+
+import hudson.EnvVars;
+import hudson.Extension;
+import hudson.model.Run;
+import hudson.model.TaskListener;
+import org.datadog.jenkins.plugins.datadog.DatadogUtilities;
+import org.datadog.jenkins.plugins.datadog.model.StepTraceData;
+import org.jenkinsci.plugins.workflow.cps.nodes.StepAtomNode;
+import org.jenkinsci.plugins.workflow.graph.FlowNode;
+import org.jenkinsci.plugins.workflow.steps.StepContext;
+import org.jenkinsci.plugins.workflow.steps.StepEnvironmentContributor;
+
+import java.io.IOException;
+import java.util.logging.Logger;
+
+@Extension
+public class TraceStepEnvironmentContributor extends StepEnvironmentContributor {
+
+    private static final Logger logger = Logger.getLogger(TraceStepEnvironmentContributor.class.getName());
+
+    @Override
+    public void buildEnvironmentFor(StepContext stepContext, EnvVars envs, TaskListener listener) throws IOException, InterruptedException {
+        if (!DatadogUtilities.getDatadogGlobalDescriptor().isEnabledCiVisibility()) {
+            return;
+        }
+
+        try {
+            final Run<?,?> run = stepContext.get(Run.class);
+            if(run == null) {
+                logger.fine("Unable to set trace ids as environment variables. Run is null");
+                return;
+            }
+
+            final BuildSpanAction buildSpanAction = run.getAction(BuildSpanAction.class);
+            if(buildSpanAction == null) {
+                logger.fine("Unable to set trace ids as environment variables. in Run '"+run.getFullDisplayName()+"'. BuildSpanAction is null");
+                return;
+            }
+
+            final StepTraceDataAction stepTraceDataAction = run.getAction(StepTraceDataAction.class);
+            if(stepTraceDataAction == null) {
+                logger.fine("Unable to set trace ids as environment variables. in Run '"+run.getFullDisplayName()+"'. StepTraceDataAction is null");
+                return;
+            }
+
+            final FlowNode flowNode = stepContext.get(FlowNode.class);
+            if(flowNode == null) {
+                logger.fine("Unable to set trace ids as environment variables. in Run '"+run.getFullDisplayName()+"'. FlowNode is null");
+                return;
+            }
+
+            if(!(flowNode instanceof StepAtomNode)){
+                return;
+            }
+
+            StepTraceData stepTraceData = stepTraceDataAction.synchronizedGet(run, flowNode);
+            if(stepTraceData == null){
+                stepTraceData = new StepTraceData(IdGenerator.generate());
+                stepTraceDataAction.synchronizedPut(run, flowNode, stepTraceData);
+            }
+
+            if(envs.get(TRACE_ID_ENVVAR_KEY) == null) {
+                final String traceIdStr = Long.toUnsignedString(buildSpanAction.getBuildSpanContext().getTraceId());
+                envs.put(TRACE_ID_ENVVAR_KEY, traceIdStr);
+                logger.fine("Set DD_CUSTOM_TRACE_ID="+traceIdStr+" for FlowNode: "+flowNode);
+            }
+
+            if(envs.get(SPAN_ID_ENVVAR_KEY) == null) {
+                final String spanIdStr  = Long.toUnsignedString(stepTraceData.getSpanId());
+                envs.put(SPAN_ID_ENVVAR_KEY, spanIdStr);
+                logger.fine("Set DD_CUSTOM_PARENT_ID="+spanIdStr+" for FlowNode: "+flowNode);
+            }
+
+        } catch (Exception ex) {
+            logger.severe("Unable to set traces IDs as environment variables before step execution. " + ex);
+        }
+    }
+
+}

src/main/java/org/datadog/jenkins/plugins/datadog/traces/TracerConstants.java

@@ -0,0 +1,12 @@
+package org.datadog.jenkins.plugins.datadog.traces;
+
+public final class TracerConstants {
+
+    private TracerConstants(){}
+
+    // EnvVar keys to support custom spans.
+    // These env vars will be used by datadog-ci CLI to continue the trace.
+    public static final String TRACE_ID_ENVVAR_KEY = "DD_CUSTOM_TRACE_ID";
+    public static final String SPAN_ID_ENVVAR_KEY = "DD_CUSTOM_PARENT_ID";
+
+}

src/main/java/org/datadog/jenkins/plugins/datadog/traces/message/TraceSpan.java

@@ -28,10 +28,10 @@ public TraceSpan(final String name, final long startNano) {
         this(name, startNano, null);
     }
 
-    public TraceSpan(final String name, final long startNano, final TraceSpanContext parent) {
+    public TraceSpan(final String name, final long startNano, final TraceSpanContext spanContext) {
         this.operationName = name;
         this.startNano = startNano;
-        this.traceSpanContext = parent != null ? new TraceSpanContext(parent) : new TraceSpanContext();
+        this.traceSpanContext = spanContext != null ? spanContext : new TraceSpanContext();
 
         // Avoid sampling
         this.metrics.put(PRIORITY_SAMPLING_KEY, 1.0);
@@ -132,10 +132,10 @@ public TraceSpanContext() {
             this.parentId = 0;
         }
 
-        public TraceSpanContext(final TraceSpanContext parent) {
-            this.traceId = parent.getTraceId();
-            this.parentId = parent.getSpanId();
-            this.spanId = IdGenerator.generate();
+        public TraceSpanContext(final long traceId, final long parentId, final long spanId) {
+            this.traceId = traceId;
+            this.parentId = parentId;
+            this.spanId = (spanId != -1L) ? spanId : IdGenerator.generate();
         }
 
         public long getTraceId() {

src/test/java/org/datadog/jenkins/plugins/datadog/listeners/DatadogTraceAbstractTest.java

@@ -14,6 +14,7 @@
 import org.datadog.jenkins.plugins.datadog.traces.CITags;
 import org.datadog.jenkins.plugins.datadog.traces.IsPipelineAction;
 import org.datadog.jenkins.plugins.datadog.traces.StepDataAction;
+import org.datadog.jenkins.plugins.datadog.traces.StepTraceDataAction;
 import org.datadog.jenkins.plugins.datadog.traces.message.TraceSpan;
 
 import java.util.Map;
@@ -46,5 +47,6 @@ protected void assertCleanupActions(Run<?,?> run) {
         assertNull(run.getAction(PipelineQueueInfoAction.class));
         assertNull(run.getAction(StageBreakdownAction.class));
         assertNull(run.getAction(IsPipelineAction.class));
+        assertNull(run.getAction(StepTraceDataAction.class));
     }
 }