Balance test_inst CI shards by measured duration instead of project-path hash

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

@@ -4,7 +4,6 @@ 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.
@@ -36,15 +35,9 @@ val Project.isInSelectedSlot: Provider<Boolean>
       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
+    // Slots are balanced by measured test duration rather than project count, see
+    // TestShardBalancer and buildSrc/src/main/resources/datadog/gradle/plugin/ci/test-shard-weights.json.
+    val projectSlot = testShardAssignment(totalSlots).getValue(project.path)
 
     project.logger.info(
       "Project {} assigned to slot {}/{}, active slot is {}",
@@ -57,6 +50,25 @@ val Project.isInSelectedSlot: Provider<Boolean>
     projectSlot == selectedSlot
   }.orElse(true)
 
+/**
+ * Returns the duration-balanced slot assignment of every project for the given number of slots.
+ *
+ * The assignment is computed once per slot count from all project paths, independently of the
+ * selected slot so that every shard agrees on it, and cached on the root project.
+ */
+private fun Project.testShardAssignment(totalSlots: Int): Map<String, Int> {
+  val extra = rootProject.extra
+  val key = "testShardAssignment.$totalSlots"
+  synchronized(extra) {
+    if (!extra.has(key)) {
+      val projectPaths = rootProject.allprojects.map { it.path }.sorted()
+      extra.set(key, TestShardBalancer.assignSlots(projectPaths, totalSlots, TestShardBalancer.loadDefaultWeights()))
+    }
+    @Suppress("UNCHECKED_CAST")
+    return extra.get(key) as Map<String, Int>
+  }
+}
+
 /**
  * Returns the task's path, given affected projects, if this task or its dependencies are affected by git changes.
  */

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

@@ -0,0 +1,68 @@
+package datadog.gradle.plugin.ci
+
+import com.fasterxml.jackson.databind.ObjectMapper
+
+/**
+ * Assigns Gradle projects to CI test shards so that the shards have similar total test duration.
+ *
+ * Kept free of Gradle types so the assignment can be unit tested in isolation.
+ */
+internal object TestShardBalancer {
+  /** Classpath location of the checked-in per-project test duration weights. */
+  const val WEIGHTS_RESOURCE = "/datadog/gradle/plugin/ci/test-shard-weights.json"
+
+  /** Weight, in seconds, of a project that is absent from the weights map. */
+  const val DEFAULT_WEIGHT_SECONDS = 20.6
+
+  /** Parses the `weights` object of a weights file; other top-level keys are ignored. */
+  fun loadWeights(json: String): Map<String, Double> {
+    val weights = ObjectMapper().readTree(json).path("weights")
+    require(weights.isObject) { "Test shard weights file must contain a 'weights' object" }
+    val result = LinkedHashMap<String, Double>()
+    weights.properties().forEach { (path, value) ->
+      require(value.isNumber && value.asDouble() >= 0.0) {
+        "Test shard weight for '$path' must be a non-negative number, got '$value'"
+      }
+      result[path] = value.asDouble()
+    }
+    return result
+  }
+
+  /** Loads the weights file bundled with buildSrc. */
+  fun loadDefaultWeights(): Map<String, Double> {
+    val stream = checkNotNull(TestShardBalancer::class.java.getResourceAsStream(WEIGHTS_RESOURCE)) {
+      "Missing test shard weights resource $WEIGHTS_RESOURCE"
+    }
+    return loadWeights(stream.bufferedReader().use { it.readText() })
+  }
+
+  /** Returns the weight of the given project, falling back to [DEFAULT_WEIGHT_SECONDS]. */
+  fun weightOf(projectPath: String, weights: Map<String, Double>): Double =
+    weights[projectPath] ?: DEFAULT_WEIGHT_SECONDS
+
+  /**
+   * Assigns every project path to a 1-based slot using greedy longest-processing-time-first
+   * bin-packing: projects are placed heaviest first (ties broken by path) into the currently
+   * least loaded slot (ties broken by lowest slot), so the result only depends on the set of
+   * paths and the weights, never on input order or hash codes.
+   */
+  fun assignSlots(projectPaths: List<String>, totalSlots: Int, weights: Map<String, Double>): Map<String, Int> {
+    require(totalSlots > 0) { "totalSlots must be positive, got $totalSlots" }
+    val ordered = projectPaths.distinct().sortedWith(
+      compareByDescending<String> { weightOf(it, weights) }.thenBy { it }
+    )
+    val loads = DoubleArray(totalSlots)
+    val assignment = LinkedHashMap<String, Int>()
+    for (path in ordered) {
+      var slot = 0
+      for (candidate in 1 until totalSlots) {
+        if (loads[candidate] < loads[slot]) {
+          slot = candidate
+        }
+      }
+      loads[slot] += weightOf(path, weights)
+      assignment[path] = slot + 1
+    }
+    return assignment
+  }
+}

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

