Fix nested object fetching and exception handling in FieldResolver (AIRBNB)

Updates `FieldResolver` to return `Value<Unit>` from `fetchObject`, `fetchObjectSerially`, and `resolveField` instead of returning `Unit` or `Value<FieldResolutionResult>`. This ensures that the execution chain properly waits for all nested object and list fetching to complete before considering the parent operation finished.

Refines exception handling in `FieldResolver` to ensure proper error bubbling:

- `fetchObject` and `fetchObjectSerially` now allow fatal errors (those not swallowed by `executeField`) to bubble up by propagating exceptions caught in `recover`.
- `executeField` continues to encapsulate `FieldFetchingException` and `InternalEngineException` (returning success for partial results).
- `maybeFetchNestedObject` now waits for all list items to be fetched.

Also updates `IViaductInstrumentation` and related classes to change the `beginFetchObject` return type to `InstrumentationContext<Unit>` to match the new `fetchObject` signature.

Github-Change-Id: 967307
GitOrigin-RevId: acc8aca219f0052faf88908de65d2f66e3e09d60

Author: skevy

engine/api/src/main/kotlin/viaduct/engine/api/instrumentation/ChainedModernGJInstrumentation.kt

@@ -11,7 +11,7 @@ open class ChainedModernGJInstrumentation(
     override fun beginFetchObject(
         parameters: InstrumentationExecutionStrategyParameters,
         state: InstrumentationState?
-    ): InstrumentationContext<Map<String, Any?>>? =
+    ): InstrumentationContext<Unit> =
         ChainedInstrumentationContext(
             gjInstrumentations.map { instr ->
                 instr.beginFetchObject(parameters, getState(instr, state))
@@ -21,7 +21,7 @@ open class ChainedModernGJInstrumentation(
     override fun beginCompleteObject(
         parameters: InstrumentationExecutionStrategyParameters,
         state: InstrumentationState?
-    ): InstrumentationContext<Any>? =
+    ): InstrumentationContext<Any> =
         ChainedInstrumentationContext(
             gjInstrumentations.map { instr ->
                 instr.beginCompleteObject(parameters, getState(instr, state))

engine/api/src/main/kotlin/viaduct/engine/api/instrumentation/IViaductInstrumentation.kt

@@ -169,14 +169,14 @@ interface IViaductInstrumentation {
         fun beginFetchObject(
             parameters: InstrumentationExecutionStrategyParameters,
             state: InstrumentationState?
-        ): InstrumentationContext<Map<String, Any?>>?
+        ): InstrumentationContext<Unit>
     }
 
     interface WithBeginCompleteObject : IViaductInstrumentation {
         fun beginCompleteObject(
             parameters: InstrumentationExecutionStrategyParameters,
             state: InstrumentationState?
-        ): InstrumentationContext<Any>?
+        ): InstrumentationContext<Any>
     }
 
     interface WithInstrumentAccessCheck : IViaductInstrumentation {

engine/api/src/main/kotlin/viaduct/engine/api/instrumentation/ViaductInstrumentationAdapter.kt

@@ -155,7 +155,7 @@ open class ViaductInstrumentationAdapter(
     override fun beginFetchObject(
         parameters: InstrumentationExecutionStrategyParameters,
         state: InstrumentationState?
-    ): InstrumentationContext<Map<String, Any?>>? =
+    ): InstrumentationContext<Unit> =
         if (viaductInstrumentation is IViaductInstrumentation.WithBeginFetchObject) {
             viaductInstrumentation.beginFetchObject(parameters, state)
         } else {
@@ -165,7 +165,7 @@ open class ViaductInstrumentationAdapter(
     override fun beginCompleteObject(
         parameters: InstrumentationExecutionStrategyParameters,
         state: InstrumentationState?
-    ): InstrumentationContext<Any>? =
+    ): InstrumentationContext<Any> =
         if (viaductInstrumentation is IViaductInstrumentation.WithBeginCompleteObject) {
             viaductInstrumentation.beginCompleteObject(parameters, state)
         } else {

engine/api/src/main/kotlin/viaduct/engine/api/instrumentation/ViaductModernInstrumentation.kt

@@ -81,7 +81,7 @@ interface ViaductModernInstrumentation {
                 override fun beginFetchObject(
                     parameters: InstrumentationExecutionStrategyParameters,
                     state: InstrumentationState?
-                ): InstrumentationContext<Map<String, Any?>>? {
+                ): InstrumentationContext<Unit> {
                     if (viaductInstrumentation is WithBeginFetchObject) {
                         return viaductInstrumentation.beginFetchObject(parameters, state)
                     }
@@ -111,7 +111,7 @@ interface ViaductModernInstrumentation {
                 override fun beginCompleteObject(
                     parameters: InstrumentationExecutionStrategyParameters,
                     state: InstrumentationState?
-                ): InstrumentationContext<Any>? {
+                ): InstrumentationContext<Any> {
                     if (viaductInstrumentation is WithBeginCompleteObject) {
                         return viaductInstrumentation.beginCompleteObject(parameters, state)
                     }
@@ -283,14 +283,14 @@ interface ViaductModernInstrumentation {
         fun beginFetchObject(
             parameters: InstrumentationExecutionStrategyParameters,
             state: InstrumentationState?
-        ): InstrumentationContext<Map<String, Any?>>?
+        ): InstrumentationContext<Unit>
     }
 
     interface WithBeginFieldExecution : ViaductModernInstrumentation {
         fun beginFieldExecution(
             parameters: InstrumentationFieldParameters,
             state: InstrumentationState?
-        ): InstrumentationContext<Any>?
+        ): InstrumentationContext<Any>
     }
 
     interface WithBeginFieldFetching : ViaductModernInstrumentation {
@@ -304,7 +304,7 @@ interface ViaductModernInstrumentation {
         fun beginCompleteObject(
             parameters: InstrumentationExecutionStrategyParameters,
             state: InstrumentationState?
-        ): InstrumentationContext<Any>?
+        ): InstrumentationContext<Any>
     }
 
     interface WithBeginFieldCompletion : ViaductModernInstrumentation {
@@ -396,14 +396,14 @@ interface ViaductModernGJInstrumentation : Instrumentation {
                 override fun beginFetchObject(
                     parameters: InstrumentationExecutionStrategyParameters,
                     state: InstrumentationState?
-                ): InstrumentationContext<Map<String, Any?>>? {
+                ): InstrumentationContext<Unit> {
                     return noOp()
                 }
 
                 override fun beginCompleteObject(
                     parameters: InstrumentationExecutionStrategyParameters,
                     state: InstrumentationState?
-                ): InstrumentationContext<Any>? {
+                ): InstrumentationContext<Any> {
                     return noOp()
                 }
 
@@ -535,14 +535,14 @@ interface ViaductModernGJInstrumentation : Instrumentation {
     fun beginFetchObject(
         parameters: InstrumentationExecutionStrategyParameters,
         state: InstrumentationState?
-    ): InstrumentationContext<Map<String, Any?>>? {
+    ): InstrumentationContext<Unit> {
         return noOp()
     }
 
     fun beginCompleteObject(
         parameters: InstrumentationExecutionStrategyParameters,
         state: InstrumentationState?
-    ): InstrumentationContext<Any>? {
+    ): InstrumentationContext<Any> {
         return noOp()
     }
 

engine/api/src/test/kotlin/viaduct/engine/api/instrumentation/ChainedModernGJInstrumentationTest.kt

@@ -29,8 +29,8 @@ class ChainedModernGJInstrumentationTest {
         val parameters = mockk<InstrumentationExecutionStrategyParameters>()
         val state1 = object : InstrumentationState {}
         val state2 = object : InstrumentationState {}
-        val context1 = mockk<InstrumentationContext<Map<String, Any?>>>()
-        val context2 = mockk<InstrumentationContext<Map<String, Any?>>>()
+        val context1 = mockk<InstrumentationContext<Unit>>()
+        val context2 = mockk<InstrumentationContext<Unit>>()
 
         val instr1 = mockk<ViaductModernGJInstrumentation> {
             every { createStateAsync(any()) } returns CompletableFuture.completedFuture(state1)

engine/api/src/test/kotlin/viaduct/engine/api/instrumentation/ViaductInstrumentationAdapterTest.kt

@@ -40,15 +40,15 @@ class ViaductInstrumentationAdapterTest {
         override fun beginFetchObject(
             parameters: InstrumentationExecutionStrategyParameters,
             state: InstrumentationState?
-        ): InstrumentationContext<Map<String, Any?>>? {
+        ): InstrumentationContext<Unit> {
             beginFetchObjectCalled = true
             return noOp()
         }
 
         override fun beginCompleteObject(
             parameters: InstrumentationExecutionStrategyParameters,
             state: InstrumentationState?
-        ): InstrumentationContext<Any>? {
+        ): InstrumentationContext<Any> {
             beginCompleteObjectCalled = true
             return noOp()
         }

engine/api/src/testFixtures/kotlin/viaduct/engine/api/mocks/FeatureTest.kt

@@ -3,6 +3,7 @@
 package viaduct.engine.api.mocks
 
 import graphql.ExecutionResult
+import graphql.execution.instrumentation.Instrumentation
 import kotlinx.coroutines.future.await
 import kotlinx.coroutines.runBlocking
 import viaduct.engine.EngineConfiguration
@@ -63,10 +64,11 @@ import viaduct.service.runtime.noderesolvers.ViaductNodeResolverAPIBootstrapper
 fun MockTenantModuleBootstrapper.runFeatureTest(
     withoutDefaultQueryNodeResolvers: Boolean = false,
     schema: ViaductSchema? = null,
+    additionalInstrumentation: Instrumentation? = null,
     block: FeatureTest.() -> Unit
 ) {
     val executableSchema = schema ?: fullSchema
-    val engine = toEngineFactory(withoutDefaultQueryNodeResolvers).create(executableSchema, fullSchema = fullSchema)
+    val engine = toEngineFactory(withoutDefaultQueryNodeResolvers, additionalInstrumentation).create(executableSchema, fullSchema = fullSchema)
     FeatureTest(engine).block()
 }
 
@@ -80,7 +82,10 @@ fun MockTenantModuleBootstrapper.runFeatureTest(
  *
  * and an [EngineConfiguration] constructed with MockFlagManager.Enabled.
  */
-private fun MockTenantModuleBootstrapper.toEngineFactory(withoutDefaultQueryNodeResolvers: Boolean): EngineFactory {
+private fun MockTenantModuleBootstrapper.toEngineFactory(
+    withoutDefaultQueryNodeResolvers: Boolean,
+    additionalInstrumentation: Instrumentation? = null
+): EngineFactory {
     val mods = listOf(this)
     val tenantAPIBootstrapper = buildList {
         add(MockTenantAPIBootstrapperBuilder(MockTenantAPIBootstrapper(mods)))
@@ -101,6 +106,8 @@ private fun MockTenantModuleBootstrapper.toEngineFactory(withoutDefaultQueryNode
     ).create(fullSchema)
     val config = EngineConfiguration.default.copy(
         flagManager = MockFlagManager.Enabled,
+        chainInstrumentationWithDefaults = true,
+        additionalInstrumentation = additionalInstrumentation
     )
     return EngineFactory(config, dispatcherRegistry)
 }

engine/runtime/src/main/kotlin/viaduct/engine/runtime/execution/FieldResolver.kt

@@ -91,13 +91,22 @@ class FieldResolver(
      * 1. Runs CollectFields on the current uncollected selection set
      * 2. Fires off field fetches for each merged object selection, in parallel
      *
+     * Note on return value: This method returns `Value<Unit>` instead of `Value<Map<String, FieldResolutionResult>>`
+     * because the actual resolved values are stored directly in the `ObjectEngineResult` associated with
+     * the parent object. The `Value<Unit>` serves as a completion signal for the orchestration layer to
+     * know when all nested fetching (including any lazy data or nested objects) has finished.
+     *
+     * If the Value returned by this method is exceptionally completed, that means that there has been
+     * a fatal error in resolving this object, and the parent ObjectEngineResult may be incomplete. Thus, callers of this
+     * method should check for exceptional completion and handle it appropriately.
+     *
      * @param parameters ExecutionParameters containing the execution context and selection set
      * @throws Exception Only if there's a fatal error in the supervisorScope itself
      */
     fun fetchObject(
         objectType: GraphQLObjectType,
         parameters: ExecutionParameters
-    ) {
+    ): Value<Unit> {
         val instrumentationParameters =
             InstrumentationExecutionStrategyParameters(parameters.executionContext, parameters.gjParameters)
         val resolveObjectCtx = nonNullCtx(
@@ -110,17 +119,20 @@ class FieldResolver(
         try {
             val results = collectFields(objectType, parameters)
                 .selections
-                .associate { field ->
+                .map { field ->
                     field as QueryPlan.CollectedField
                     val newParams = parameters.forField(objectType, field)
-                    field.responseKey to resolveField(newParams, field)
+                    resolveField(newParams, field)
                 }
 
             // Wait for all values to be completed.
-            // Errors will be folded into the result map
-            Value.waitAll(results.values)
-                .thenApply { _, t ->
-                    resolveObjectCtx.onCompleted(results, t)
+            return Value.waitAll(results)
+                .map {
+                    resolveObjectCtx.onCompleted(Unit, null)
+                    it
+                }.recover { t ->
+                    resolveObjectCtx.onCompleted(null, t)
+                    Value.fromThrowable(t)
                 }
         } catch (e: Exception) {
             resolveObjectCtx.onCompleted(null, e)
@@ -137,13 +149,22 @@ class FieldResolver(
      *   initiated when the previous selection has completed fetching (either
      *   successfully or exceptionally)
      *
+     * Note on return value: This method returns `Value<Unit>` instead of `Value<Map<String, FieldResolutionResult>>`
+     * because the actual resolved values are stored directly in the `ObjectEngineResult` associated with
+     * the parent object. The `Value<Unit>` serves as a completion signal for the orchestration layer to
+     * know when all nested fetching (including any lazy data or nested objects) has finished.
+     *
+     * If the Value returned by this method is exceptionally completed, that means that there has been
+     * a fatal error in resolving this object, and the parent ObjectEngineResult may be incomplete. Thus, callers of this
+     * method should check for exceptional completion and handle it appropriately.
+     *
      * @param parameters ExecutionParameters containing the execution context and selection set
      * @throws Exception Only if there's a fatal error in the supervisorScope itself
      */
     fun fetchObjectSerially(
         objectType: GraphQLObjectType,
         parameters: ExecutionParameters
-    ) {
+    ): Value<Unit> {
         val instrumentationParameters =
             InstrumentationExecutionStrategyParameters(parameters.executionContext, parameters.gjParameters)
         val resolveObjectCtx = nonNullCtx(
@@ -156,24 +177,21 @@ class FieldResolver(
         try {
             val fields = collectFields(objectType, parameters).selections
             val initial: Value<Unit> = Value.fromValue(Unit)
-            val results = mutableMapOf<String, Value<FieldResolutionResult>>()
 
             // iterate over each field to build a chained execution
             // Each field will kick off only after the previous one completes
-            fields.fold(initial) { acc, field ->
+            return fields.fold(initial) { acc, field ->
                 field as QueryPlan.CollectedField
                 acc.flatMap { _ ->
                     val fieldParameters = parameters.forField(objectType, field)
-                    val value = resolveField(fieldParameters, field)
-                    results.put(field.responseKey, value)
-                    // ignore any errors thrown by the field resolver so that the chained execution can continue
-                    // these field errors will be surfaced during field completion
-                    value.thenApply { _, _ ->
-                        Unit
-                    }
+                    resolveField(fieldParameters, field)
                 }
-            }.thenApply { _, t ->
-                resolveObjectCtx.onCompleted(results, t)
+            }.map {
+                resolveObjectCtx.onCompleted(Unit, null)
+                it
+            }.recover { t ->
+                resolveObjectCtx.onCompleted(null, t)
+                Value.fromThrowable(t)
             }
         } catch (e: Exception) {
             resolveObjectCtx.onCompleted(null, e)
@@ -196,7 +214,7 @@ class FieldResolver(
     fun resolveField(
         parameters: ExecutionParameters,
         field: QueryPlan.CollectedField
-    ): Value<FieldResolutionResult> {
+    ): Value<Unit> {
         field.childPlans.forEach { launchQueryPlan(parameters, it) }
         return executeField(parameters)
     }
@@ -245,7 +263,7 @@ class FieldResolver(
      * @param parameters The execution parameters containing field and context information
      */
     @Suppress("UNCHECKED_CAST")
-    private fun executeField(parameters: ExecutionParameters): Value<FieldResolutionResult> {
+    private fun executeField(parameters: ExecutionParameters): Value<Unit> {
         val field = checkNotNull(parameters.field) { "Expected field to be non-null." }
 
         // We're fetching an individual field; the current engine result will always be an ObjectEngineResult
@@ -303,8 +321,12 @@ class FieldResolver(
         return fieldResolutionResultValue.thenCompose { v, e ->
             fieldInstrumentationCtx.onCompleted(v, e)
             if (e != null) {
-                Value.fromThrowable(e)
+                // if the field resolution failed, don't attempt to fetch lazy data or nested objects
+                // and mark this field as completed
+                Value.fromValue(Unit)
             } else {
+                // otherwise, proceed with lazy data fetching and nested object resolution
+
                 // if the result contains lazy data, begin fetching it
                 maybeFetchLazyData(
                     v!!,
@@ -321,7 +343,6 @@ class FieldResolver(
                         executionStepInfo = executionStepInfoForField,
                     )
                 )
-                Value.fromValue(v)
             }
         }
     }
@@ -497,6 +518,7 @@ class FieldResolver(
                 val result = rawValue.getOrThrow()
                 result as? FieldResolutionResult ?: throw IllegalStateException("Expected FieldResolutionResult but got ${result!!::class}")
             }
+
             else -> throw IllegalStateException(
                 "Expected the raw value slot to contain a Value.Sync<FieldResolutionResult>, but got ${rawValue::class}"
             )
@@ -522,23 +544,25 @@ class FieldResolver(
         outputType: GraphQLOutputType,
         field: QueryPlan.CollectedField,
         parameters: ExecutionParameters,
-    ) {
+    ): Value<Unit> {
         // if engineResult is null, then there is no nested object to fetch and we can return early
-        if (fieldResolutionResult.engineResult == null) return
-        when (outputType) {
+        if (fieldResolutionResult.engineResult == null) return Value.fromValue(Unit)
+        return when (outputType) {
             is GraphQLNonNull -> maybeFetchNestedObject(fieldResolutionResult, GraphQLTypeUtil.unwrapOneAs(outputType), field, parameters)
             is GraphQLList -> {
                 val engineResult = checkNotNull(fieldResolutionResult.engineResult as? Iterable<*>) { "Expected iterable engineResult but got ${fieldResolutionResult.engineResult}" }
-                engineResult.forEachIndexed { i, item ->
+                val values = engineResult.mapIndexed { i, item ->
                     check(item is Cell) { "Expected engine result to be a Cell." }
                     val frr = extractFieldResolutionResult(item)
                     val newParams = updateListItemParameters(parameters, i)
                     maybeFetchNestedObject(frr, GraphQLTypeUtil.unwrapOneAs(outputType), field, newParams)
                 }
+                Value.waitAll(values)
             }
+
             else -> {
                 // if engineResult is a scalar or simple value, then no nesting is possible and we can return
-                val oer = fieldResolutionResult.engineResult as? ObjectEngineResultImpl ?: return
+                val oer = fieldResolutionResult.engineResult as? ObjectEngineResultImpl ?: return Value.fromValue(Unit)
                 fetchObject(
                     oer.graphQLObjectType,
                     parameters.forObjectTraversal(field, oer, fieldResolutionResult.localContext, fieldResolutionResult.originalSource, fieldResolutionResult.resolutionPolicy)