Batching based on sharding target

lib/src/main/java/graphql/nadel/NextgenEngine.kt

@@ -198,7 +198,7 @@ internal class NextgenEngine(
             )
 
             val result: ExecutionResult = try {
-                val fields = fieldToService.getServicesForTopLevelFields(operation, executionHints)
+                val fields = fieldToService.getServicesForTopLevelFields(executionContext)
                 val results = coroutineScope {
                     fields
                         .map { (fields, service) ->

lib/src/main/java/graphql/nadel/engine/transform/query/NadelFieldToService.kt

@@ -3,6 +3,7 @@ package graphql.nadel.engine.transform.query
 import graphql.introspection.Introspection
 import graphql.nadel.NadelExecutionHints
 import graphql.nadel.Service
+import graphql.nadel.engine.NadelExecutionContext
 import graphql.nadel.engine.blueprint.IntrospectionService
 import graphql.nadel.engine.blueprint.NadelIntrospectionRunnerFactory
 import graphql.nadel.engine.blueprint.NadelOverallExecutionBlueprint
@@ -24,9 +25,11 @@ internal class NadelFieldToService(
     private val introspectionService = IntrospectionService(querySchema, introspectionRunnerFactory)
 
     fun getServicesForTopLevelFields(
-        query: ExecutableNormalizedOperation,
-        executionHints: NadelExecutionHints,
+        executionContext: NadelExecutionContext,
     ): List<NadelFieldAndService> {
+        val query = executionContext.query
+        val executionHints = executionContext.hints
+
         // Feature flag: when root-field batching is globally disabled, keep the original behaviour
         // of one entry (i.e. one service call) per root field.
         if (!executionHints.batchRootFields()) {
@@ -39,9 +42,12 @@ internal class NadelFieldToService(
             }
         }
 
-        // Group batch-eligible root fields per service into a single entry (one service call).
+        // Group batch-eligible root fields into a single entry (one service call) per (service, shard).
+        // Two root fields can be owned by the same service but routed to different shards (e.g. different
+        // cloud IDs), so they must not be batched together. Nadel stays generic here: the sharding target
+        // comes from NadelExecutionHooks.getShardingTarget and is used purely as an opaque grouping key.
         val result = mutableListOf<NadelFieldAndService>()
-        val batchedByService = LinkedHashMap<Service, MutableList<ExecutableNormalizedField>>()
+        val batchedByGroup = LinkedHashMap<NadelBatchGroup, MutableList<ExecutableNormalizedField>>()
         for (topLevelField in query.topLevelFields) {
             if (isNamespacedField(topLevelField)) {
                 result += getServicePairsForNamespacedFields(topLevelField, executionHints)
@@ -50,14 +56,15 @@ internal class NadelFieldToService(
 
             val service = getService(topLevelField)
             if (canBatchRootField(topLevelField, service, executionHints)) {
-                batchedByService.getOrPut(service) { mutableListOf() }.add(topLevelField)
+                val shardingTarget = executionContext.hooks.getShardingTarget(executionContext, service, topLevelField)
+                batchedByGroup.getOrPut(NadelBatchGroup(service, shardingTarget)) { mutableListOf() }.add(topLevelField)
             } else {
                 result += NadelFieldAndService(fields = listOf(topLevelField), service = service)
             }
         }
 
-        batchedByService.forEach { (service, batchedFields) ->
-            result += NadelFieldAndService(fields = batchedFields, service = service)
+        batchedByGroup.forEach { (group, batchedFields) ->
+            result += NadelFieldAndService(fields = batchedFields, service = group.service)
         }
 
         return result
@@ -155,3 +162,14 @@ data class NadelFieldAndService(
     val service: Service,
 )
 
+/**
+ * Grouping key for root-field batching: root fields are only batched together when they share the
+ * same [service] and the same [shardingTarget]. [shardingTarget] is an opaque value supplied by
+ * [graphql.nadel.hooks.NadelExecutionHooks.getShardingTarget] (`null` when the field is not bound to
+ * a specific shard).
+ */
+private data class NadelBatchGroup(
+    val service: Service,
+    val shardingTarget: Any?,
+)
+

lib/src/main/java/graphql/nadel/hooks/NadelExecutionHooks.kt

@@ -44,6 +44,29 @@ interface NadelExecutionHooks {
         return null
     }
 
+    /**
+     * Returns an opaque key identifying the target "shard" that a top level (root) [field] will be
+     * routed to. It is used to decide which root fields may be batched into a single service call.
+     *
+     * Root fields are only batched together when they share the same [service] AND the same sharding
+     * target. Two fields can belong to the same [service] but be routed to different shards (e.g.
+     * different cloud IDs), in which case they cannot be sent in the same request.
+     *
+     * Nadel treats the returned value as opaque - it groups by it via [equals]/[hashCode] and never
+     * interprets it - so Atlassian-specific concerns such as cloud ID / ARI parsing stay entirely in
+     * the implementation.
+     *
+     * Return `null` when the field is not bound to a specific shard; such fields are grouped together
+     * (per service) separately from shard-specific ones.
+     */
+    fun getShardingTarget(
+        executionContext: NadelExecutionContext,
+        service: Service,
+        field: ExecutableNormalizedField,
+    ): Any? {
+        return null
+    }
+
     fun <T : NadelGenericHydrationInstruction> getHydrationInstruction(
         virtualField: ExecutableNormalizedField,
         instructions: List<T>,

test/src/test/kotlin/graphql/nadel/tests/next/fixtures/batching/BatchRootFieldsByShardTest.kt

@@ -0,0 +1,70 @@
+package graphql.nadel.tests.next.fixtures.batching
+
+import graphql.nadel.Nadel
+import graphql.nadel.NadelExecutionHints
+import graphql.nadel.engine.NadelExecutionContext
+import graphql.nadel.hooks.NadelExecutionHooks
+import graphql.nadel.tests.next.NadelIntegrationTest
+import graphql.normalized.ExecutableNormalizedField
+
+/**
+ * Root fields owned by the same service but routed to different shards must NOT be batched together.
+ *
+ * Here [getShardingTarget] returns the field's `cloudId` argument as the (opaque) shard key, so:
+ *  - `a` + `b` (cloudId "site-1") share a shard and batch into a single call
+ *  - `c` (cloudId "site-2") is a different shard and is sent separately
+ *
+ * The snapshot therefore records TWO calls to the "issues" service: one for site-1 (a + b) and one
+ * for site-2 (c).
+ */
+class BatchRootFieldsByShardTest : NadelIntegrationTest(
+    query = """
+        query {
+          a: issueById(cloudId: "site-1") { id }
+          b: issueById(cloudId: "site-1") { id }
+          c: issueById(cloudId: "site-2") { id }
+        }
+    """.trimIndent(),
+    services = listOf(
+        Service(
+            name = "issues",
+            overallSchema = """
+                type Query {
+                  issueById(cloudId: String!): Issue
+                }
+                type Issue {
+                  id: ID
+                }
+            """.trimIndent(),
+            runtimeWiring = { wiring ->
+                wiring
+                    .type("Query") { type ->
+                        type
+                            .dataFetcher("issueById") { env ->
+                                mapOf("id" to env.getArgument<String>("cloudId"))
+                            }
+                    }
+            },
+        ),
+    ),
+) {
+    override fun makeExecutionHints(): NadelExecutionHints.Builder {
+        return super.makeExecutionHints()
+            .batchRootFields { true }
+    }
+
+    override fun makeNadel(): Nadel.Builder {
+        return super.makeNadel()
+            .executionHooks(
+                object : NadelExecutionHooks {
+                    override fun getShardingTarget(
+                        executionContext: NadelExecutionContext,
+                        service: graphql.nadel.Service,
+                        field: ExecutableNormalizedField,
+                    ): Any? {
+                        return field.resolvedArguments["cloudId"]
+                    }
+                },
+            )
+    }
+}

test/src/test/kotlin/graphql/nadel/tests/next/fixtures/batching/BatchRootFieldsByShardTestSnapshot.kt

@@ -0,0 +1,129 @@
+// @formatter:off
+package graphql.nadel.tests.next.fixtures.batching
+
+import graphql.nadel.tests.next.ExpectedNadelResult
+import graphql.nadel.tests.next.ExpectedServiceCall
+import graphql.nadel.tests.next.TestSnapshot
+import graphql.nadel.tests.next.listOfJsonStrings
+import kotlin.Suppress
+import kotlin.collections.List
+import kotlin.collections.listOf
+
+private suspend fun main() {
+    graphql.nadel.tests.next.update<BatchRootFieldsByShardTest>()
+}
+
+/**
+ * This class is generated. Do NOT modify.
+ *
+ * Refer to [graphql.nadel.tests.next.UpdateTestSnapshots]
+ */
+@Suppress("unused")
+public class BatchRootFieldsByShardTestSnapshot : TestSnapshot() {
+    /**
+     * Query
+     *
+     * ```graphql
+     * query {
+     *   a: issueById(cloudId: "site-1") { id }
+     *   b: issueById(cloudId: "site-1") { id }
+     *   c: issueById(cloudId: "site-2") { id }
+     * }
+     * ```
+     *
+     * Variables
+     *
+     * ```json
+     * {}
+     * ```
+     */
+    override val calls: List<ExpectedServiceCall> = listOf(
+            ExpectedServiceCall(
+                service = "issues",
+                query = """
+                | {
+                |   a: issueById(cloudId: "site-1") {
+                |     id
+                |   }
+                |   b: issueById(cloudId: "site-1") {
+                |     id
+                |   }
+                | }
+                """.trimMargin(),
+                variables = "{}",
+                result = """
+                | {
+                |   "data": {
+                |     "a": {
+                |       "id": "site-1"
+                |     },
+                |     "b": {
+                |       "id": "site-1"
+                |     }
+                |   }
+                | }
+                """.trimMargin(),
+                delayedResults = listOfJsonStrings(
+                ),
+            ),
+            ExpectedServiceCall(
+                service = "issues",
+                query = """
+                | {
+                |   c: issueById(cloudId: "site-2") {
+                |     id
+                |   }
+                | }
+                """.trimMargin(),
+                variables = "{}",
+                result = """
+                | {
+                |   "data": {
+                |     "c": {
+                |       "id": "site-2"
+                |     }
+                |   }
+                | }
+                """.trimMargin(),
+                delayedResults = listOfJsonStrings(
+                ),
+            ),
+        )
+
+    /**
+     * ```json
+     * {
+     *   "data": {
+     *     "a": {
+     *       "id": "site-1"
+     *     },
+     *     "b": {
+     *       "id": "site-1"
+     *     },
+     *     "c": {
+     *       "id": "site-2"
+     *     }
+     *   }
+     * }
+     * ```
+     */
+    override val result: ExpectedNadelResult = ExpectedNadelResult(
+            result = """
+            | {
+            |   "data": {
+            |     "a": {
+            |       "id": "site-1"
+            |     },
+            |     "b": {
+            |       "id": "site-1"
+            |     },
+            |     "c": {
+            |       "id": "site-2"
+            |     }
+            |   }
+            | }
+            """.trimMargin(),
+            delayedResults = listOfJsonStrings(
+            ),
+        )
+}