Decouple Muzzle planning from CI slot selection

buildSrc/src/main/kotlin/datadog/gradle/plugin/muzzle/MuzzleExtension.kt

@@ -3,17 +3,13 @@ package datadog.gradle.plugin.muzzle
 import org.gradle.api.Action
 import org.gradle.api.model.ObjectFactory
 import org.gradle.kotlin.dsl.newInstance
-import org.gradle.kotlin.dsl.property
 import java.util.Locale
 import javax.inject.Inject
 
 /**
  * Muzzle extension containing all pass and fail directives.
  */
 abstract class MuzzleExtension @Inject constructor(private val objectFactory: ObjectFactory) {
-    /** Whether aggregate runMuzzle invocations should plan checks for this project. */
-    val includeInAggregate = objectFactory.property<Boolean>().convention(true)
-
     val directives: MutableList<MuzzleDirective> = ArrayList()
     private val additionalRepositories: MutableList<Triple<String, String, String>> = ArrayList()
 

buildSrc/src/main/kotlin/datadog/gradle/plugin/muzzle/MuzzlePlugin.kt

@@ -9,6 +9,7 @@ import org.eclipse.aether.artifact.Artifact
 import org.gradle.api.NamedDomainObjectProvider
 import org.gradle.api.Plugin
 import org.gradle.api.Project
+import org.gradle.api.Task
 import org.gradle.api.artifacts.Configuration
 import org.gradle.api.tasks.TaskProvider
 import org.gradle.kotlin.dsl.create
@@ -26,8 +27,7 @@ class MuzzlePlugin : Plugin<Project> {
    *
    * This function sets up the necessary configurations, dependencies, and tasks for muzzle validation.
    * It registers the muzzle extension, configures bootstrap/tooling dependencies, and sets up tasks for
-   * compiling, running, printing, and reporting muzzle checks. It also optimizes configuration overhead
-   * by only proceeding if relevant muzzle tasks are requested.
+   * compiling, running, printing, and reporting muzzle checks. Validation tasks are registered lazily.
    *
    * @param project The Gradle project to apply the plugin to.
    */
@@ -101,46 +101,22 @@ class MuzzlePlugin : Plugin<Project> {
     }
     project.publishMuzzleReport(report.flatMap { it.versionsFile })
 
-    val relevantTasks = project.gradle.startParameter.taskNames.filter { taskName ->
-      val taskProjectPath = taskName.substringBeforeLast(":", "")
-      val taskNameOnly = taskName.substringAfterLast(":")
-      val isAggregate = taskNameOnly.equals("runMuzzle", ignoreCase = true)
-      val isRelevantForProject = taskProjectPath.isEmpty() || taskProjectPath == project.path ||
-        (isAggregate && project.path.startsWith("$taskProjectPath:"))
-
-      isRelevantForProject && !taskNameOnly.equals("compileMuzzle", ignoreCase = true) &&
-        taskNameOnly.endsWith("muzzle", ignoreCase = true)
-    }
-    if (relevantTasks.isEmpty()) {
-      // Adding muzzle dependencies has a large config overhead. Stop unless muzzle is explicitly run.
-      project.logger.info("No muzzle tasks invoked for ${project.path}, skipping muzzle task planification")
-      return
-    }
-
-    // We only get here if we are running muzzle, so let's start timing things
-    val startTime = System.currentTimeMillis()
-
-    project.afterEvaluate {
-      if (relevantTasks.all { it.substringAfterLast(":").equals("runMuzzle", ignoreCase = true) } &&
-        !extension.includeInAggregate.get()) {
-        project.logger.info("No muzzle tasks invoked for ${project.path}, skipping muzzle task planification")
-        return@afterEvaluate
-      }
+    val validationTasks = MuzzleValidationTasks(project.providers) {
+      val startTime = System.currentTimeMillis()
       val system = MuzzleMavenRepoUtils.newRepositorySystem()
       val session = MuzzleMavenRepoUtils.newRepositorySystemSession(system)
       val taskPlanner = MuzzleTaskPlanner.from(system, session)
-      // use runAfter to set up task finalizers in version order
-      var runAfter: TaskProvider<MuzzleTask> = muzzleTask
-      val muzzleReportTasks = mutableListOf<TaskProvider<MuzzleTask>>()
       val directives = project.extensions.getByType<MuzzleExtension>().directives
-      taskPlanner.plan(directives).forEach { plan ->
-        runAfter = registerMuzzleTask(plan.directive, plan.artifact, project, runAfter, muzzleBootstrap, muzzleTooling)
-        muzzleReportTasks.add(runAfter)
+      val checks = taskPlanner.plan(directives).map { plan ->
+        val check = registerMuzzleTask(plan.directive, plan.artifact, project, muzzleBootstrap, muzzleTooling)
         project.logger.info("configured ${plan.directive}")
+        check
       }
 
-      if (muzzleReportTasks.isEmpty() && !directives.any { it.assertPass }) {
-        muzzleReportTasks.add(muzzleTask)
+      val muzzleReportTasks = if (checks.isEmpty() && !directives.any { it.assertPass }) {
+        listOf(muzzleTask)
+      } else {
+        checks
       }
 
       val timingTask = project.tasks.register<MuzzleEndTask>("muzzle-end") {
@@ -152,10 +128,17 @@ class MuzzlePlugin : Plugin<Project> {
         )
         muzzleResultFiles.from(muzzleReportTasks.map { it.flatMap { task -> task.result } })
       }
-      // last muzzle task to run
-      runAfter.configure {
-        finalizedBy(timingTask)
+
+      val executionChain: List<TaskProvider<out Task>> = checks + timingTask
+      executionChain.zipWithNext { current, next ->
+        current.configure { finalizedBy(next) }
       }
+
+      // Note this is unordered set, the execution order is configured via the finalizer relationships
+      setOf(executionChain.first().get())
+    }
+    muzzleTask.configure {
+      finalizedBy(validationTasks.tasks)
     }
   }
 
@@ -169,7 +152,6 @@ class MuzzlePlugin : Plugin<Project> {
      * @param muzzleDirective The directive describing the dependency and test parameters.
      * @param versionArtifact The artifact representing the dependency version to test (may be null when `muzzleDirective.coreJdk` is true).
      * @param instrumentationProject The Gradle project to register the task in.
-     * @param runAfterTask The task provider for the task that this muzzle task should run after.
      * @param muzzleBootstrap The configuration provider for agent bootstrap dependencies.
      * @param muzzleTooling The configuration provider for agent tooling dependencies.
      * @return The muzzle task provider.
@@ -178,7 +160,6 @@ class MuzzlePlugin : Plugin<Project> {
       muzzleDirective: MuzzleDirective,
       versionArtifact: Artifact?,
       instrumentationProject: Project,
-      runAfterTask: TaskProvider<MuzzleTask>,
       muzzleBootstrap: NamedDomainObjectProvider<Configuration>,
       muzzleTooling: NamedDomainObjectProvider<Configuration>
     ): TaskProvider<MuzzleTask> {
@@ -246,9 +227,6 @@ class MuzzlePlugin : Plugin<Project> {
         this.muzzleTooling.set(muzzleTooling)
       }
 
-      runAfterTask.configure {
-        finalizedBy(muzzleTask)
-      }
       return muzzleTask
     }
   }

buildSrc/src/main/kotlin/datadog/gradle/plugin/muzzle/MuzzleValidationTasks.kt

@@ -0,0 +1,41 @@
+package datadog.gradle.plugin.muzzle
+
+import org.gradle.api.Task
+import org.gradle.api.provider.Provider
+import org.gradle.api.provider.ProviderFactory
+import org.gradle.api.tasks.TaskProvider
+
+/**
+ * Registers validation tasks when Gradle resolves a module's muzzle finalizers.
+ *
+ * [tasks] is a plain [Provider] used as dependency notation by [Task.finalizedBy]. During normal
+ * task-graph construction, Gradle resolves its task-set value after project evaluation and before
+ * execution. Registration therefore runs outside task configuration actions and can use completed
+ * directives to discover checks whose names were unknown when the plugin was applied.
+ *
+ * A [TaskProvider] references an already registered task. A task-backed [Provider.map] can supply
+ * producer dependencies without evaluating its mapping, so registration must not be a side effect
+ * of that mapping. [ProviderFactory.provider] supplies the task-set value directly instead.
+ *
+ * Providers can be queried repeatedly. Cache both success and failure to avoid duplicate task
+ * registration, including retries after partial registration.
+ *
+ * Note that [org.gradle.api.tasks.TaskDependency], a mechanism to supplying directly the **concrete**
+ * tasks when assembling the execution graph, exists today but is on the side-line possibly
+ * getting deprecated for removal.
+ *
+ * shortcut: generated-check CLI selectors remain unsupported; revisit if individual checks
+ * need to be selected or excluded.
+ *
+ * @see <a href="https://docs.gradle.org/current/javadoc/org/gradle/api/provider/ProviderFactory.html#provider(java.util.concurrent.Callable)">ProviderFactory.provider API</a>
+ * @see <a href="https://docs.gradle.org/current/javadoc/org/gradle/api/tasks/TaskProvider.html">TaskProvider API</a>
+ * @see <a href="https://github.com/gradle/gradle/issues/30534#issuecomment-2360291115">Discussion on removing older Gradle interface API</a>
+ */
+internal class MuzzleValidationTasks(
+  providerFactory: ProviderFactory,
+  registerValidationTasks: () -> Set<Task>
+) {
+  private val registration by lazy { runCatching(registerValidationTasks) }
+
+  val tasks: Provider<Set<Task>> = providerFactory.provider { registration.getOrThrow() }
+}

buildSrc/src/main/kotlin/dd-trace-java.ci-jobs.gradle.kts

@@ -1,5 +1,4 @@
 import datadog.gradle.plugin.ci.isInSelectedSlot
-import datadog.gradle.plugin.muzzle.MuzzleExtension
 import org.gradle.api.tasks.testing.Test
 import java.io.File
 
@@ -16,12 +15,6 @@ if (project != rootProject) {
 }
 
 allprojects {
-  pluginManager.withPlugin("dd-trace-java.muzzle") {
-    extensions.configure<MuzzleExtension> {
-      includeInAggregate.convention(project.isInSelectedSlot)
-    }
-  }
-
   // Enable tests only on the selected slot (if -Pslot=n/t is provided)
   tasks.withType<Test>().configureEach {
     onlyIf("Project is in selected slot") {
@@ -133,7 +126,7 @@ tasks.register("runMuzzle") {
   val muzzleSubprojects = subprojects.filter { p ->
     p.plugins.hasPlugin("java") &&
       p.plugins.hasPlugin("dd-trace-java.muzzle") &&
-      p.extensions.getByType<MuzzleExtension>().includeInAggregate.get()
+      p.isInSelectedSlot.get()
   }
   dependsOn(muzzleSubprojects.map { p -> "${p.path}:muzzle" })
 }

buildSrc/src/test/kotlin/datadog/gradle/plugin/muzzle/MuzzlePluginFunctionalTest.kt

@@ -8,8 +8,9 @@ import org.junit.jupiter.params.provider.ValueSource
 import kotlin.io.path.readText
 
 class MuzzlePluginFunctionalTest : MuzzlePluginTestFixture() {
-  @Test
-  fun `nested aggregate plans descendant instrumentation checks`() {
+  @ParameterizedTest
+  @ValueSource(strings = [":dd-java-agent:runMuzzle", "dd-java-agent:runMuzzle"])
+  fun `nested aggregate plans descendant instrumentation checks`(taskName: String) {
     writeProject(
       """
       plugins {
@@ -25,14 +26,14 @@ class MuzzlePluginFunctionalTest : MuzzlePluginTestFixture() {
       tasks.register("runMuzzle") { dependsOn(":dd-java-agent:instrumentation:demo:muzzle") }
       """
     )
-    val result = run(":dd-java-agent:runMuzzle")
+    val result = run(taskName)
     assertThat(result.output).contains("BUILD SUCCESSFUL")
     assertThat(result.task(":dd-java-agent:instrumentation:demo:muzzle-AssertPass-core-jdk")?.outcome).isEqualTo(SUCCESS)
   }
 
   @ParameterizedTest
-  @ValueSource(strings = ["muzzle", ":dd-java-agent:instrumentation:demo:muzzle", "runMuzzle"])
-  fun `detects muzzle invocation with various task names`(taskName: String) {
+  @ValueSource(strings = ["muzzle", ":dd-java-agent:instrumentation:demo:muzzle", "runMuzzle", "verifyCompatibility"])
+  fun `runs validation through different task entry points`(taskName: String) {
     writeProject(
       """
       plugins {
@@ -54,21 +55,23 @@ class MuzzlePluginFunctionalTest : MuzzlePluginTestFixture() {
       tasks.register("runMuzzle") {
         dependsOn(":dd-java-agent:instrumentation:demo:muzzle")
       }
+      tasks.register("verifyCompatibility") {
+        dependsOn(":dd-java-agent:instrumentation:demo:muzzle")
+      }
       """
     )
 
     writeNoopScanPlugin()
 
     val result = run(taskName, "--stacktrace")
 
-    assertThat(result.tasks)
-      .withFailMessage("Should create muzzle tasks when '$taskName' is requested")
-      .anyMatch { it.path.contains("muzzle") }
-    assertThat(result.output)
-      .withFailMessage("Should not skip muzzle task planification when '$taskName' is requested")
-      .doesNotContain("No muzzle tasks invoked, skipping muzzle task planification")
-    assertThat(result.tasks).withFailMessage("Should execute muzzle tasks when '$taskName' is requested")
-      .anyMatch { it.path == ":dd-java-agent:instrumentation:demo:muzzle" || it.path.contains("muzzle-Assert") }
+    assertThat(result.output).contains("BUILD SUCCESSFUL")
+    assertThat(result.task(":dd-java-agent:instrumentation:demo:muzzle-AssertPass-core-jdk")?.outcome)
+      .withFailMessage("Should execute validation through '$taskName'")
+      .isEqualTo(SUCCESS)
+    assertThat(result.task(":dd-java-agent:instrumentation:demo:muzzle-end")?.outcome)
+      .withFailMessage("Should execute reporting through '$taskName'")
+      .isEqualTo(SUCCESS)
   }
 
   @Test