Fixed race condition on dumping future cleanup.

build.gradle.kts

@@ -3,14 +3,13 @@ import com.diffplug.gradle.spotless.SpotlessExtension
 plugins {
   id("datadog.gradle-debug")
   id("datadog.dependency-locking")
+  id("datadog.tracer-version")
+  id("datadog.dump-hanged-test")
 
   id("com.diffplug.spotless") version "6.13.0"
   id("com.github.spotbugs") version "5.0.14"
   id("de.thetaphi.forbiddenapis") version "3.8"
-
-  id("tracer-version")
   id("io.github.gradle-nexus.publish-plugin") version "2.0.0"
-
   id("com.gradleup.shadow") version "8.3.6" apply false
   id("me.champeau.jmh") version "0.7.3" apply false
   id("org.gradle.playframework") version "0.13" apply false

buildSrc/build.gradle.kts

@@ -27,9 +27,13 @@ gradlePlugin {
       implementationClass = "datadog.gradle.plugin.CallSiteInstrumentationPlugin"
     }
     create("tracer-version-plugin") {
-      id = "tracer-version"
+      id = "datadog.tracer-version"
       implementationClass = "datadog.gradle.plugin.version.TracerVersionPlugin"
     }
+    create("dump-hanged-test-plugin") {
+      id = "datadog.dump-hanged-test"
+      implementationClass = "datadog.gradle.plugin.dump.DumpHangedTestPlugin"
+    }
     create("supported-config-generation") {
       id = "supported-config-generator"
       implementationClass = "datadog.gradle.plugin.config.SupportedConfigPlugin"

buildSrc/src/integTest/kotlin/datadog/gradle/plugin/version/DumpHangedTestIntegrationTest.kt

@@ -0,0 +1,123 @@
+package datadog.gradle.plugin.version
+
+import org.gradle.testkit.runner.GradleRunner
+import org.gradle.testkit.runner.UnexpectedBuildFailure
+import org.junit.jupiter.api.Assertions.assertFalse
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertNotNull
+import org.junit.jupiter.api.io.TempDir
+import java.io.File
+import java.nio.file.Paths
+
+class DumpHangedTestIntegrationTest {
+  @Test
+  fun `should not take dumps`(@TempDir projectDir: File) {
+    val output = runGradleTest(projectDir, testSleep = 1000)
+
+    // Assert Gradle output has no evidence of taking dumps.
+    assertFalse(output.contains("Taking dumps after 15 seconds delay for :test"))
+    assertFalse(output.contains("Requesting stop of task ':test' as it has exceeded its configured timeout of 20s."))
+
+    assertTrue(file(projectDir, "build").exists()) // Assert build happened.
+    assertFalse(file(projectDir, "build", "dumps").exists()) // Assert no dumps created.
+  }
+
+  @Test
+  fun `should take dumps`(@TempDir projectDir: File) {
+    val output = runGradleTest(projectDir, testSleep = 25_0000)
+
+    // Assert Gradle output has evidence of taking dumps.
+    assertTrue(output.contains("Taking dumps after 15 seconds delay for :test"))
+    assertTrue(output.contains("Requesting stop of task ':test' as it has exceeded its configured timeout of 20s."))
+
+    assertTrue(file(projectDir, "build").exists()) // Assert build happened.
+
+    val dumps = file(projectDir, "build", "dumps")
+    assertTrue(dumps.exists()) // Assert dumps created.
+
+    // Assert actual dumps created.
+    val dumpFiles = dumps.list()
+    assertNotNull(dumpFiles.find { it.endsWith(".hprof") })
+    assertNotNull(dumpFiles.find { it.startsWith("all-thread-dumps") })
+  }
+
+  private fun runGradleTest(projectDir: File, testSleep: Long): List<String> {
+    file(projectDir, "settings.gradle.kts").writeText(
+      """
+      rootProject.name = "test-project"
+      """.trimIndent()
+    )
+
+    file(projectDir, "build.gradle.kts").writeText(
+      """
+      import java.time.Duration
+
+      plugins {
+        id("java")
+        id("datadog.dump-hanged-test")
+      }
+
+      group = "datadog.dump.test"
+
+      repositories {
+        mavenCentral()
+      }
+
+      dependencies {
+        testImplementation(platform("org.junit:junit-bom:5.10.0"))
+        testImplementation("org.junit.jupiter:junit-jupiter")
+        testRuntimeOnly("org.junit.platform:junit-platform-launcher")
+      }
+
+      dumpHangedTest {
+        // Set the dump offset for 5 seconds to trigger taking dumps after 15 seconds.
+        dumpOffset.set(5)
+      }
+
+      tasks.withType<Test>().configureEach {
+        // Set test timeout after 20 seconds.
+        timeout.set(Duration.ofSeconds(20))
+
+        useJUnitPlatform()
+      }
+      """.trimIndent()
+    )
+
+    file(projectDir, "src", "test", "java", "SimpleTest.java", makeDirectory = true).writeText(
+      """
+      import org.junit.jupiter.api.Test;
+
+      public class SimpleTest {
+          @Test
+          public void test() throws InterruptedException {
+              Thread.sleep($testSleep);
+          }
+      }
+      """.trimIndent()
+    )
+
+    try {
+      val buildResult = GradleRunner.create()
+        .forwardOutput()
+        .withPluginClasspath()
+        .withArguments("test")
+        .withProjectDir(projectDir)
+        .build()
+
+      return buildResult.output.lines()
+    } catch (e: UnexpectedBuildFailure) {
+      return e.buildResult.output.lines()
+    }
+  }
+
+  private fun file(projectDir: File, vararg parts: String, makeDirectory: Boolean = false): File {
+    val f = Paths.get(projectDir.absolutePath, *parts).toFile()
+
+    if (makeDirectory) {
+      f.parentFile.mkdirs()
+    }
+
+    return f
+  }
+}

buildSrc/src/integTest/kotlin/datadog/gradle/plugin/version/TracerVersionIntegrationTest.kt

@@ -265,7 +265,7 @@ class TracerVersionIntegrationTest {
     File(projectDir, "build.gradle.kts").writeText(
       """
       plugins {
-        id("tracer-version")
+        id("datadog.tracer-version")
       }
 
       tasks.register("printVersion") {

buildSrc/src/main/kotlin/datadog/gradle/plugin/dump/DumpHangedTestPlugin.kt

@@ -0,0 +1,177 @@
+package datadog.gradle.plugin.dump
+
+import org.gradle.api.Plugin
+import org.gradle.api.Project
+import org.gradle.api.Task
+import org.gradle.api.model.ObjectFactory
+import org.gradle.api.provider.Property
+import org.gradle.api.provider.Provider
+import org.gradle.api.services.BuildService
+import org.gradle.api.services.BuildServiceParameters
+import org.gradle.api.tasks.testing.Test
+import org.gradle.kotlin.dsl.extra
+import org.gradle.kotlin.dsl.withType
+import java.io.File
+import java.io.IOException
+import java.lang.ProcessBuilder.Redirect
+import java.time.Duration
+import java.util.concurrent.Executors
+import java.util.concurrent.ScheduledExecutorService
+import java.util.concurrent.ScheduledFuture
+import java.util.concurrent.TimeUnit
+import javax.inject.Inject
+
+/**
+ * Plugin to collect thread and heap dumps for hanged tests.
+ */
+class DumpHangedTestPlugin : Plugin<Project> {
+  companion object {
+    private const val DUMP_FUTURE_KEY = "dumping_future"
+  }
+
+  /** Plugin properties */
+  abstract class DumpHangedTestProperties @Inject constructor(objects: ObjectFactory) {
+    // Time offset (in seconds) before a test reaches its timeout at which dumps should be started.
+    // Defaults to 60 seconds.
+    val dumpOffset: Property<Long> = objects.property(Long::class.java)
+  }
+
+  /** Executor wrapped with proper Gradle lifecycle. */
+  abstract class DumpSchedulerService : BuildService<BuildServiceParameters.None>, AutoCloseable {
+    private val executor: ScheduledExecutorService =
+      Executors.newSingleThreadScheduledExecutor { r -> Thread(r, "hanged-test-dump").apply { isDaemon = true } }
+
+    fun schedule(task: () -> Unit, delay: Duration): ScheduledFuture<*> =
+      executor.schedule(task, delay.toMillis(), TimeUnit.MILLISECONDS)
+
+    override fun close() {
+      executor.shutdownNow()
+    }
+  }
+
+  override fun apply(project: Project) {
+    if (project.rootProject != project) {
+      return
+    }
+
+    val scheduler = project.gradle.sharedServices
+      .registerIfAbsent("dumpHangedTestScheduler", DumpSchedulerService::class.java)
+
+    // Create plugin properties.
+    val props = project.extensions.create("dumpHangedTest", DumpHangedTestProperties::class.java)
+
+    fun configure(p: Project) {
+      p.tasks.withType<Test>().configureEach {
+        doFirst { schedule(this, scheduler, props) }
+        doLast { cleanup(this) }
+      }
+    }
+
+    configure(project)
+
+    project.subprojects(::configure)
+  }
+
+  private fun schedule(t: Task, scheduler: Provider<DumpSchedulerService>, props: DumpHangedTestProperties) {
+    val taskName = t.path
+
+    if (t.extra.has(DUMP_FUTURE_KEY)) {
+      t.logger.info("Taking dumps already scheduled for $taskName")
+      return
+    }
+
+    val dumpOffset = props.dumpOffset.getOrElse(60)
+    val delay = t.timeout.map { it.minusSeconds(dumpOffset) }.orNull
+
+    if (delay == null || delay.seconds < 0) {
+      t.logger.info("Taking dumps has invalid timeout configured for $taskName")
+      return
+    }
+
+    val future = scheduler.get().schedule({
+      t.logger.quiet("Taking dumps after ${delay.seconds} seconds delay for $taskName")
+
+      takeDump(t)
+    }, delay)
+
+    t.extra.set(DUMP_FUTURE_KEY, future)
+  }
+
+  private fun takeDump(t: Task) {
+    try {
+      // Use Gradle's build dir and adjust for CI artifacts collection if needed.
+      val dumpsDir: File = t.project.layout.buildDirectory
+        .dir("dumps")
+        .map { dir ->
+          if (t.project.providers.environmentVariable("CI").isPresent) {
+            // Move reports into the folder collected by the collect_reports.sh script.
+            File(
+              dir.asFile.absolutePath.replace(
+                "dd-trace-java/dd-java-agent",
+                "dd-trace-java/workspace/dd-java-agent"
+              )
+            )
+          } else {
+            dir.asFile
+          }
+        }
+        .get()
+
+      dumpsDir.mkdirs()
+
+      fun file(name: String, ext: String = "log") =
+        File(dumpsDir, "$name-${System.currentTimeMillis()}.$ext")
+
+      // For simplicity, use `0` as the PID, which collects all thread dumps across JVMs.
+      val allThreadsFile = file("all-thread-dumps")
+      runCmd(Redirect.to(allThreadsFile), "jcmd", "0", "Thread.print", "-l")
+
+      // Collect all JVMs pids.
+      val allJavaProcessesFile = file("all-java-processes")
+      runCmd(Redirect.to(allJavaProcessesFile), "jcmd", "-l")
+
+      // Collect pids for 'Gradle Test Executor'.
+      val pids = allJavaProcessesFile.readLines()
+        .filter { it.contains("Gradle Test Executor") }
+        .map { it.substringBefore(' ') }
+
+      pids.forEach { pid ->
+        // Collect heap dump by pid.
+        val heapDumpPath = file("${pid}-heap-dump", "hprof").absolutePath
+        runCmd(Redirect.INHERIT, "jcmd", pid, "GC.heap_dump", heapDumpPath)
+
+        // Collect thread dump by pid.
+        val threadDumpFile = file("${pid}-thread-dump")
+        runCmd(Redirect.to(threadDumpFile), "jcmd", pid, "Thread.print", "-l")
+      }
+    } catch (e: Throwable) {
+      t.logger.warn("Taking dumps failed with error: ${e.message}, for ${t.path}")
+    }
+  }
+
+  private fun cleanup(t: Task) {
+    val future = t.extra
+      .takeIf { it.has(DUMP_FUTURE_KEY) }
+      ?.get(DUMP_FUTURE_KEY) as? ScheduledFuture<*>
+
+    if (future != null && !future.isDone) {
+      t.logger.info("Taking dump canceled with remaining delay of ${future.getDelay(TimeUnit.SECONDS)} seconds for ${t.path}")
+      future.cancel(false)
+    }
+  }
+
+  private fun runCmd(
+    redirectTo: Redirect,
+    vararg args: String
+  ) {
+    val exitCode = ProcessBuilder(*args)
+      .redirectErrorStream(true)
+      .redirectOutput(redirectTo)
+      .start()
+      .waitFor()
+
+    if (exitCode != 0) {
+      throw IOException("Process failed: ${args.joinToString(" ")}, exit code: $exitCode")
+    }
+  }
+}

gradle/configure_tests.gradle

@@ -1,8 +1,6 @@
 import java.time.Duration
 import java.time.temporal.ChronoUnit
 
-apply from: "$rootDir/gradle/dump_hanging_test.gradle"
-
 def isTestingInstrumentation(Project project) {
   return [
     "junit-4.10",