Simplify tasks partitioning UX

.gitlab-ci.yml

@@ -109,10 +109,7 @@ default:
 .normalize_node_index: &normalize_node_index
   - if [ "$CI_NO_SPLIT" == "true" ] ; then CI_NODE_INDEX=1; CI_NODE_TOTAL=1; fi # A job uses parallel but doesn't intend to split by index
   - if [ -n "$CI_SPLIT" ]; then CI_NODE_INDEX="${CI_SPLIT%%/*}"; CI_NODE_TOTAL="${CI_SPLIT##*/}"; fi
-  - echo "CI_NODE_TOTAL=${CI_NODE_TOTAL}, CI_NODE_INDEX=$CI_NODE_INDEX"
-  - export NORMALIZED_NODE_TOTAL=${CI_NODE_TOTAL:-1}
-  - ONE_INDEXED_NODE_INDEX=${CI_NODE_INDEX:-1}; export NORMALIZED_NODE_INDEX=$((ONE_INDEXED_NODE_INDEX - 1))
-  - echo "NORMALIZED_NODE_TOTAL=${NORMALIZED_NODE_TOTAL}, NORMALIZED_NODE_INDEX=$NORMALIZED_NODE_INDEX"
+  - echo "CI_NODE_INDEX=$CI_NODE_INDEX, CI_NODE_TOTAL=${CI_NODE_TOTAL}"
 
 .cgroup_info: &cgroup_info
   - source .gitlab/gitlab-utils.sh
@@ -420,7 +417,7 @@ test_published_artifacts:
   script:
     - *gitlab_base_ref_params
     - ./gradlew --version
-    - ./gradlew $GRADLE_TARGET $GRADLE_PARAMS -PskipTests -PrunBuildSrcTests -PskipSpotless -PtaskPartitionCount=$NORMALIZED_NODE_TOTAL -PtaskPartition=$NORMALIZED_NODE_INDEX $GRADLE_ARGS
+    - ./gradlew $GRADLE_TARGET $GRADLE_PARAMS -PskipTests -PrunBuildSrcTests -PskipSpotless -Pslot=$CI_NODE_INDEX/$CI_NODE_TOTAL $GRADLE_ARGS
   after_script:
     - *cgroup_info
     - source .gitlab/gitlab-utils.sh
@@ -488,7 +485,7 @@ muzzle:
   script:
     - export SKIP_BUILDSCAN="true"
     - ./gradlew --version
-    - ./gradlew :runMuzzle -PtaskPartitionCount=$NORMALIZED_NODE_TOTAL -PtaskPartition=$NORMALIZED_NODE_INDEX $GRADLE_ARGS
+    - ./gradlew :runMuzzle -Pslot=$CI_NODE_INDEX/$CI_NODE_TOTAL $GRADLE_ARGS
   after_script:
     - *cgroup_info
     - source .gitlab/gitlab-utils.sh
@@ -570,7 +567,7 @@ muzzle-dep-report:
     - *prepare_test_env
     - export GRADLE_OPTS="-Dorg.gradle.jvmargs='-Xms$GRADLE_MEM -Xmx$GRADLE_MEM $PROFILER_COMMAND -XX:ErrorFile=/tmp/hs_err_pid%p.log -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/tmp' -Ddatadog.forkedMaxHeapSize=1024M -Ddatadog.forkedMinHeapSize=128M"
     - ./gradlew --version
-    - ./gradlew $GRADLE_TARGET $GRADLE_PARAMS -PtestJvm=$testJvm -PtaskPartitionCount=$NORMALIZED_NODE_TOTAL -PtaskPartition=$NORMALIZED_NODE_INDEX $GRADLE_ARGS --continue || $CONTINUE_ON_FAILURE
+    - ./gradlew $GRADLE_TARGET $GRADLE_PARAMS -PtestJvm=$testJvm -Pslot=$CI_NODE_INDEX/$CI_NODE_TOTAL $GRADLE_ARGS --continue || $CONTINUE_ON_FAILURE
   after_script:
     - *restore_pretest_env
     - *set_datadog_api_keys

buildSrc/src/main/kotlin/datadog/gradle/plugin/ci/CIJobsExtensions.kt

@@ -2,7 +2,54 @@ package datadog.gradle.plugin.ci
 
 import org.gradle.api.Project
 import org.gradle.api.Task
+import org.gradle.api.provider.Provider
 import org.gradle.kotlin.dsl.extra