@@ -1,58 +1,60 @@
 package datadog.gradle.plugin.muzzle
 
 import org.gradle.api.GradleException
-import org.gradle.api.file.FileCollection
 import org.gradle.workers.WorkAction
+import java.io.File
 import java.lang.reflect.Method
 import java.net.URLClassLoader
 
 abstract class MuzzleAction : WorkAction<MuzzleWorkParameters> {
     companion object {
         private val lock = Any()
-        @Volatile
-        private var lastBootCL: ClassLoader? = null
-        @Volatile
-        private var lastToolCL: ClassLoader? = null
-        @Volatile
-        private var lastBuildStamp: Long = 0
-        @Volatile
-        private var lastBuildPathCount: Int = 0
+        private var cachedBuildStamp: Long = 0
+        private val cachedClassLoaders = HashMap<Pair<List<File>, List<File>>, Pair<ClassLoader, ClassLoader>>()
 
         /**
          * Creates a classloader for [cp] with the given [parent]. By default, the platform
          * classloader keeps JDK classes visible without exposing Gradle's bundled libraries.
          */
         fun createClassLoader(
-          cp: FileCollection,
+          cp: Iterable<File>,
           parent: ClassLoader = ClassLoader.getPlatformClassLoader()
         ): ClassLoader {
             val urls = cp.map { it.toURI().toURL() }.toTypedArray()
             return URLClassLoader(urls, parent)
         }
-    }
 
-    override fun execute() {
-        val buildStamp = parameters.buildStartedTime.get()
-        val buildPathCount = parameters.bootstrapClassPath.count() + parameters.toolingClassPath.count()
-        var bootCL : ClassLoader? = lastBootCL
-        var toolCL : ClassLoader? = lastToolCL
-        // cache boot and tool classloaders for each run; rebuild if either class-path is extended mid-build
-        if (bootCL == null || toolCL == null || lastBuildStamp < buildStamp || lastBuildPathCount < buildPathCount) {
-            synchronized(lock) {
-                bootCL = lastBootCL
-                toolCL = lastToolCL
-                if (bootCL == null || toolCL == null || lastBuildStamp < buildStamp || lastBuildPathCount < buildPathCount) {
-                    bootCL = createClassLoader(parameters.bootstrapClassPath)
-                    toolCL = createClassLoader(parameters.toolingClassPath, bootCL!!)
-                    lastBootCL = bootCL
-                    lastToolCL = toolCL
-                    lastBuildStamp = buildStamp
-                    lastBuildPathCount = buildPathCount
-                }
+        /**
+         * Returns the boot and tool classloaders for the given class-paths, cached for the current build.
+         *
+         * Projects can extend the bootstrap class-path (see `muzzleBootstrap`), so classloaders are cached
+         * per distinct class-path: reusing the ones built for another project would hide classes this
+         * project declares, or expose classes it does not declare.
+         */
+        internal fun bootAndToolClassLoaders(
+          buildStamp: Long,
+          bootstrapClassPath: List<File>,
+          toolingClassPath: List<File>
+        ): Pair<ClassLoader, ClassLoader> = synchronized(lock) {
+            if (cachedBuildStamp < buildStamp) {
+                cachedClassLoaders.clear()
+                cachedBuildStamp = buildStamp
+            }
+            cachedClassLoaders.getOrPut(bootstrapClassPath to toolingClassPath) {
+                val bootCL = createClassLoader(bootstrapClassPath)
+                bootCL to createClassLoader(toolingClassPath, bootCL)
             }
         }
-        val instCL = createClassLoader(parameters.instrumentationClassPath, toolCL!!)
-        val testCL = createClassLoader(parameters.testApplicationClassPath, bootCL!!)
+    }
+
+    override fun execute() {
+        val (bootCL, toolCL) = bootAndToolClassLoaders(
+          parameters.buildStartedTime.get(),
+          parameters.bootstrapClassPath.files.toList(),
+          parameters.toolingClassPath.files.toList()
+        )
+        val instCL = createClassLoader(parameters.instrumentationClassPath, toolCL)
+        val testCL = createClassLoader(parameters.testApplicationClassPath, bootCL)
         val assertPass = parameters.assertPass.get()
         val muzzleDirective = parameters.muzzleDirective.orNull
         val assertionMethod: Method = instCL.loadClass("datadog.trace.agent.tooling.muzzle.MuzzleVersionScanPlugin")

buildSrc/src/main/resources/datadog/gradle/plugin/ci/test-shard-weights.json

@@ -0,0 +1,57 @@
+{
+  "_comment": [
+    "Per-project test_inst execution weights used to balance CI test shards.",
+    "Unit: mean seconds of test execution time per pipeline run.",
+    "Source: Datadog CI Visibility, sum(@duration) over test modules grouped by",
+    "Gradle project, JVM 11, job test_inst, 7-day window ending 2026-10-11,",
+    "divided by the 491 observed pipelines.",
+    "Projects absent from this map fall back to the default weight.",
+    "See docs/ci-test-shard-weights.md to regenerate."
+  ],
+  "_metadata": {
+    "generated": "2026-10-12",
+    "window_days": 7,
+    "jvm": "11",
+    "job": "test_inst",
+    "pipelines_observed": 491,
+    "projects_measured": 36
+  },
+  "weights": {
+    ":dd-java-agent:instrumentation:jdbc": 1116.7,
+    ":dd-java-agent:instrumentation:lettuce:lettuce-5.0": 978.1,
+    ":dd-java-agent:instrumentation:akka:akka-http:akka-http-10.0": 612.3,
+    ":dd-java-agent:instrumentation:spring:spring-webflux:spring-webflux-5.0": 509.4,
+    ":dd-java-agent:instrumentation:kafka:kafka-clients-0.11": 379.1,
+    ":dd-java-agent:instrumentation:akka:akka-http:akka-http-10.6": 258.0,
+    ":dd-java-agent:instrumentation:aws-java:aws-java-sdk-2.2": 220.7,
+    ":dd-java-agent:instrumentation:mule-4.5": 217.1,
+    ":dd-java-agent:instrumentation:grpc-1.5": 196.7,
+    ":dd-java-agent:instrumentation:junit:junit-5:junit-5.3": 195.4,
+    ":dd-java-agent:instrumentation:liberty:liberty-20.0": 187.3,
+    ":dd-java-agent:instrumentation:datastax-cassandra:datastax-cassandra-3.8": 175.4,
+    ":dd-java-agent:instrumentation:maven:maven-3.2.1": 172.2,
+    ":dd-java-agent:instrumentation:mongo:mongo-driver:mongo-driver-4.0": 171.3,
+    ":dd-java-agent:instrumentation:couchbase:couchbase-3.2": 165.8,
+    ":dd-java-agent:instrumentation:rabbitmq-amqp-2.7": 159.7,
+    ":dd-java-agent:instrumentation:aws-java:aws-java-sqs-1.0": 152.1,
+    ":dd-java-agent:instrumentation:couchbase:couchbase-3.1": 151.2,
+    ":dd-java-agent:instrumentation:aws-java:aws-java-sqs-2.0": 148.7,
+    ":dd-java-agent:instrumentation:vertx:vertx-web:vertx-web-4.0": 140.9,
+    ":dd-java-agent:instrumentation:spark:spark_2.12": 113.9,
+    ":dd-java-agent:instrumentation:aws-java:aws-java-sdk-1.11": 104.0,
+    ":dd-java-agent:instrumentation:hibernate:hibernate-core-3.3": 93.6,
+    ":dd-java-agent:instrumentation:jersey:jersey-2.0": 84.8,
+    ":dd-java-agent:instrumentation:couchbase:couchbase-2.0": 82.5,
+    ":dd-java-agent:instrumentation:akka:akka-actor-2.5": 66.8,
+    ":dd-java-agent:instrumentation:jetty:jetty-server:jetty-server-9.4.21": 64.6,
+    ":dd-java-agent:instrumentation:lettuce:lettuce-4.0": 47.3,
+    ":dd-java-agent:instrumentation:jedis:jedis-3.0": 42.0,
+    ":dd-java-agent:instrumentation:mongo:mongo-driver:mongo-driver-3:mongo-driver-3.1": 34.7,
+    ":dd-java-agent:instrumentation:elasticsearch:elasticsearch-rest:elasticsearch-rest-7.0": 31.7,
+    ":dd-java-agent:instrumentation:couchbase:couchbase-2.6": 25.6,
+    ":dd-java-agent:instrumentation:synapse-3.0": 24.5,
+    ":dd-java-agent:instrumentation:spring:spring-webmvc:spring-webmvc-5.3": 21.9,
+    ":dd-java-agent:instrumentation:spring:spring-boot-1.3": 16.6,
+    ":dd-java-agent:instrumentation:vertx:vertx-web:vertx-web-3.9": 12.6
+  }
+}