Trace Maven and Gradle build tasks (#7721)

Author: nikita-tkachenko-datadog

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

@@ -3,6 +3,7 @@
 import datadog.trace.api.civisibility.domain.BuildModuleLayout;
 import datadog.trace.api.civisibility.domain.BuildSessionSettings;
 import datadog.trace.api.civisibility.domain.JavaAgent;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.civisibility.config.JvmInfo;
 import java.nio.file.Path;
 import java.util.Collection;
@@ -25,6 +26,8 @@ BuildSystemModule testModuleStart(
       @Nullable Collection<Path> classpath,
       @Nullable JavaAgent jacocoAgent);
 
+  AgentSpan testTaskStart(String taskName);
+
   BuildSessionSettings getSettings();
 
   interface Factory {

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

@@ -1,5 +1,7 @@
 package datadog.trace.civisibility.domain.buildsystem;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.civisibility.CIConstants;
@@ -202,6 +204,11 @@ public BuildSystemModuleImpl testModuleStart(
         this::onModuleFinish);
   }
 
+  @Override
+  public AgentSpan testTaskStart(String taskName) {
+    return startSpan("ci_visibility", taskName, span.context());
+  }
+
   private void onModuleFinish(AgentSpan moduleSpan) {
     // multiple modules can finish in parallel
     synchronized (tagPropagationLock) {

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/BuildEventsHandlerImpl.java

@@ -5,6 +5,7 @@
 import datadog.trace.api.civisibility.domain.BuildSessionSettings;
 import datadog.trace.api.civisibility.domain.JavaAgent;
 import datadog.trace.api.civisibility.events.BuildEventsHandler;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.civisibility.config.JvmInfo;
 import datadog.trace.civisibility.config.JvmInfoFactory;
@@ -23,7 +24,10 @@ public class BuildEventsHandlerImpl<SessionKey> implements BuildEventsHandler<Se
   private final ConcurrentMap<SessionKey, BuildSystemSession> inProgressTestSessions =
       new ConcurrentHashMap<>();
 
-  private final ConcurrentMap<TestModuleDescriptor<SessionKey>, BuildSystemModule>
+  private final ConcurrentMap<BuildTaskDescriptor<SessionKey>, AgentSpan> inProgressBuildTasks =
+      new ConcurrentHashMap<>();
+
+  private final ConcurrentMap<BuildTaskDescriptor<SessionKey>, BuildSystemModule>
       inProgressTestModules = new ConcurrentHashMap<>();
 
   private final BuildSystemSession.Factory sessionFactory;
@@ -79,6 +83,47 @@ public void onTestSessionFinish(final SessionKey sessionKey) {
     testSession.end(null);
   }
 
+  @Override
+  public void onBuildTaskStart(
+      SessionKey sessionKey, String taskName, Map<String, Object> additionalTags) {
+    BuildSystemSession testSession = inProgressTestSessions.get(sessionKey);
+    if (testSession == null) {
+      throw new IllegalStateException("Could not find session span for key: " + sessionKey);
+    }
+    AgentSpan buildTask = testSession.testTaskStart(taskName);
+    for (Map.Entry<String, Object> e : additionalTags.entrySet()) {
+      buildTask.setTag(e.getKey(), e.getValue());
+    }
+    inProgressBuildTasks.put(new BuildTaskDescriptor<>(sessionKey, taskName), buildTask);
+  }
+
+  @Override
+  public void onBuildTaskFail(SessionKey sessionKey, String taskName, Throwable throwable) {
+    AgentSpan buildTask = inProgressBuildTasks.get(new BuildTaskDescriptor<>(sessionKey, taskName));
+    if (buildTask == null) {
+      throw new IllegalStateException(
+          "Could not find build task span for session key "
+              + sessionKey
+              + " and task name "
+              + taskName);
+    }
+    buildTask.setError(true);
+    buildTask.addThrowable(throwable);
+  }
+
+  @Override
+  public void onBuildTaskFinish(SessionKey sessionKey, String taskName) {
+    AgentSpan buildTask = inProgressBuildTasks.get(new BuildTaskDescriptor<>(sessionKey, taskName));
+    if (buildTask == null) {
+      throw new IllegalStateException(
+          "Could not find build task span for session key "
+              + sessionKey
+              + " and task name "
+              + taskName);
+    }
+    buildTask.finish();
+  }
+
   @Override
   public BuildModuleSettings onTestModuleStart(
       final SessionKey sessionKey,
@@ -104,8 +149,8 @@ public BuildModuleSettings onTestModuleStart(
       }
     }
 
-    TestModuleDescriptor<SessionKey> testModuleDescriptor =
-        new TestModuleDescriptor<>(sessionKey, moduleName);
+    BuildTaskDescriptor<SessionKey> testModuleDescriptor =
+        new BuildTaskDescriptor<>(sessionKey, moduleName);
     inProgressTestModules.put(testModuleDescriptor, testModule);
 
     return testModule.getSettings();
@@ -126,8 +171,8 @@ public void onTestModuleFail(
   }
 
   private BuildSystemModule getTestModule(final SessionKey sessionKey, final String moduleName) {
-    TestModuleDescriptor<SessionKey> testModuleDescriptor =
-        new TestModuleDescriptor<>(sessionKey, moduleName);
+    BuildTaskDescriptor<SessionKey> testModuleDescriptor =
+        new BuildTaskDescriptor<>(sessionKey, moduleName);
     BuildSystemModule testModule = inProgressTestModules.get(testModuleDescriptor);
     if (testModule == null) {
       throw new IllegalStateException(
@@ -138,8 +183,8 @@ private BuildSystemModule getTestModule(final SessionKey sessionKey, final Strin
 
   @Override
   public void onTestModuleFinish(SessionKey sessionKey, String moduleName) {
-    TestModuleDescriptor<SessionKey> testModuleDescriptor =
-        new TestModuleDescriptor<>(sessionKey, moduleName);
+    BuildTaskDescriptor<SessionKey> testModuleDescriptor =
+        new BuildTaskDescriptor<>(sessionKey, moduleName);
     BuildSystemModule testModule = inProgressTestModules.remove(testModuleDescriptor);
     if (testModule == null) {
       throw new IllegalStateException(

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/BuildTaskDescriptor.java

@@ -0,0 +1,31 @@
+package datadog.trace.civisibility.events;
+
+import java.util.Objects;
+
+public class BuildTaskDescriptor<T> {
+
+  private final T sessionKey;
+  private final String taskName;
+
+  public BuildTaskDescriptor(T sessionKey, String taskName) {
+    this.sessionKey = sessionKey;
+    this.taskName = taskName;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    BuildTaskDescriptor<?> that = (BuildTaskDescriptor<?>) o;
+    return sessionKey.equals(that.sessionKey) && taskName.equals(that.taskName);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(sessionKey, taskName);
+  }
+}

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/TestModuleDescriptor.java

@@ -8,6 +8,7 @@
 import datadog.trace.api.civisibility.domain.JavaAgent;
 import datadog.trace.api.civisibility.events.BuildEventsHandler;
 import java.nio.file.Path;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -95,14 +96,15 @@ static final class TestTaskExecutionListener implements TaskExecutionListener {
 
     @Override
     public void beforeExecute(@Nonnull Task task) {
-      if (!GradleUtils.isTestTask(task)) {
-        return;
-      }
-
       Project project = task.getProject();
       Gradle gradle = project.getGradle();
       String taskPath = task.getPath();
 
+      if (!GradleUtils.isTestTask(task)) {
+        buildEventsHandler.onBuildTaskStart(gradle, taskPath, Collections.emptyMap());
+        return;
+      }
+
       List<String> sourceSetNames = Config.get().getCiVisibilityJacocoGradleSourceSets();
       BuildModuleLayout moduleLayout = GradleUtils.getModuleLayout(project, sourceSetNames);
       Path jvmExecutable = GradleUtils.getEffectiveExecutable(task);
@@ -118,15 +120,19 @@ public void beforeExecute(@Nonnull Task task) {
 
     @Override
     public void afterExecute(@Nonnull Task task, @Nonnull TaskState state) {
-      if (!GradleUtils.isTestTask(task)) {
-        return;
-      }
-
       Project project = task.getProject();
       Gradle gradle = project.getGradle();
       String taskPath = task.getPath();
-
       Throwable failure = state.getFailure();
+
+      if (!GradleUtils.isTestTask(task)) {
+        if (failure != null) {
+          buildEventsHandler.onBuildTaskFail(gradle, taskPath, failure);
+        }
+        buildEventsHandler.onBuildTaskFinish(gradle, taskPath);
+        return;
+      }
+
       if (failure != null) {
         buildEventsHandler.onTestModuleFail(gradle, taskPath, failure);
 

dd-java-agent/instrumentation/gradle-3.0/src/main/groovy/datadog/trace/instrumentation/gradle/legacy/GradleBuildListener.java

@@ -173,13 +173,13 @@ public void projectsEvaluated(Gradle gradle) {
   @SuppressWarnings("unchecked")
   @Override
   public void beforeExecute(TaskIdentity<?> taskIdentity) {
+    String taskPath = taskIdentity.getTaskPath();
     if (!Test.class.isAssignableFrom(taskIdentity.getTaskType())) {
+      ciVisibilityService.onBuildTaskStart(taskPath);
       return;
     }
 
     String projectPath = taskIdentity.getProjectPath();
-    String taskPath = taskIdentity.getTaskPath();
-
     Project project = gradle.getRootProject().project(projectPath);
     Test task = (Test) project.getTasks().getByName(taskIdentity.name);
 
@@ -226,12 +226,14 @@ private JavaAgent getJacocoAgent(Test task) {
 
   @Override
   public void afterExecute(TaskIdentity<?> taskIdentity, TaskState state) {
+    String taskPath = taskIdentity.getTaskPath();
+    Throwable failure = state.getFailure();
+
     if (!Test.class.isAssignableFrom(taskIdentity.getTaskType())) {
+      ciVisibilityService.onBuildTaskFinish(taskPath, failure);
       return;
     }
 
-    String taskPath = taskIdentity.getTaskPath();
-    Throwable failure = state.getFailure();
     String reason = state.getSkipped() || !state.getDidWork() ? state.getSkipMessage() : null;
     ciVisibilityService.onModuleFinish(taskPath, failure, reason);
   }

dd-java-agent/instrumentation/gradle-8.3/src/main/java/datadog/trace/instrumentation/gradle/CiVisibilityGradleListener.java

@@ -104,6 +104,17 @@ public void onBuildStart(
         SESSION_KEY, buildPath, projectRoot, startCommand, "gradle", gradleVersion, additionalTags);
   }
 
+  public void onBuildTaskStart(String taskPath) {
+    buildEventsHandler.onBuildTaskStart(SESSION_KEY, taskPath, Collections.emptyMap());
+  }
+
+  public void onBuildTaskFinish(String taskPath, @Nullable Throwable failure) {
+    if (failure != null) {
+      buildEventsHandler.onBuildTaskFail(SESSION_KEY, taskPath, failure);
+    }
+    buildEventsHandler.onBuildTaskFinish(SESSION_KEY, taskPath);
+  }
+
   public void onModuleStart(
       String taskPath,
       BuildModuleLayout moduleLayout,