diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedDataLoaderDispatcher.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedDataLoaderDispatcher.kt index d646b44187..7d66eaa21f 100644 --- a/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedDataLoaderDispatcher.kt +++ b/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedDataLoaderDispatcher.kt @@ -21,6 +21,7 @@ import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state. import org.dataloader.DataLoader import org.dataloader.instrumentation.DataLoaderInstrumentation import org.dataloader.instrumentation.DataLoaderInstrumentationContext +import java.util.concurrent.atomic.AtomicInteger /** * Custom [DataLoaderInstrumentation] implementation that helps to calculate the state of [DataLoader]s in the @@ -29,20 +30,23 @@ import org.dataloader.instrumentation.DataLoaderInstrumentationContext class DataLoaderSyncExecutionExhaustedDataLoaderDispatcher( private val syncExecutionExhaustedState: SyncExecutionExhaustedState ) : DataLoaderInstrumentation { - private val contextForSyncExecutionExhausted: DataLoaderInstrumentationContext = + override fun beginLoad( + dataLoader: DataLoader<*, *>, + key: Any, + loadContext: Any? + ): DataLoaderInstrumentationContext = object : DataLoaderInstrumentationContext { + /** + * Counter the load was added to, cleared on completion so the load is only decreased once + */ + private var loadCounter: AtomicInteger? = null + override fun onDispatched() { - syncExecutionExhaustedState.onDataLoaderLoadDispatched() + loadCounter = syncExecutionExhaustedState.trackDataLoaderLoad() } override fun onCompleted(result: Any?, t: Throwable?) { - syncExecutionExhaustedState.onDataLoaderLoadCompleted() + loadCounter?.let(syncExecutionExhaustedState::onDataLoaderLoadCompleted) + loadCounter = null } } - - override fun beginLoad( - dataLoader: DataLoader<*, *>, - key: Any, - loadContext: Any? - ): DataLoaderInstrumentationContext = - contextForSyncExecutionExhausted } diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/state/DataLoaderRegistryState.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/state/DataLoaderRegistryState.kt index 1d0eb4fc94..353b76868b 100644 --- a/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/state/DataLoaderRegistryState.kt +++ b/executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/state/DataLoaderRegistryState.kt @@ -26,22 +26,27 @@ import java.util.concurrent.atomic.AtomicInteger */ class DataLoaderRegistryState { /** - * Count [DataLoader.load] invocations + * Count [DataLoader.load] invocations that are not completed yet and were invoked + * after the last [DataLoaderRegistry.dispatchAll] */ - private val loadCounter = AtomicInteger(0) + @Volatile + private var loadCounter = AtomicInteger(0) /** - * Snapshot of [loadCounter] when [DataLoaderRegistry.dispatchAll] is invoked, - * then on every load complete decrease it + * Count [DataLoader.load] invocations that are not completed yet and were invoked + * before the last [DataLoaderRegistry.dispatchAll] */ - private val onDispatchAllLoadCounter = AtomicInteger(0) + @Volatile + private var onDispatchAllLoadCounter = AtomicInteger(0) /** - * Take snapshot of [loadCounter] when [DataLoaderRegistry.dispatchAll] is invoked + * Take snapshot of [loadCounter] when [DataLoaderRegistry.dispatchAll] is invoked, + * loads tracked by [trackDataLoaderLoad] before the snapshot will decrease [onDispatchAllLoadCounter] when completed */ + @Synchronized fun takeSnapshot() { - onDispatchAllLoadCounter.set(loadCounter.get()) - loadCounter.set(0) + onDispatchAllLoadCounter = loadCounter + loadCounter = AtomicInteger(0) } /** @@ -59,6 +64,11 @@ class DataLoaderRegistryState { /** * Increase [loadCounter] when [DataLoader.load] is invoked */ + @Deprecated( + "Loads are tracked by SyncExecutionExhaustedState through DataLoaderSyncExecutionExhaustedDataLoaderDispatcher, " + + "register it with KotlinDataLoaderRegistryFactory.generate instead of calling this method directly. " + + "Will be removed in the next major version." + ) fun onDataLoaderLoadDispatched() { loadCounter.incrementAndGet() } @@ -66,7 +76,34 @@ class DataLoaderRegistryState { /** * Decrease [onDispatchAllLoadCounter] when [DataLoader.load] returned [CompletableFuture] completes */ + @Deprecated( + "Loads are tracked by SyncExecutionExhaustedState through DataLoaderSyncExecutionExhaustedDataLoaderDispatcher, " + + "register it with KotlinDataLoaderRegistryFactory.generate instead of calling this method directly. " + + "This method decreases onDispatchAllLoadCounter even for loads that were not included in the last snapshot. " + + "Will be removed in the next major version." + ) fun onDataLoaderLoadCompleted() { onDispatchAllLoadCounter.decrementAndGet() } + + /** + * Increase [loadCounter] when [DataLoader.load] is invoked + * + * @return the counter that the load was added to, it needs to be provided to [onDataLoaderLoadCompleted] + * when the [CompletableFuture] returned by [DataLoader.load] completes + */ + @Synchronized + internal fun trackDataLoaderLoad(): AtomicInteger = + loadCounter.also(AtomicInteger::incrementAndGet) + + /** + * Decrease the counter that the load was added to when [DataLoader.load] returned [CompletableFuture] completes, + * a load that completes before the next snapshot, for example a cache hit on a completed [CompletableFuture], + * will decrease [loadCounter] instead of [onDispatchAllLoadCounter] + * + * @param loadCounter the counter returned by [trackDataLoaderLoad] + */ + internal fun onDataLoaderLoadCompleted(loadCounter: AtomicInteger) { + loadCounter.decrementAndGet() + } } 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 index 67e182dddc..9de9888aa2 100644 --- 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 @@ -157,6 +157,12 @@ class SyncExecutionExhaustedState( /** * This is invoked right after a [DataLoader.load] was dispatched */ + @Deprecated( + "Loads are tracked by DataLoaderSyncExecutionExhaustedDataLoaderDispatcher, register it with " + + "KotlinDataLoaderRegistryFactory.generate instead of calling this method directly. " + + "Will be removed in the next major version." + ) + @Suppress("DEPRECATION") fun onDataLoaderLoadDispatched() { dataLoadersDispatchState.onDataLoaderLoadDispatched() } @@ -164,6 +170,13 @@ class SyncExecutionExhaustedState( /** * This is invoked right after a [DataLoader.load] was completed */ + @Deprecated( + "Loads are tracked by DataLoaderSyncExecutionExhaustedDataLoaderDispatcher, register it with " + + "KotlinDataLoaderRegistryFactory.generate instead of calling this method directly. " + + "This method decreases the dispatched loads counter even for loads that were not included in the last dispatch. " + + "Will be removed in the next major version." + ) + @Suppress("DEPRECATION") fun onDataLoaderLoadCompleted() { dataLoadersDispatchState.onDataLoaderLoadCompleted() if (allSyncExecutionsExhausted()) { @@ -171,6 +184,26 @@ class SyncExecutionExhaustedState( } } + /** + * This is invoked right after a [DataLoader.load] was dispatched + * + * @return the counter that the load was added to, it needs to be provided to [onDataLoaderLoadCompleted] + */ + internal fun trackDataLoaderLoad(): AtomicInteger = + dataLoadersDispatchState.trackDataLoaderLoad() + + /** + * This is invoked right after a [DataLoader.load] was completed + * + * @param loadCounter the counter returned by [trackDataLoaderLoad] when the load was dispatched + */ + internal fun onDataLoaderLoadCompleted(loadCounter: AtomicInteger) { + dataLoadersDispatchState.onDataLoaderLoadCompleted(loadCounter) + if (allSyncExecutionsExhausted()) { + dataLoaderRegistryProvider.invoke().dispatchAll() + } + } + /** * check if 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 diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedCacheHitTest.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedCacheHitTest.kt new file mode 100644 index 0000000000..d5fb854e83 --- /dev/null +++ b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedCacheHitTest.kt @@ -0,0 +1,259 @@ +/* + * Copyright 2026 Expedia, Inc + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion + +import com.expediagroup.graphql.dataloader.KotlinDataLoader +import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistryFactory +import com.expediagroup.graphql.dataloader.instrumentation.extensions.dispatchIfNeeded +import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState +import graphql.ExecutionInput +import graphql.ExecutionResult +import graphql.GraphQL +import graphql.GraphQLContext +import graphql.schema.DataFetcher +import graphql.schema.DataFetchingEnvironment +import graphql.schema.idl.RuntimeWiring +import graphql.schema.idl.SchemaGenerator +import graphql.schema.idl.SchemaParser +import graphql.schema.idl.TypeRuntimeWiring +import org.dataloader.DataLoader +import org.dataloader.DataLoaderFactory +import org.dataloader.DataLoaderRegistry +import org.junit.jupiter.api.Test +import java.util.concurrent.CompletableFuture +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.LinkedBlockingQueue +import java.util.concurrent.TimeUnit +import java.util.concurrent.TimeoutException +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertNull +import kotlin.test.fail + +/** + * A [DataLoader.load] that hits the cache still notifies the [DataLoaderSyncExecutionExhaustedDataLoaderDispatcher], + * and when the cached future is already completed it does so synchronously, before any dispatch happens. + * These tests make sure that such loads do not corrupt the dispatch state and leave later loads undispatched. + * + * Every future is completed by the test itself so each scenario always runs in the same order: + * batch loads are completed with [completeBatchLoad] and fields that resolve asynchronously + * are released with [openGate]. + */ +class DataLoaderSyncExecutionExhaustedCacheHitTest { + private val batchLoads = LinkedBlockingQueue() + private val gates = ConcurrentHashMap>() + + private val entityDataLoader = object : KotlinDataLoader { + override val dataLoaderName: String = "EntityDataLoader" + override fun getDataLoader(graphQLContext: GraphQLContext): DataLoader = + DataLoaderFactory.newDataLoader { keys: List -> + CompletableFuture>().also { future -> batchLoads.add(BatchLoad(keys.toList(), future)) } + } + } + + private val entityDataFetcher = DataFetcher { environment -> + environment.entityDataLoader().load(environment.getArgument("key")!!) + } + + private val gatedEntityDataFetcher = DataFetcher { environment -> + val key = environment.getArgument("key")!! + gate(environment.getArgument("gate")!!) + .thenCompose { environment.entityDataLoader().load(key).dispatchIfNeeded(environment) } + } + + private val graphQL = GraphQL.newGraphQL( + SchemaGenerator().makeExecutableSchema( + SchemaParser().parse( + """ + type Query { + entity(key: String!): String + gatedEntity(key: String!, gate: String!): String + } + """.trimIndent() + ), + RuntimeWiring.newRuntimeWiring() + .type( + TypeRuntimeWiring.newTypeWiring("Query") + .dataFetcher("entity", entityDataFetcher) + .dataFetcher("gatedEntity", gatedEntityDataFetcher) + ) + .build() + ) + ) + .instrumentation(GraphQLSyncExecutionExhaustedDataLoaderDispatcher()) + .doNotAutomaticallyDispatchDataLoader() + .build() + + @Test + fun `Instrumentation should dispatch a load invoked after a cache hit on a completed future`() { + val results = executeOperations( + "{ entity(key: \"a\") }", + "{ gatedEntity(key: \"a\", gate: \"cacheHit\") }", + "{ gatedEntity(key: \"b\", gate: \"newKey\") }" + ) + + completeBatchLoad(listOf("a")) + openGate("cacheHit") + openGate("newKey") + completeBatchLoad(listOf("b")) + + assertEquals( + listOf( + mapOf("entity" to "value-a"), + mapOf("gatedEntity" to "value-a"), + mapOf("gatedEntity" to "value-b") + ), + results.awaitData() + ) + assertNoMoreBatchLoads() + } + + @Test + fun `Instrumentation should dispatch a load invoked after a cache hit on a completed future in a single operation`() { + val results = executeOperations( + "{ first: entity(key: \"a\") second: gatedEntity(key: \"a\", gate: \"cacheHit\") third: gatedEntity(key: \"b\", gate: \"newKey\") }" + ) + + completeBatchLoad(listOf("a")) + openGate("cacheHit") + openGate("newKey") + completeBatchLoad(listOf("b")) + + assertEquals( + listOf(mapOf("first" to "value-a", "second" to "value-a", "third" to "value-b")), + results.awaitData() + ) + assertNoMoreBatchLoads() + } + + @Test + fun `Instrumentation should dispatch a load invoked after a cache hit on a pending future`() { + val results = executeOperations( + "{ entity(key: \"a\") }", + "{ gatedEntity(key: \"a\", gate: \"cacheHit\") }", + "{ gatedEntity(key: \"b\", gate: \"newKey\") }" + ) + + val firstBatchLoad = awaitBatchLoad() + assertEquals(listOf("a"), firstBatchLoad.keys) + openGate("cacheHit") + firstBatchLoad.complete() + openGate("newKey") + completeBatchLoad(listOf("b")) + + assertEquals( + listOf( + mapOf("entity" to "value-a"), + mapOf("gatedEntity" to "value-a"), + mapOf("gatedEntity" to "value-b") + ), + results.awaitData() + ) + assertNoMoreBatchLoads() + } + + @Test + fun `Instrumentation should dispatch every load invoked after a cache hit on a completed future`() { + val results = executeOperations( + "{ entity(key: \"a\") }", + "{ gatedEntity(key: \"a\", gate: \"cacheHit\") }", + "{ first: gatedEntity(key: \"b\", gate: \"newKeys\") second: gatedEntity(key: \"c\", gate: \"newKeys\") }" + ) + + completeBatchLoad(listOf("a")) + openGate("cacheHit") + openGate("newKeys") + val secondBatchLoad = awaitBatchLoad().also(BatchLoad::complete) + val thirdBatchLoad = awaitBatchLoad().also(BatchLoad::complete) + + assertEquals(setOf("b", "c"), (secondBatchLoad.keys + thirdBatchLoad.keys).toSet()) + assertEquals( + listOf( + mapOf("entity" to "value-a"), + mapOf("gatedEntity" to "value-a"), + mapOf("first" to "value-b", "second" to "value-c") + ), + results.awaitData() + ) + assertNoMoreBatchLoads() + } + + private fun DataFetchingEnvironment.entityDataLoader(): DataLoader = + getDataLoader("EntityDataLoader")!! + + private fun gate(name: String): CompletableFuture = + gates.computeIfAbsent(name) { CompletableFuture() } + + private fun openGate(name: String) { + gate(name).complete(Unit) + } + + private fun awaitBatchLoad(): BatchLoad = + assertNotNull( + batchLoads.poll(HANG_TIMEOUT_MS, TimeUnit.MILLISECONDS), + "expected a batch load but none was dispatched" + ) + + private fun completeBatchLoad(expectedKeys: List) { + val batchLoad = awaitBatchLoad() + assertEquals(expectedKeys, batchLoad.keys) + batchLoad.complete() + } + + private fun assertNoMoreBatchLoads() { + assertNull(batchLoads.poll(), "unexpected batch load") + } + + private fun executeOperations(vararg queries: String): List> { + val graphQLContext = GraphQLContext.getDefault() + val syncExecutionExhaustedState = SyncExecutionExhaustedState(queries.size) { + graphQLContext.get(DataLoaderRegistry::class) + } + graphQLContext.put(SyncExecutionExhaustedState::class, syncExecutionExhaustedState) + val dataLoaderRegistry = KotlinDataLoaderRegistryFactory(listOf(entityDataLoader)).generate( + graphQLContext, + DataLoaderSyncExecutionExhaustedDataLoaderDispatcher(syncExecutionExhaustedState) + ) + + return queries.map { query -> + graphQL.executeAsync( + ExecutionInput.newExecutionInput(query) + .dataLoaderRegistry(dataLoaderRegistry) + .graphQLContext { it.of(graphQLContext) } + .build() + ) + } + } + + private fun List>.awaitData(): List = + try { + CompletableFuture.allOf(*toTypedArray()).get(HANG_TIMEOUT_MS, TimeUnit.MILLISECONDS) + map { it.join().getData() } + } catch (e: TimeoutException) { + fail("operations did not complete") + } + + private class BatchLoad(val keys: List, private val future: CompletableFuture>) { + fun complete() { + future.complete(keys.map { key -> "value-$key" }) + } + } + + companion object { + private const val HANG_TIMEOUT_MS = 5_000L + } +} diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedDataLoaderDispatcherTest.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedDataLoaderDispatcherTest.kt new file mode 100644 index 0000000000..7b0e07e93e --- /dev/null +++ b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedDataLoaderDispatcherTest.kt @@ -0,0 +1,62 @@ +/* + * Copyright 2026 Expedia, Inc + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion + +import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState +import io.mockk.mockk +import org.dataloader.DataLoader +import org.junit.jupiter.api.Test +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +class DataLoaderSyncExecutionExhaustedDataLoaderDispatcherTest { + private val syncExecutionExhaustedState = SyncExecutionExhaustedState(1) { error("no execution started, registry should not be dispatched") } + private val dispatcher = DataLoaderSyncExecutionExhaustedDataLoaderDispatcher(syncExecutionExhaustedState) + private val dataLoader = mockk>() + + @Test + fun `Dispatcher should track a load from dispatch until completion`() { + val load = dispatcher.beginLoad(dataLoader, "a", null) + + load.onDispatched() + assertTrue(syncExecutionExhaustedState.dataLoadersLoadInvokedAfterDispatchAll()) + + load.onCompleted("value-a", null) + assertFalse(syncExecutionExhaustedState.dataLoadersLoadInvokedAfterDispatchAll()) + } + + @Test + fun `Dispatcher should decrease a load only once when its completion is notified twice`() { + val load = dispatcher.beginLoad(dataLoader, "a", null) + load.onDispatched() + load.onCompleted("value-a", null) + load.onCompleted("value-a", null) + + dispatcher.beginLoad(dataLoader, "b", null).onDispatched() + + assertTrue(syncExecutionExhaustedState.dataLoadersLoadInvokedAfterDispatchAll()) + } + + @Test + fun `Dispatcher should ignore the completion of a load that was never dispatched`() { + dispatcher.beginLoad(dataLoader, "a", null).onCompleted(null, IllegalStateException("load failed")) + + dispatcher.beginLoad(dataLoader, "b", null).onDispatched() + + assertTrue(syncExecutionExhaustedState.dataLoadersLoadInvokedAfterDispatchAll()) + } +} diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/state/DataLoaderRegistryStateTest.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/state/DataLoaderRegistryStateTest.kt new file mode 100644 index 0000000000..3e1d7ea219 --- /dev/null +++ b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/state/DataLoaderRegistryStateTest.kt @@ -0,0 +1,110 @@ +/* + * Copyright 2026 Expedia, Inc + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state + +import org.junit.jupiter.api.Test +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +class DataLoaderRegistryStateTest { + private val state = DataLoaderRegistryState() + + @Test + fun `State should start with no loads to dispatch and no dispatched loads pending`() { + assertFalse(state.dataLoadersLoadInvokedAfterDispatchAll()) + assertTrue(state.onDispatchAllFuturesCompleted()) + } + + @Test + fun `State should report a load invoked after the last dispatch until it is included in a snapshot`() { + state.trackDataLoaderLoad() + assertTrue(state.dataLoadersLoadInvokedAfterDispatchAll()) + assertTrue(state.onDispatchAllFuturesCompleted()) + + state.takeSnapshot() + assertFalse(state.dataLoadersLoadInvokedAfterDispatchAll()) + assertFalse(state.onDispatchAllFuturesCompleted()) + } + + @Test + fun `State should complete the dispatched loads when every load included in the snapshot completes`() { + val first = state.trackDataLoaderLoad() + val second = state.trackDataLoaderLoad() + state.takeSnapshot() + + state.onDataLoaderLoadCompleted(first) + assertFalse(state.onDispatchAllFuturesCompleted()) + + state.onDataLoaderLoadCompleted(second) + assertTrue(state.onDispatchAllFuturesCompleted()) + } + + @Test + fun `State should not be affected by a load that completes before the next snapshot`() { + state.takeSnapshot() + + val cacheHit = state.trackDataLoaderLoad() + state.onDataLoaderLoadCompleted(cacheHit) + + assertFalse(state.dataLoadersLoadInvokedAfterDispatchAll()) + assertTrue(state.onDispatchAllFuturesCompleted()) + + state.trackDataLoaderLoad() + assertTrue(state.dataLoadersLoadInvokedAfterDispatchAll()) + assertTrue(state.onDispatchAllFuturesCompleted()) + } + + @Test + fun `State should keep waiting for dispatched loads when a load invoked after the snapshot completes first`() { + val dispatched = state.trackDataLoaderLoad() + state.takeSnapshot() + + val cacheHit = state.trackDataLoaderLoad() + state.onDataLoaderLoadCompleted(cacheHit) + assertFalse(state.onDispatchAllFuturesCompleted()) + + state.onDataLoaderLoadCompleted(dispatched) + assertTrue(state.onDispatchAllFuturesCompleted()) + } + + @Test + fun `State should decrease the snapshot a load was included in even when newer loads were invoked`() { + val dispatched = state.trackDataLoaderLoad() + state.takeSnapshot() + val pending = state.trackDataLoaderLoad() + + state.onDataLoaderLoadCompleted(dispatched) + assertTrue(state.onDispatchAllFuturesCompleted()) + assertTrue(state.dataLoadersLoadInvokedAfterDispatchAll()) + + state.onDataLoaderLoadCompleted(pending) + assertFalse(state.dataLoadersLoadInvokedAfterDispatchAll()) + } + + @Test + @Suppress("DEPRECATION") + fun `State should keep the previous behavior for the deprecated load callbacks`() { + state.onDataLoaderLoadDispatched() + assertTrue(state.dataLoadersLoadInvokedAfterDispatchAll()) + + state.takeSnapshot() + assertFalse(state.onDispatchAllFuturesCompleted()) + + state.onDataLoaderLoadCompleted() + assertTrue(state.onDispatchAllFuturesCompleted()) + } +}