feat(batching): v7 synchronize executions when attempting to remove an entry (#2013)

samuelAndalon · Samuel Vazquez · commit 3a91e410e935 · 2024-07-12T15:34:42.000-07:00
### 📝 Description cherry-pick #2012
diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/extensions/CompletableFutureExtensions.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/extensions/CompletableFutureExtensions.kt
@@ -37,23 +37,25 @@ fun <V> CompletableFuture<V>.dispatchIfNeeded(
     val dataLoaderRegistry = environment.dataLoaderRegistry as? KotlinDataLoaderRegistry ?: throw MissingKotlinDataLoaderRegistryException()
 
     if (dataLoaderRegistry.dataLoadersInvokedOnDispatch()) {
-        val cantContinueExecution = when {
+        when {
             environment.graphQlContext.hasKey(ExecutionLevelDispatchedState::class) -> {
-                environment
-                    .graphQlContext.get<ExecutionLevelDispatchedState>(ExecutionLevelDispatchedState::class)
-                    .allExecutionsDispatched(Level(environment.executionStepInfo.path.level))
+                val cantContinueExecution =
+                    environment
+                        .graphQlContext.get<ExecutionLevelDispatchedState>(ExecutionLevelDispatchedState::class)
+                        .allExecutionsDispatched(Level(environment.executionStepInfo.path.level))
+                if (cantContinueExecution) {
+                    dataLoaderRegistry.dispatchAll()
+                }
             }
             environment.graphQlContext.hasKey(SyncExecutionExhaustedState::class) -> {
                 environment
                     .graphQlContext.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
-                    .allSyncExecutionsExhausted()
+                    .ifAllSyncExecutionsExhausted {
+                        dataLoaderRegistry.dispatchAll()
+                    }
             }
             else -> throw MissingInstrumentationStateException()
         }
-
-        if (cantContinueExecution) {
-            dataLoaderRegistry.dispatchAll()
-        }
     }
     return this
 }
diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/state/SyncExecutionExhaustedState.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/state/SyncExecutionExhaustedState.kt
@@ -84,11 +84,12 @@ class SyncExecutionExhaustedState(
             override fun onCompleted(result: ExecutionResult?, t: Throwable?) {
                 if ((result != null && result.errors.size > 0) || t != null) {
                     if (executions.containsKey(parameters.executionInput.executionId)) {
-                        executions.remove(parameters.executionInput.executionId)
-                        totalExecutions.set(totalExecutions.get() - 1)
-                        val allSyncExecutionsExhausted = allSyncExecutionsExhausted()
-                        if (allSyncExecutionsExhausted) {
-                            onSyncExecutionExhausted(executions.keys().toList())
+                        synchronized(executions) {
+                            executions.remove(parameters.executionInput.executionId)
+                            totalExecutions.set(totalExecutions.get() - 1)
+                        }
+                        ifAllSyncExecutionsExhausted { executionIds ->
+                            onSyncExecutionExhausted(executionIds)
                         }
                     }
                 }
@@ -147,9 +148,8 @@ class SyncExecutionExhaustedState(
                     executionState
                 }
 
-                val allSyncExecutionsExhausted = allSyncExecutionsExhausted()
-                if (allSyncExecutionsExhausted) {
-                    onSyncExecutionExhausted(executions.keys().toList())
+                ifAllSyncExecutionsExhausted { executionIds ->
+                    onSyncExecutionExhausted(executionIds)
                 }
             }
             override fun onCompleted(result: Any?, t: Throwable?) {
@@ -158,26 +158,26 @@ class SyncExecutionExhaustedState(
                     executionState
                 }
 
-                val allSyncExecutionsExhausted = allSyncExecutionsExhausted()
-                if (allSyncExecutionsExhausted) {
-                    onSyncExecutionExhausted(executions.keys().toList())
+                ifAllSyncExecutionsExhausted { executionIds ->
+                    onSyncExecutionExhausted(executionIds)
                 }
             }
         }
     }
 
     /**
-     * Provide the information about when all [ExecutionInput] sharing a [GraphQLContext] exhausted their execution
+     * execute a given [predicate] when all [ExecutionInput] sharing a [GraphQLContext] exhausted their execution.
      * A Synchronous Execution is considered Exhausted when all [DataFetcher]s of all paths were executed up until
      * a scalar leaf or a [DataFetcher] that returns a [CompletableFuture]
      */
-    fun allSyncExecutionsExhausted(): Boolean = synchronized(executions) {
-        val operationsToExecute = totalExecutions.get()
-        when {
-            executions.size < operationsToExecute || !dataLoaderRegistry.onDispatchFuturesHandled() -> false
-            else -> {
-                executions.values.all(ExecutionBatchState::isSyncExecutionExhausted)
+    fun ifAllSyncExecutionsExhausted(predicate: (List<ExecutionId>) -> Unit) =
+        synchronized(executions) {
+            val operationsToExecute = totalExecutions.get()
+            if (executions.size < operationsToExecute || !dataLoaderRegistry.onDispatchFuturesHandled())
+                return@synchronized
+
+            if (executions.values.all(ExecutionBatchState::isSyncExecutionExhausted)) {
+                predicate(executions.keys().toList())
             }
         }
-    }
 }
diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedInstrumentationTest.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedInstrumentationTest.kt
@@ -616,9 +616,9 @@ class DataLoaderSyncExecutionExhaustedInstrumentationTest {
     fun `Instrumentation should not consider executions that thrown exceptions`() {
         val executions = listOf(
             ExecutionInput.newExecutionInput("query test1 { astronaut(id: 1) { id name } }").operationName("test1").build(),
-            ExecutionInput.newExecutionInput("query test2 { astronaut(id: 2) { id name } }").operationName("test2").build(),
-            ExecutionInput.newExecutionInput("query test3 { mission(id: 3) { id designation } }").operationName("test3").build(),
-            ExecutionInput.newExecutionInput("query test4 { mission(id: 4) { designation } }").operationName("OPERATION_NOT_IN_DOCUMENT").build()
+            ExecutionInput.newExecutionInput("query test2 { astronaut(id: 2) { id name } }").operationName("OPERATION_NOT_IN_DOCUMENT").build(),
+            ExecutionInput.newExecutionInput("query test3 { mission(id: 3) { id designation } }").operationName("OPERATION_NOT_IN_DOCUMENT").build(),
+            ExecutionInput.newExecutionInput("query test4 { mission(id: 4) { designation } }").operationName("test4").build()
         )
 
         val (results, kotlinDataLoaderRegistry) = AstronautGraphQL.execute(
@@ -633,7 +633,7 @@ class DataLoaderSyncExecutionExhaustedInstrumentationTest {
         val missionStatistics = kotlinDataLoaderRegistry.dataLoadersMap["MissionDataLoader"]?.statistics
 
         assertEquals(1, astronautStatistics?.batchInvokeCount)
-        assertEquals(2, astronautStatistics?.batchLoadCount)
+        assertEquals(1, astronautStatistics?.batchLoadCount)
 
         assertEquals(1, missionStatistics?.batchInvokeCount)
         assertEquals(1, missionStatistics?.batchLoadCount)

Original file line number	Diff line number	Diff line change
`@@ -84,11 +84,12 @@ class SyncExecutionExhaustedState(`
`84`	`84`	`override fun onCompleted(result: ExecutionResult?, t: Throwable?) {`
`85`	`85`	`if ((result != null && result.errors.size > 0) \|\| t != null) {`
`86`	`86`	`if (executions.containsKey(parameters.executionInput.executionId)) {`
`87`		`- executions.remove(parameters.executionInput.executionId)`
`88`		`- totalExecutions.set(totalExecutions.get() - 1)`
`89`		`- val allSyncExecutionsExhausted = allSyncExecutionsExhausted()`
`90`		`- if (allSyncExecutionsExhausted) {`
`91`		`- onSyncExecutionExhausted(executions.keys().toList())`
	`87`	`+ synchronized(executions) {`
	`88`	`+ executions.remove(parameters.executionInput.executionId)`
	`89`	`+ totalExecutions.set(totalExecutions.get() - 1)`
	`90`	`+ }`
	`91`	`+ ifAllSyncExecutionsExhausted { executionIds ->`
	`92`	`+ onSyncExecutionExhausted(executionIds)`
`92`	`93`	`}`
`93`	`94`	`}`
`94`	`95`	`}`
`@@ -147,9 +148,8 @@ class SyncExecutionExhaustedState(`
`147`	`148`	`executionState`
`148`	`149`	`}`
`149`	`150`
`150`		`- val allSyncExecutionsExhausted = allSyncExecutionsExhausted()`
`151`		`- if (allSyncExecutionsExhausted) {`
`152`		`- onSyncExecutionExhausted(executions.keys().toList())`
	`151`	`+ ifAllSyncExecutionsExhausted { executionIds ->`
	`152`	`+ onSyncExecutionExhausted(executionIds)`
`153`	`153`	`}`
`154`	`154`	`}`
`155`	`155`	`override fun onCompleted(result: Any?, t: Throwable?) {`
`@@ -158,26 +158,26 @@ class SyncExecutionExhaustedState(`
`158`	`158`	`executionState`
`159`	`159`	`}`
`160`	`160`
`161`		`- val allSyncExecutionsExhausted = allSyncExecutionsExhausted()`
`162`		`- if (allSyncExecutionsExhausted) {`
`163`		`- onSyncExecutionExhausted(executions.keys().toList())`
	`161`	`+ ifAllSyncExecutionsExhausted { executionIds ->`
	`162`	`+ onSyncExecutionExhausted(executionIds)`
`164`	`163`	`}`
`165`	`164`	`}`
`166`	`165`	`}`
`167`	`166`	`}`
`168`	`167`
`169`	`168`	`/**`
`170`		`- * Provide the information about when all [ExecutionInput] sharing a [GraphQLContext] exhausted their execution`
	`169`	`+ * execute a given [predicate] when all [ExecutionInput] sharing a [GraphQLContext] exhausted their execution.`
`171`	`170`	`* A Synchronous Execution is considered Exhausted when all [DataFetcher]s of all paths were executed up until`
`172`	`171`	`* a scalar leaf or a [DataFetcher] that returns a [CompletableFuture]`
`173`	`172`	`*/`
`174`		`- fun allSyncExecutionsExhausted(): Boolean = synchronized(executions) {`
`175`		`- val operationsToExecute = totalExecutions.get()`
`176`		`- when {`
`177`		`- executions.size < operationsToExecute \|\| !dataLoaderRegistry.onDispatchFuturesHandled() -> false`
`178`		`- else -> {`
`179`		`- executions.values.all(ExecutionBatchState::isSyncExecutionExhausted)`
	`173`	`+ fun ifAllSyncExecutionsExhausted(predicate: (List<ExecutionId>) -> Unit) =`
	`174`	`+ synchronized(executions) {`
	`175`	`+ val operationsToExecute = totalExecutions.get()`
	`176`	`+ if (executions.size < operationsToExecute \|\| !dataLoaderRegistry.onDispatchFuturesHandled())`
	`177`	`+ return@synchronized`
	`178`	`+`
	`179`	`+ if (executions.values.all(ExecutionBatchState::isSyncExecutionExhausted)) {`
	`180`	`+ predicate(executions.keys().toList())`
`180`	`181`	`}`
`181`	`182`	`}`
`182`		`- }`
`183`	`183`	`}`