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

samuelAndalon · web-flow · commit c7e23e0d6e21 · 2024-07-12T23:33:02.000+01: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
@@ -63,11 +63,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)
                         }
                     }
                 }
@@ -126,9 +127,8 @@ class SyncExecutionExhaustedState(
                     executionState
                 }
 
-                val allSyncExecutionsExhausted = allSyncExecutionsExhausted()
-                if (allSyncExecutionsExhausted) {
-                    onSyncExecutionExhausted(executions.keys().toList())
+                ifAllSyncExecutionsExhausted { executionIds ->
+                    onSyncExecutionExhausted(executionIds)
                 }
             }
             override fun onCompleted(result: Any?, t: Throwable?) {
@@ -137,26 +137,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
@@ -614,9 +614,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(
@@ -631,7 +631,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
`@@ -63,11 +63,12 @@ class SyncExecutionExhaustedState(`
`63`	`63`	`override fun onCompleted(result: ExecutionResult?, t: Throwable?) {`
`64`	`64`	`if ((result != null && result.errors.size > 0) \|\| t != null) {`
`65`	`65`	`if (executions.containsKey(parameters.executionInput.executionId)) {`
`66`		`- executions.remove(parameters.executionInput.executionId)`
`67`		`- totalExecutions.set(totalExecutions.get() - 1)`
`68`		`- val allSyncExecutionsExhausted = allSyncExecutionsExhausted()`
`69`		`- if (allSyncExecutionsExhausted) {`
`70`		`- onSyncExecutionExhausted(executions.keys().toList())`
	`66`	`+ synchronized(executions) {`
	`67`	`+ executions.remove(parameters.executionInput.executionId)`
	`68`	`+ totalExecutions.set(totalExecutions.get() - 1)`
	`69`	`+ }`
	`70`	`+ ifAllSyncExecutionsExhausted { executionIds ->`
	`71`	`+ onSyncExecutionExhausted(executionIds)`
`71`	`72`	`}`
`72`	`73`	`}`
`73`	`74`	`}`
`@@ -126,9 +127,8 @@ class SyncExecutionExhaustedState(`
`126`	`127`	`executionState`
`127`	`128`	`}`
`128`	`129`
`129`		`- val allSyncExecutionsExhausted = allSyncExecutionsExhausted()`
`130`		`- if (allSyncExecutionsExhausted) {`
`131`		`- onSyncExecutionExhausted(executions.keys().toList())`
	`130`	`+ ifAllSyncExecutionsExhausted { executionIds ->`
	`131`	`+ onSyncExecutionExhausted(executionIds)`
`132`	`132`	`}`
`133`	`133`	`}`
`134`	`134`	`override fun onCompleted(result: Any?, t: Throwable?) {`
`@@ -137,26 +137,26 @@ class SyncExecutionExhaustedState(`
`137`	`137`	`executionState`
`138`	`138`	`}`
`139`	`139`
`140`		`- val allSyncExecutionsExhausted = allSyncExecutionsExhausted()`
`141`		`- if (allSyncExecutionsExhausted) {`
`142`		`- onSyncExecutionExhausted(executions.keys().toList())`
	`140`	`+ ifAllSyncExecutionsExhausted { executionIds ->`
	`141`	`+ onSyncExecutionExhausted(executionIds)`
`143`	`142`	`}`
`144`	`143`	`}`
`145`	`144`	`}`
`146`	`145`	`}`
`147`	`146`
`148`	`147`	`/**`
`149`		`- * Provide the information about when all [ExecutionInput] sharing a [GraphQLContext] exhausted their execution`
	`148`	`+ * execute a given [predicate] when all [ExecutionInput] sharing a [GraphQLContext] exhausted their execution.`
`150`	`149`	`* A Synchronous Execution is considered Exhausted when all [DataFetcher]s of all paths were executed up until`
`151`	`150`	`* a scalar leaf or a [DataFetcher] that returns a [CompletableFuture]`
`152`	`151`	`*/`
`153`		`- fun allSyncExecutionsExhausted(): Boolean = synchronized(executions) {`
`154`		`- val operationsToExecute = totalExecutions.get()`
`155`		`- when {`
`156`		`- executions.size < operationsToExecute \|\| !dataLoaderRegistry.onDispatchFuturesHandled() -> false`
`157`		`- else -> {`
`158`		`- executions.values.all(ExecutionBatchState::isSyncExecutionExhausted)`
	`152`	`+ fun ifAllSyncExecutionsExhausted(predicate: (List<ExecutionId>) -> Unit) =`
	`153`	`+ synchronized(executions) {`
	`154`	`+ val operationsToExecute = totalExecutions.get()`
	`155`	`+ if (executions.size < operationsToExecute \|\| !dataLoaderRegistry.onDispatchFuturesHandled())`
	`156`	`+ return@synchronized`
	`157`	`+`
	`158`	`+ if (executions.values.all(ExecutionBatchState::isSyncExecutionExhausted)) {`
	`159`	`+ predicate(executions.keys().toList())`
`159`	`160`	`}`
`160`	`161`	`}`
`161`		`- }`
`162`	`162`	`}`