Skip to content

Changes CI test sharding from per-project to per-task granularity. #11847

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
AlexeyKuznetsov-DD wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

Changes CI test sharding from per-project to per-task granularity. #11847

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
921e3a2
Trigger CI baseline (no-op) for per-task sharding A/B
AlexeyKuznetsov-DD Jul 2, 2026
3394c28
Shard CI test slots per-task instead of per-project
AlexeyKuznetsov-DD Jul 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
225 changes: 144 additions & 81 deletions buildSrc/src/main/kotlin/datadog/gradle/plugin/ci/CIJobsExtensions.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,60 +2,95 @@ package datadog.gradle.plugin.ci

import org.gradle.api.Project
import org.gradle.api.Task
import org.gradle.api.provider.Provider
import org.gradle.api.tasks.testing.Test
import org.gradle.kotlin.dsl.extra
import kotlin.math.abs

/** Parsed `-Pslot=X/Y` selection: 1-based [selected] slot out of [total] (> 0). */
private data class SlotSelection(val selected: Int, val total: Int) {
/** Whether [key] falls into this slot. Bucket = `abs(key.hashCode() % total) + 1`, matching Java's truncating `%`. */
fun selects(key: String): Boolean = abs(key.hashCode() % total) + 1 == selected
}

/** Boxes a parsed selection (possibly null) so "absent" can be cached distinctly from "not computed yet". */
private class SlotHolder(val selection: SlotSelection?)

private const val SLOT_HOLDER_KEY = "datadog.ci.slotSelection"
private const val FORCE_COVERAGE_PREFIXES_KEY = "datadog.ci.forceCoveragePrefixes"

/**
* 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.
* The `-Pslot=X/Y` selection, parsed once per build and cached on the root project, or null when the
* property is absent, empty, or malformed — meaning no slot filtering (everything is selected).
*/
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
}
private val Project.slotSelection: SlotSelection?
get() {
val root = rootProject
(root.extra.properties[SLOT_HOLDER_KEY] as? SlotHolder)?.let { return it.selection }
val parsed = parseSlotSelection(root.providers.gradleProperty("slot").orNull, root)
root.extra.set(SLOT_HOLDER_KEY, SlotHolder(parsed))
return parsed
}

// When CI_NODE_INDEX or CI_NODE_TOTAL is unset in non-parallel jobs, one part may be empty (e.g. slot="/1") — treat as no filtering
if (parts[0].isBlank() || parts[1].isBlank()) {
project.logger.info("Incomplete slot value '{}', CI_NODE_INDEX or CI_NODE_TOTAL not set. Treating all projects as selected.", slot)
return@map true
}
private fun parseSlotSelection(raw: String?, root: Project): SlotSelection? {
if (raw == null) return null
val parts = raw.split("/")
if (parts.size != 2) {
root.logger.warn("Invalid slot format '{}', expected 'X/Y'. Treating all as selected.", raw)
return null
}
// When CI_NODE_INDEX or CI_NODE_TOTAL is unset in non-parallel jobs, one part may be empty
// (e.g. slot="/1") — treat as no filtering.
if (parts[0].isBlank() || parts[1].isBlank()) return null
val selected = parts[0].toIntOrNull()
val total = parts[1].toIntOrNull()
if (selected == null || total == null || total <= 0) {
root.logger.warn("Invalid slot values '{}', expected numeric 'X/Y' with Y > 0. Treating all as selected.", raw)
return null
}
return SlotSelection(selected, total)
}

val selectedSlot = parts[0].toIntOrNull()
val totalSlots = parts[1].toIntOrNull()
/**
* Module path prefixes registered (via [testAggregate] with `forceCoverage = true`) that always
* collect coverage. Their test tasks must stay whole-project-slotted even without `-PcheckCoverage`.
*/
@Suppress("UNCHECKED_CAST")
private val Project.forceCoveragePrefixes: List<String>
get() = (rootProject.extra.properties[FORCE_COVERAGE_PREFIXES_KEY] as? List<String>) ?: emptyList()

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
}
/**
* Whether this module collects coverage — via `-PcheckCoverage` or a forceCoverage aggregate. When
* true the module must run whole (project-level slotting) so per-module JaCoCo sees complete
* execution data. [createRootTask] uses the same notion so the two stay consistent.
*/
private val Project.coverageEnabled: Boolean
get() = rootProject.providers.gradleProperty("checkCoverage").isPresent ||
forceCoveragePrefixes.any { path.startsWith(it) }

