Addressed PR comments

Author: rboustany-crypto

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

@@ -206,7 +206,7 @@ internal class NextgenEngine(
                                 try {
                                     val resolvedService = fieldToService.resolveDynamicService(fields, service)
                                     executeTopLevelField(
-                                        topLevelField = fields,
+                                        topLevelFields = fields,
                                         service = resolvedService,
                                         executionContext = executionContext,
                                     )
@@ -221,7 +221,7 @@ internal class NextgenEngine(
                 }.awaitAll()
 
                 if (executionHints.newResultMergerAndNamespacedTypename()) {
-                    NadelResultMerger.mergeResults(fields, engineSchema, results)
+                    NadelResultMerger.mergeResults(operation.topLevelFields, engineSchema, results)
                 } else {
                     graphql.nadel.engine.util.mergeResults(results)
                 }
@@ -260,7 +260,7 @@ internal class NextgenEngine(
     ): ServiceExecutionResult {
         return try {
             executeTopLevelField(
-                topLevelField = listOf(topLevelField),
+                topLevelFields = listOf(topLevelField),
                 service = service,
                 executionContext = executionContext.copy(
                     hydrationDetails = hydrationDetails,
@@ -283,7 +283,7 @@ internal class NextgenEngine(
         executionContext: NadelExecutionContext,
     ): ServiceExecutionResult {
         return executeTopLevelField(
-            topLevelField = listOf(topLevelField),
+            topLevelFields = listOf(topLevelField),
             service = service,
             executionContext = executionContext.copy(
                 isPartitionedCall = true,
@@ -292,7 +292,7 @@ internal class NextgenEngine(
     }
 
     private suspend fun executeTopLevelField(
-        topLevelField: List<ExecutableNormalizedField>,
+        topLevelFields: List<ExecutableNormalizedField>,
         service: Service,
         executionContext: NadelExecutionContext,
     ): ServiceExecutionResult {
@@ -305,7 +305,7 @@ internal class NextgenEngine(
                 serviceExecutionContext = serviceExecutionContext,
                 services = services,
                 service = service,
-                rootField = topLevelField,
+                rootFields = topLevelFields,
                 serviceHydrationDetails = executionContext.hydrationDetails,
             )
         }
@@ -315,7 +315,7 @@ internal class NextgenEngine(
                 executionContext = executionContext,
                 serviceExecutionContext = serviceExecutionContext,
                 executionPlan = executionPlan,
-                fields = topLevelField
+                fields = topLevelFields
             )
         }
         val result: ServiceExecutionResult = timer.time(step = RootStep.ServiceExecution.child(service.name)) {
@@ -354,7 +354,7 @@ internal class NextgenEngine(
         val transformedResult: ServiceExecutionResult = when {
             // Introspection fields are never batched with other fields (see NadelFieldToService),
             // so an all-introspection batch needs no result transformation.
-            topLevelField.all { it.name.startsWith("__") } -> result
+            topLevelFields.all { it.name.startsWith("__") } -> result
             else -> timer.time(step = RootStep.ResultTransforming) {
                 resultTransformer.transform(
                     executionContext = executionContext,

lib/src/main/java/graphql/nadel/engine/plan/NadelExecutionPlanFactory.kt

@@ -53,24 +53,22 @@ internal class NadelExecutionPlanFactory(
 
     /**
      * This derives an execution plan from with the main input parameters being the
-     * [rootField] and [executionBlueprint].
+     * [rootFields] and [executionBlueprint].
      */
     suspend fun create(
         executionContext: NadelExecutionContext,
         serviceExecutionContext: NadelServiceExecutionContext,
         services: Map<String, Service>,
         service: Service,
-        rootField: List<ExecutableNormalizedField>,
+        rootFields: List<ExecutableNormalizedField>,
         serviceHydrationDetails: ServiceExecutionHydrationDetails?,
     ): NadelExecutionPlan {
         val executionSteps: MutableMap<ExecutableNormalizedField, List<NadelExecutionPlan.Step<Any>>> =
             mutableMapOf()
         val transformContexts: MutableMap<NadelTransform<Any>, NadelTransformServiceExecutionContext?> =
             mutableMapOf()
         executionContext.timer.batch { timer ->
-            // Plan across every root field: batched root fields (see NadelBatchRootFieldsHint) share
-            // one service call, so all their transform steps belong to this single execution plan.
-            traverseQuery(rootField) { field ->
+            traverseQuery(rootFields) { field ->
                 val steps = transformsWithTimingStepInfo.mapNotNull { transformWithTimingInfo ->
                     val transform = transformWithTimingInfo.transform
                     // This is a patch to prevent errors
@@ -89,9 +87,7 @@ internal class NadelExecutionPlanFactory(
                                     executionBlueprint,
                                     services,
                                     service,
-                                    // buildContext is memoized once per service call; the first root
-                                    // field is passed as the representative root field of the batch.
-                                    rootField.first(),
+                                    rootFields,
                                     serviceHydrationDetails
                                 )
                             }

lib/src/main/java/graphql/nadel/engine/transform/NadelTransform.kt

@@ -25,7 +25,8 @@ interface NadelTransform<State : Any> {
      * of the transform on all the fields.
      *
      * @param executionBlueprint the [NadelOverallExecutionBlueprint] of the Nadel instance being operated on
-     * @param rootField the root [ExecutableNormalizedField] of the operation this [NadelTransform] runs on
+     * @param rootFields the root [ExecutableNormalizedField]s of the operation this [NadelTransform] runs on.
+     * There is more than one when sibling root fields are batched into a single service call.
      * @param hydrationDetails the [ServiceExecutionHydrationDetails] when the [NadelTransform] is applied to fields inside
      * hydrations, `null` otherwise
      *
@@ -39,7 +40,7 @@ interface NadelTransform<State : Any> {
         executionBlueprint: NadelOverallExecutionBlueprint,
         services: Map<String, Service>,
         service: Service,
-        rootField: ExecutableNormalizedField,
+        rootFields: List<ExecutableNormalizedField>,
         hydrationDetails: ServiceExecutionHydrationDetails?,
     ): NadelTransformServiceExecutionContext? {
         return null

lib/src/main/java/graphql/nadel/engine/transform/NadelTransformJavaCompat.kt

@@ -113,7 +113,7 @@ interface NadelTransformJavaCompat<State : Any> {
                     executionBlueprint: NadelOverallExecutionBlueprint,
                     services: Map<String, Service>,
                     service: Service,
-                    rootField: ExecutableNormalizedField,
+                    rootFields: List<ExecutableNormalizedField>,
                     hydrationDetails: ServiceExecutionHydrationDetails?,
                 ): NadelTransformServiceExecutionContext? {
                     return compat.buildContext(
@@ -122,7 +122,8 @@ interface NadelTransformJavaCompat<State : Any> {
                         executionBlueprint,
                         services,
                         service,
-                        rootField,
+                        // The Java compat API predates root-field batching and takes a single root field.
+                        rootFields.first(),
                         hydrationDetails
                     ).asDeferred().await()
                 }

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

@@ -27,45 +27,39 @@ internal class NadelFieldToService(
         query: ExecutableNormalizedOperation,
         executionHints: NadelExecutionHints,
     ): List<NadelFieldAndService> {
-        val topLevelFields = query.topLevelFields
-
-        // Resolve the owning service for each non-namespaced root field, and bucket the ones that
-        // are eligible for batching by service. LinkedHashMap preserves first-occurrence order.
-        val fieldToService = HashMap<ExecutableNormalizedField, Service>()
-        val batchedByService = LinkedHashMap<Service, MutableList<ExecutableNormalizedField>>()
-        for (topLevelField in topLevelFields) {
-            if (isNamespacedField(topLevelField)) {
-                continue
-            }
-            val service = getService(topLevelField)
-            fieldToService[topLevelField] = service
-            if (canBatchRootField(topLevelField, service, executionHints)) {
-                batchedByService.getOrPut(service) { mutableListOf() }.add(topLevelField)
+        // 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()) {
+            return query.topLevelFields.flatMap { topLevelField ->
+                if (isNamespacedField(topLevelField)) {
+                    getServicePairsForNamespacedFields(topLevelField, executionHints)
+                } else {
+                    listOf(NadelFieldAndService(field = listOf(topLevelField), service = getService(topLevelField)))
+                }
             }
         }
 
-        // Emit entries in query order. A batched service's root fields are emitted as a single
-        // entry (one service call) at the position of the service's first batchable root field;
-        // everything else keeps its own single-field entry, exactly as before.
+        // Group batch-eligible root fields per service into a single entry (one service call).
         val result = mutableListOf<NadelFieldAndService>()
-        val emittedBatches = HashSet<Service>()
-        for (topLevelField in topLevelFields) {
+        val batchedByService = LinkedHashMap<Service, MutableList<ExecutableNormalizedField>>()
+        for (topLevelField in query.topLevelFields) {
             if (isNamespacedField(topLevelField)) {
                 result += getServicePairsForNamespacedFields(topLevelField, executionHints)
                 continue
             }
 
-            val service = fieldToService.getValue(topLevelField)
-            val batchedFields = batchedByService[service]
-            if (batchedFields != null && topLevelField in batchedFields) {
-                if (emittedBatches.add(service)) {
-                    result += NadelFieldAndService(field = batchedFields, service = service)
-                }
+            val service = getService(topLevelField)
+            if (canBatchRootField(topLevelField, service, executionHints)) {
+                batchedByService.getOrPut(service) { mutableListOf() }.add(topLevelField)
             } else {
                 result += NadelFieldAndService(field = listOf(topLevelField), service = service)
             }
         }
 
+        batchedByService.forEach { (service, batchedFields) ->
+            result += NadelFieldAndService(field = batchedFields, service = service)
+        }
+
         return result
     }
 

lib/src/main/java/graphql/nadel/hints/NadelBatchRootFieldsHint.kt

@@ -4,14 +4,19 @@ import graphql.nadel.Service
 
 fun interface NadelBatchRootFieldsHint {
     /**
-     * Determines whether multiple top level (root) fields destined for the same service should be
-     * combined into a single service call (i.e. one request), rather than one call per root field.
+     * Global feature flag for root-field batching. When false, Nadel keeps its original behaviour
+     * of making one service call per root field. Use this to feature flag the whole feature.
+     */
+    operator fun invoke(): Boolean
+
+    /**
+     * Per-service opt-in for root-field batching. Defaults to the global flag [invoke].
      *
-     * This only affects sibling root fields with no shared namespace wrapper (e.g. the prefixed
+     * Only affects sibling root fields with no shared namespace wrapper (e.g. the prefixed
      * `jira_foo`, `jira_bar` form). Namespaced fields are already batched per service regardless.
      *
      * @param service the service the root fields would be sent to
      * @return true to batch this service's root fields into a single call
      */
-    operator fun invoke(service: Service): Boolean
+    operator fun invoke(service: Service): Boolean = invoke()
 }

lib/src/main/java/graphql/nadel/result/NadelResultMerger.kt

@@ -5,7 +5,6 @@ import graphql.ExecutionResultImpl
 import graphql.GraphQLError
 import graphql.introspection.Introspection
 import graphql.nadel.ServiceExecutionResult
-import graphql.nadel.engine.transform.query.NadelFieldAndService
 import graphql.nadel.engine.transform.result.NadelResultKey
 import graphql.nadel.engine.util.AnyMap
 import graphql.nadel.engine.util.JsonMap
@@ -20,7 +19,7 @@ import graphql.schema.GraphQLSchema
 
 internal object NadelResultMerger {
     fun mergeResults(
-        fields: List<NadelFieldAndService>,
+        topLevelFields: List<ExecutableNormalizedField>,
         engineSchema: GraphQLSchema,
         results: List<ServiceExecutionResult>,
     ): ExecutionResult {
@@ -35,7 +34,7 @@ internal object NadelResultMerger {
         }
 
         return ExecutionResultImpl.newExecutionResult()
-            .data(fixData(fields, engineSchema, data))
+            .data(fixData(topLevelFields, engineSchema, data))
             .extensions(extensions.let {
                 @Suppress("UNCHECKED_CAST") // .extensions should take in a Map<*, *> instead of strictly Map<Any?, Any?>
                 it as Map<Any?, Any?>
@@ -47,17 +46,15 @@ internal object NadelResultMerger {
     }
 
     private fun fixData(
-        fields: List<NadelFieldAndService>,
+        topLevelFields: List<ExecutableNormalizedField>,
         engineSchema: GraphQLSchema,
         data: MutableJsonMap,
     ): MutableJsonMap? {
-        val requiredFieldMap = buildRequiredFieldMap(fields, engineSchema)
+        val requiredFieldMap = buildRequiredFieldMap(topLevelFields, engineSchema)
 
         for ((topLevelResultKey, children) in requiredFieldMap) {
             val topLevelFieldDef by lazy {
-                fields
-                    .asSequence()
-                    .flatMap { it.field }
+                topLevelFields
                     .first { it.resultKey == topLevelResultKey.value }
                     .getFieldDefinitions(engineSchema)
                     .single() // This is under Query, Mutation etc. so there is only one field
@@ -116,14 +113,14 @@ internal object NadelResultMerger {
     }
 
     private fun buildRequiredFieldMap(
-        fields: List<NadelFieldAndService>,
+        topLevelFields: List<ExecutableNormalizedField>,
         engineSchema: GraphQLSchema,
     ): MutableMap<NadelResultKey, MutableList<ExecutableNormalizedField>> {
         val requiredFields = mutableMapOf<NadelResultKey, MutableList<ExecutableNormalizedField>>()
 
         // NOTE: please ensure all fields are from object types and will NOT have multiple field defs
         // Other code in this file relies on this contract
-        for (field in fields.flatMap { it.field }) {
+        for (field in topLevelFields) {
             val requiredChildFields = requiredFields
                 .computeIfAbsent(NadelResultKey(field.resultKey)) {
                     mutableListOf()

lib/src/test/kotlin/graphql/nadel/transform/NadelTransformJavaCompatTest.kt

@@ -86,7 +86,7 @@ class NadelTransformJavaCompatTest : DescribeSpec({
                 executionBlueprint,
                 services,
                 service,
-                overallField,
+                listOf(overallField),
                 hydrationDetails = hydrationDetails,
             )
 

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

@@ -1,6 +1,7 @@
 package graphql.nadel.tests.next.fixtures.batching
 
 import graphql.nadel.NadelExecutionHints
+import graphql.nadel.hints.NadelBatchRootFieldsHint
 import graphql.nadel.tests.next.NadelIntegrationTest
 
 /**
@@ -58,6 +59,9 @@ class BatchRootFieldsPerServiceTest : NadelIntegrationTest(
 ) {
     override fun makeExecutionHints(): NadelExecutionHints.Builder {
         return super.makeExecutionHints()
-            .batchRootFields { service -> service.name == "batched" }
+            .batchRootFields(object : NadelBatchRootFieldsHint {
+                override fun invoke(): Boolean = true
+                override fun invoke(service: graphql.nadel.Service): Boolean = service.name == "batched"
+            })
     }
 }

test/src/test/kotlin/graphql/nadel/tests/next/fixtures/execution/ServiceExecutionContextTest.kt

@@ -158,11 +158,11 @@ class ServiceExecutionContextTest : NadelIntegrationTest(
                             executionBlueprint: NadelOverallExecutionBlueprint,
                             services: Map<String, graphql.nadel.Service>,
                             service: graphql.nadel.Service,
-                            rootField: ExecutableNormalizedField,
+                            rootFields: List<ExecutableNormalizedField>,
                             hydrationDetails: ServiceExecutionHydrationDetails?,
                         ): NadelTransformServiceExecutionContext? {
                             val testTransformServiceExecutionContext =
-                                TestTransformServiceExecutionContext(rootField.toExecutionString())
+                                TestTransformServiceExecutionContext(rootFields.single().toExecutionString())
                             transformServiceExecutionContexts.add(testTransformServiceExecutionContext)
                             return testTransformServiceExecutionContext
                         }