Implement setMain in common code (#2967)

Fixes #1720

Author: dkhalanskyjb

kotlinx-coroutines-core/js/src/Dispatchers.kt

@@ -8,8 +8,22 @@ import kotlin.coroutines.*
 
 public actual object Dispatchers {
     public actual val Default: CoroutineDispatcher = createDefaultDispatcher()
-    public actual val Main: MainCoroutineDispatcher = JsMainDispatcher(Default, false)
+    public actual val Main: MainCoroutineDispatcher
+        get() = injectedMainDispatcher ?: mainDispatcher
     public actual val Unconfined: CoroutineDispatcher = kotlinx.coroutines.Unconfined
+
+    private val mainDispatcher = JsMainDispatcher(Default, false)
+    private var injectedMainDispatcher: MainCoroutineDispatcher? = null
+
+    @PublishedApi
+    internal fun injectMain(dispatcher: MainCoroutineDispatcher) {
+        injectedMainDispatcher = dispatcher
+    }
+
+    @PublishedApi
+    internal fun resetInjectedMain() {
+        injectedMainDispatcher = null
+    }
 }
 
 private class JsMainDispatcher(

kotlinx-coroutines-core/native/src/Dispatchers.kt

@@ -6,10 +6,26 @@ package kotlinx.coroutines
 
 import kotlin.coroutines.*
 
+/** Not inside [Dispatchers], as otherwise mutating this throws an `InvalidMutabilityException`. */
+private var injectedMainDispatcher: MainCoroutineDispatcher? = null
+
 public actual object Dispatchers {
     public actual val Default: CoroutineDispatcher = createDefaultDispatcher()
-    public actual val Main: MainCoroutineDispatcher = NativeMainDispatcher(Default)
+    public actual val Main: MainCoroutineDispatcher
+        get() = injectedMainDispatcher ?: mainDispatcher
     public actual val Unconfined: CoroutineDispatcher get() = kotlinx.coroutines.Unconfined // Avoid freezing
+
+    private val mainDispatcher = NativeMainDispatcher(Default)
+
+    @PublishedApi
+    internal fun injectMain(dispatcher: MainCoroutineDispatcher) {
+        injectedMainDispatcher = dispatcher
+    }
+
+    @PublishedApi
+    internal fun resetInjectedMain() {
+        injectedMainDispatcher = null
+    }
 }
 
 private class NativeMainDispatcher(val delegate: CoroutineDispatcher) : MainCoroutineDispatcher() {

kotlinx-coroutines-test/common/src/TestDispatchers.kt

@@ -1,10 +1,13 @@
 /*
  * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
+@file:JvmName("TestDispatchers")
 
 package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.test.internal.*
+import kotlin.jvm.*
 
 /**
  * Sets the given [dispatcher] as an underlying dispatcher of [Dispatchers.Main].
@@ -13,7 +16,10 @@ import kotlinx.coroutines.*
  * It is unsafe to call this method if alive coroutines launched in [Dispatchers.Main] exist.
  */
 @ExperimentalCoroutinesApi
-public expect fun Dispatchers.setMain(dispatcher: CoroutineDispatcher)
+public fun Dispatchers.setMain(dispatcher: CoroutineDispatcher) {
+    require(dispatcher !is TestMainDispatcher) { "Dispatchers.setMain(Dispatchers.Main) is prohibited, probably Dispatchers.resetMain() should be used instead" }
+    getTestMainDispatcher().setDispatcher(dispatcher)
+}
 
 /**
  * Resets state of the [Dispatchers.Main] to the original main dispatcher.
@@ -23,4 +29,6 @@ public expect fun Dispatchers.setMain(dispatcher: CoroutineDispatcher)
  * It is unsafe to call this method if alive coroutines launched in [Dispatchers.Main] exist.
  */
 @ExperimentalCoroutinesApi
-public expect fun Dispatchers.resetMain()
+public fun Dispatchers.resetMain() {
+    getTestMainDispatcher().resetDispatcher()
+}

kotlinx-coroutines-test/common/src/internal/TestMainDispatcher.kt

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.test.internal
+import kotlinx.coroutines.*
+import kotlin.coroutines.*
+
+/**
+ * The testable main dispatcher used by kotlinx-coroutines-test.
+ * It is a [MainCoroutineDispatcher] that delegates all actions to a settable delegate.
+ */
+internal class TestMainDispatcher(private var delegate: CoroutineDispatcher):
+    MainCoroutineDispatcher(),
+    Delay by (delegate as? Delay ?: defaultDelay)
+{
+    private val mainDispatcher = delegate // the initial value passed to the constructor
+
+    override val immediate: MainCoroutineDispatcher
+        get() = (delegate as? MainCoroutineDispatcher)?.immediate ?: this
+
+    override fun dispatch(context: CoroutineContext, block: Runnable) = delegate.dispatch(context, block)
+
+    override fun isDispatchNeeded(context: CoroutineContext): Boolean = delegate.isDispatchNeeded(context)
+
+    override fun dispatchYield(context: CoroutineContext, block: Runnable) = delegate.dispatchYield(context, block)
+
+    fun setDispatcher(dispatcher: CoroutineDispatcher) {
+        delegate = dispatcher
+    }
+
+    fun resetDispatcher() {
+        delegate = mainDispatcher
+    }
+}
+
+@Suppress("INVISIBLE_MEMBER")
+private val defaultDelay
+    inline get() = DefaultDelay
+
+@Suppress("INVISIBLE_MEMBER")
+internal expect fun Dispatchers.getTestMainDispatcher(): TestMainDispatcher

kotlinx-coroutines-test/common/test/TestCoroutineDispatcherOrderTest.kt

@@ -15,7 +15,6 @@ class TestCoroutineDispatcherOrderTest {
 
     private fun expect(index: Int) {
         val wasIndex = actionIndex.incrementAndGet()
-        // println("expect($index), wasIndex=$wasIndex")
         check(index == wasIndex) { "Expecting action index $index but it is actually $wasIndex" }
     }
 

kotlinx-coroutines-test/jvm/test/TestDispatchersTest.kt renamed to kotlinx-coroutines-test/common/test/TestDispatchersTest.kt

@@ -14,7 +14,6 @@ class TestDispatchersTest {
 
     private fun expect(index: Int) {
         val wasIndex = actionIndex.incrementAndGet()
-        println("expect($index), wasIndex=$wasIndex")
         check(index == wasIndex) { "Expecting action index $index but it is actually $wasIndex" }
     }
 
@@ -69,4 +68,4 @@ class TestDispatchersTest {
             block.run()
         }
     }
-}
+}

kotlinx-coroutines-test/common/test/TestRunBlockingOrderTest.kt

@@ -15,7 +15,6 @@ class TestRunBlockingOrderTest {
 
     private fun expect(index: Int) {
         val wasIndex = actionIndex.incrementAndGet()
-        // println("expect($index), wasIndex=$wasIndex")
         check(index == wasIndex) { "Expecting action index $index but it is actually $wasIndex" }
     }
 

kotlinx-coroutines-test/js/src/TestDispatchers.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.test.internal
+import kotlinx.coroutines.*
+
+@Suppress("INVISIBLE_MEMBER")
+internal actual fun Dispatchers.getTestMainDispatcher(): TestMainDispatcher =
+    when (val mainDispatcher = Main) {
+        is TestMainDispatcher -> mainDispatcher
+        else -> TestMainDispatcher(mainDispatcher).also { injectMain(it) }
+    }

kotlinx-coroutines-test/js/src/internal/TestMainDispatcher.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.test.internal
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.internal.*
+
+internal class TestMainDispatcherFactory : MainDispatcherFactory {
+
+    override fun createDispatcher(allFactories: List<MainDispatcherFactory>): MainCoroutineDispatcher {
+        val otherFactories = allFactories.filter { it !== this }
+        val secondBestFactory = otherFactories.maxByOrNull { it.loadPriority } ?: MissingMainCoroutineDispatcherFactory
+        val dispatcher = secondBestFactory.tryCreateDispatcher(otherFactories)
+        return TestMainDispatcher(dispatcher)
+    }
+
+    /**
+     * [Int.MAX_VALUE] -- test dispatcher always wins no matter what factories are present in the classpath.
+     * By default, all actions are delegated to the second-priority dispatcher, so that it won't be the issue.
+     */
+    override val loadPriority: Int
+        get() = Int.MAX_VALUE
+}
+
+internal actual fun Dispatchers.getTestMainDispatcher(): TestMainDispatcher {
+    val mainDispatcher = Main
+    require(mainDispatcher is TestMainDispatcher) { "TestMainDispatcher is not set as main dispatcher, have $mainDispatcher instead." }
+    return mainDispatcher
+}

kotlinx-coroutines-test/jvm/src/TestDispatchers.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.test.internal
+import kotlinx.coroutines.*
+
+@Suppress("INVISIBLE_MEMBER")
+internal actual fun Dispatchers.getTestMainDispatcher(): TestMainDispatcher =
+    when (val mainDispatcher = Main) {
+        is TestMainDispatcher -> mainDispatcher
+        else -> TestMainDispatcher(mainDispatcher).also { injectMain(it) }
+    }