// 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)
/**
* Whether the current project is in the selected slot, at coarse one-task-per-project granularity.
* Used by whole-project aggregates such as `runMuzzle`.
*/
val Project.isInSelectedSlot: Boolean
get() = slotSelection?.selects(path) ?: true

/**
* Whether this Test task is in the selected slot.
*
* Sharding at the *task* level lets a module's test variants — e.g. jdbc's
* `test`/`forkedTest`/`oldH2Test`/`oldPostgresTest` — spread across different CI slots instead of
* serializing in one job. The key is `"<projectPath>:<taskName>"`.
*
* Exception: when the module collects coverage (see [coverageEnabled]) all its test tasks must stay
* in one slot so per-module JaCoCo sees complete execution data, so we fall back to the project-level
* key. This MUST match the coverage decision in [createRootTask].
*/
val Task.isInSelectedSlot: Boolean
get() {
val slot = project.slotSelection ?: return true
val key = if (project.coverageEnabled) project.path else "${project.path}:$name"
return slot.selects(key)
}

/**
* Returns the task's path, given affected projects, if this task or its dependencies are affected by git changes.
Expand Down Expand Up @@ -87,54 +122,64 @@ internal fun findAffectedTaskPath(baseTask: Task, affectedProjects: Map<Project,
}

/**
* Creates a single aggregate root task that depends on matching subproject tasks
* Creates a single aggregate root task that depends on matching subproject tasks.
*
* When [testTaskFilter] is non-null and coverage is off, the aggregate shards at *task* granularity:
* it depends only on the in-slot Test tasks the umbrella would run, so out-of-slot modules aren't
* pulled into this slot at all (their compile/test work does not run here).
*
* Otherwise — the `check` aggregate, or any coverage build — it shards at *project* granularity.
* `onlyIf` on a Test task only skips that task's own action, never its dependencies or non-Test
* siblings (spotbugs, forbiddenApis, codenarc, jacoco...), and JaCoCo needs a module's complete
* execution data, so the whole module must stay in one slot.
*/
private fun Project.createRootTask(
rootTaskName: String,
subProjTaskName: String,
includePrefixes: List<String>,
excludePrefixes: List<String>,
forceCoverage: Boolean
forceCoverage: Boolean,
testTaskFilter: ((Test) -> Boolean)?
) {
val coverage = forceCoverage || rootProject.providers.gradleProperty("checkCoverage").isPresent
val perTaskShardable = !coverage && testTaskFilter != null
val slot = slotSelection
tasks.register(rootTaskName) {
subprojects.forEach { subproject ->
subprojects.forEach forEachSub@{ subproject ->
if (
subproject.isInSelectedSlot.get() &&
includePrefixes.any { subproject.path.startsWith(it) } &&
!excludePrefixes.any { subproject.path.startsWith(it) }
!includePrefixes.any { subproject.path.startsWith(it) } ||
excludePrefixes.any { subproject.path.startsWith(it) }
) {
val testTask = subproject.tasks.findByName(subProjTaskName)
var isAffected = true

if (testTask != null) {
val useGitChanges = rootProject.extra.get("useGitChanges") as Boolean
if (useGitChanges) {
@Suppress("UNCHECKED_CAST")
val affectedProjects = rootProject.extra.get("affectedProjects") as Map<Project, Set<String>>
val affectedTaskPath = findAffectedTaskPath(testTask, affectedProjects)
if (affectedTaskPath != null) {
logger.warn("Selecting ${subproject.path}:$subProjTaskName (affected by $affectedTaskPath)")
} else {
logger.warn("Skipping ${subproject.path}:$subProjTaskName (not affected by changed files)")
isAffected = false
}
}
if (isAffected) {
dependsOn(testTask)
}
return@forEachSub
}

val subProjTask = subproject.tasks.findByName(subProjTaskName)

// Git-change filtering, keyed off the umbrella task at module granularity (unchanged behavior).
if (subProjTask != null && rootProject.extra.get("useGitChanges") as Boolean) {
@Suppress("UNCHECKED_CAST")
val affectedProjects = rootProject.extra.get("affectedProjects") as Map<Project, Set<String>>
val affectedTaskPath = findAffectedTaskPath(subProjTask, affectedProjects)
if (affectedTaskPath == null) {
logger.warn("Skipping ${subproject.path}:$subProjTaskName (not affected by changed files)")
return@forEachSub
}
logger.warn("Selecting ${subproject.path}:$subProjTaskName (affected by $affectedTaskPath)")
}

if (isAffected && coverage) {
val coverageTask = subproject.tasks.findByName("jacocoTestReport")
if (coverageTask != null) {
dependsOn(coverageTask)
}
val verificationTask = subproject.tasks.findByName("jacocoTestCoverageVerification")
if (verificationTask != null) {
dependsOn(verificationTask)
if (perTaskShardable) {
// Depend only on the in-slot Test tasks the umbrella would run; leave the rest to other slots.
subproject.tasks.withType(Test::class.java).matching { testTaskFilter!!(it) }.forEach { testTask ->
if (slot == null || slot.selects("${subproject.path}:${testTask.name}")) {
dependsOn(testTask)
}
}
} else if (subProjTask != null && (slot == null || slot.selects(subproject.path))) {
dependsOn(subProjTask)
if (coverage) {
subproject.tasks.findByName("jacocoTestReport")?.let { dependsOn(it) }
subproject.tasks.findByName("jacocoTestCoverageVerification")?.let { dependsOn(it) }
}
}
}
}
Expand All @@ -154,7 +199,25 @@ fun Project.testAggregate(
excludePrefixes: List<String> = emptyList(),
forceCoverage: Boolean = false
) {
createRootTask("${baseTaskName}Test", "allTests", includePrefixes, excludePrefixes, forceCoverage)
createRootTask("${baseTaskName}LatestDepTest", "allLatestDepTests", includePrefixes, excludePrefixes, forceCoverage)
createRootTask("${baseTaskName}Check", "check", includePrefixes, excludePrefixes, forceCoverage)
if (forceCoverage) {
registerForceCoveragePrefixes(includePrefixes)
}
// The two Test umbrellas mirror the membership filters in dd-trace-java.configure-tests.gradle.kts.
createRootTask("${baseTaskName}Test", "allTests", includePrefixes, excludePrefixes, forceCoverage) {
!it.name.contains("latest", ignoreCase = true) && it.name != "traceAgentTest"
}
createRootTask("${baseTaskName}LatestDepTest", "allLatestDepTests", includePrefixes, excludePrefixes, forceCoverage) {
it.name.contains("latest", ignoreCase = true)
}
createRootTask("${baseTaskName}Check", "check", includePrefixes, excludePrefixes, forceCoverage, testTaskFilter = null)
}

@Suppress("UNCHECKED_CAST")
private fun Project.registerForceCoveragePrefixes(prefixes: List<String>) {
val existing = rootProject.extra.properties[FORCE_COVERAGE_PREFIXES_KEY] as? MutableList<String>
if (existing != null) {
existing.addAll(prefixes)
} else {
rootProject.extra.set(FORCE_COVERAGE_PREFIXES_KEY, prefixes.toMutableList())
}
}
10 changes: 6 additions & 4 deletions buildSrc/src/main/kotlin/dd-trace-java.ci-jobs.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,12 @@ if (project != rootProject) {
}

allprojects {
// Enable tests only on the selected slot (if -Pslot=n/t is provided)
// Enable tests only on the selected slot (if -Pslot=n/t is provided).
// Sharding is per-task: a module's test variants (e.g. jdbc's test/forkedTest/oldH2Test/
// oldPostgresTest) hash to independent slots instead of serializing within one job.
tasks.withType<Test>().configureEach {
onlyIf("Project is in selected slot") {
project.isInSelectedSlot.get()
onlyIf("Task is in selected slot") {
isInSelectedSlot
}
}
}
Expand Down Expand Up @@ -124,7 +126,7 @@ if (gitBaseRefProvider.isPresent) {

tasks.register("runMuzzle") {
val muzzleSubprojects = subprojects.filter { p ->
p.isInSelectedSlot.get()
p.isInSelectedSlot
&& p.plugins.hasPlugin("java")
&& p.plugins.hasPlugin("dd-trace-java.muzzle")
}
Expand Down
Loading