+import kotlin.math.abs
+
+/**
+ * Determines if the current project is in the selected slot.
+ *
+ * The "slot" property should be provided in the format "X/Y", where X is the selected slot (1-based)
+ * and Y is the total number of slots.
+ *
+ * If the "slot" property is not provided, all projects are considered to be in the selected slot.
+ */
+val Project.isInSelectedSlot: Provider<Boolean>
+  get() = rootProject.providers.gradleProperty("slot").map { slot ->
+    val parts = slot.split("/")
+    if (parts.size != 2) {
+      project.logger.warn("Invalid slot format '{}', expected 'X/Y'. Treating all projects as selected.", slot)
+      return@map true
+    }
+
+    val selectedSlot = parts[0].toIntOrNull()
+    val totalSlots = parts[1].toIntOrNull()
+
+    if (selectedSlot == null || totalSlots == null || totalSlots <= 0) {
+      project.logger.warn("Invalid slot values '{}', expected numeric 'X/Y' with Y > 0. Treating all projects as selected.", slot)
+      return@map true
+    }
+
+    // Distribution numbers when running on rootProject.allprojects indicates
+    // bucket sizes are reasonably balanced:
+    //
+    // * size  4 distribution: {2=146, 0=143, 1=157, 3=145}
+    // * size  6 distribution: {4=100, 0=92, 3=97, 2=97, 1=108, 5=97}
+    // * size  8 distribution: {2=62, 4=72, 0=71, 5=70, 7=78, 6=84, 1=87, 3=67}
+    // * size 10 distribution: {8=62, 0=65, 5=70, 9=59, 3=54, 1=56, 6=63, 4=47, 2=52, 7=63}
+    // * size 12 distribution: {10=55, 0=47, 4=45, 9=46, 8=51, 3=51, 2=46, 1=59, 5=52, 7=49, 11=45, 6=45}
+    val projectSlot = abs(project.path.hashCode() % totalSlots) + 1 // Convert to 1-based
+
+    project.logger.info(
+      "Project {} assigned to slot {}/{}, active slot is {}",
+      project.path,
+      projectSlot,
+      totalSlots,
+      selectedSlot,
+    )
+
+    projectSlot == selectedSlot
+  }.orElse(true)
 
 /**
  * Returns the task's path, given affected projects, if this task or its dependencies are affected by git changes.
@@ -46,9 +93,8 @@ private fun Project.createRootTask(
   val coverage = forceCoverage || rootProject.providers.gradleProperty("checkCoverage").isPresent
   tasks.register(rootTaskName) {
     subprojects.forEach { subproject ->
-      val activePartition = subproject.extra.get("activePartition") as Boolean
       if (
-        activePartition &&
+        isInSelectedSlot.get() &&
         includePrefixes.any { subproject.path.startsWith(it) } &&
         !excludePrefixes.any { subproject.path.startsWith(it) }
       ) {

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

@@ -1,33 +1,25 @@
+import datadog.gradle.plugin.ci.isInSelectedSlot
+import org.gradle.api.tasks.testing.Test
+import java.io.File
+
 /*
- * This plugin defines a set of tasks to be used in CI. These aggregate tasks support partitioning (to parallelize
- * jobs) with -PtaskPartitionCount and -PtaskPartition, and limiting tasks to those affected by git changes
- * with -PgitBaseRef.
+ * This plugin defines a set of tasks to be used in CI.
+ *
+ * These aggregate tasks support partitioning (to parallelize jobs) with
+ * `-Pslot=x/y`, and limiting tasks to those affected by git changes with
+ * `-PgitBaseRef`.
  */
 
-import org.gradle.api.tasks.testing.Test
-import java.io.File
-import kotlin.math.abs
+if (project != rootProject) {
+  logger.error("This plugin has been applied on a non-root project: ${project.path}")
+}
 
-// Set up activePartition property on all projects
 allprojects {
-  extra.set("activePartition", true)
-
-  val taskPartitionCountProvider = rootProject.providers.gradleProperty("taskPartitionCount")
-  val taskPartitionProvider = rootProject.providers.gradleProperty("taskPartition")
-  if (taskPartitionCountProvider.isPresent && taskPartitionProvider.isPresent) {
-    val taskPartitionCount = taskPartitionCountProvider.get()
-    val taskPartition = taskPartitionProvider.get()
-    val currentTaskPartition = abs(project.path.hashCode() % taskPartitionCount.toInt())
-    extra.set("activePartition", currentTaskPartition == taskPartition.toInt())
-  }
-
-  // Disable test tasks if not in active partition
-  val activePartitionProvider = providers.provider {
-    project.extra.properties["activePartition"] as? Boolean ?: true
-  }
-
+  // Enable tests only on the selected slot (if -Pslot=n/t is provided)
   tasks.withType<Test>().configureEach {
-    enabled = activePartitionProvider.get()
+    onlyIf("Project is in selected slot") {
+      project.isInSelectedSlot.get()
+    }
   }
 }
 
@@ -132,8 +124,9 @@ if (gitBaseRefProvider.isPresent) {
 
 tasks.register("runMuzzle") {
   val muzzleSubprojects = subprojects.filter { p ->
-    val activePartition = p.extra.get("activePartition") as Boolean
-    activePartition && p.plugins.hasPlugin("java") && p.plugins.hasPlugin("dd-trace-java.muzzle")
+    p.isInSelectedSlot.get()
+        && p.plugins.hasPlugin("java")
+        && p.plugins.hasPlugin("dd-trace-java.muzzle")
   }
   dependsOn(muzzleSubprojects.map { p -> "${p.path}:muzzle" })
 }