make MainDispatcherChecker resilient to a change of Main dispatcher t…

…hread (#1208)

For example, in Swing it is not guaranteed that the Event Dispatch
Thread would never change. Suggested check works with just a few simple
comparisions on the happy path and only requires dispatching coroutines
if the thread doesn't match.

lifecycle/lifecycle-runtime/build.gradle

@@ -73,6 +73,12 @@ kotlin {
             dependsOn(jvmMain)
         }
 
+        desktopTest {
+            dependencies {
+                implementation(libs.kotlinCoroutinesSwing)
+            }
+        }
+
         androidMain {
             dependsOn(jvmMain)
             dependencies {

lifecycle/lifecycle-runtime/src/desktopMain/kotlin/androidx/lifecycle/MainDispatcherChecker.desktop.kt

@@ -16,30 +16,41 @@
 package androidx.lifecycle
 
 import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 
 internal object MainDispatcherChecker {
-    private val isMainDispatcherAvailable: Boolean
-    private var isMainDispatcherThread = ThreadLocal.withInitial { false }
+    private var isMainDispatcherAvailable: Boolean = true
+    @Volatile
+    private var mainDispatcherThread: Thread? = null
 
-    init {
-        isMainDispatcherAvailable = try {
-            runBlocking {
-                launch(Dispatchers.Main.immediate) {
-                    isMainDispatcherThread.set(true)
-                }
+    private fun storeMainDispatcherThread() {
+        try {
+            runBlocking(Dispatchers.Main.immediate) {
+                mainDispatcherThread = Thread.currentThread()
             }
-            true
         } catch (_: IllegalStateException) {
             // No main dispatchers are present in the classpath
-            false
+            isMainDispatcherAvailable = false
         }
     }
 
-    fun isMainDispatcherThread(): Boolean = if (isMainDispatcherAvailable) {
-        isMainDispatcherThread.get()
-    } else {
-        true
+    fun isMainDispatcherThread(): Boolean {
+        if (!isMainDispatcherAvailable) {
+            return true
+        }
+        val currentThread = Thread.currentThread()
+        // if the thread has already been retrieved,
+        // we can just check whether we are currently running on the same thread
+        if (currentThread === mainDispatcherThread) {
+            return true
+        }
+        // if threads do not match, it is either:
+        // * field is not initialized yet
+        // * Swing's EDT may have changed
+        // * it is not the main thread indeed
+        // let's recheck to make sure the field has an actual value
+        // it is potentially a long operation, but it happens only not on the happy path
+        storeMainDispatcherThread()
+        return !isMainDispatcherAvailable || currentThread === mainDispatcherThread
     }
 }

lifecycle/lifecycle-runtime/src/desktopTest/kotlin/MainDispatcherCheckerTest.kt

@@ -0,0 +1,112 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * 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
+ *
+ *      http://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.
+ */
+
+import androidx.lifecycle.MainDispatcherChecker
+import java.util.concurrent.ExecutorService
+import java.util.concurrent.Executors
+import kotlin.coroutines.CoroutineContext
+import kotlin.test.Test
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.MainCoroutineDispatcher
+import kotlinx.coroutines.Runnable
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * 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
+ *
+ *      http://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.
+ */
+
+class MainDispatcherCheckerTest {
+    @Test
+    fun checkMainDispatcher() {
+        runBlocking(Dispatchers.Main) {
+            assertTrue(MainDispatcherChecker.isMainDispatcherThread())
+        }
+        runBlocking(Dispatchers.Main.immediate) {
+            assertTrue(MainDispatcherChecker.isMainDispatcherThread())
+        }
+    }
+
+    @Test
+    fun checkNonMainDispatcher() {
+        runBlocking(Dispatchers.IO) {
+            assertFalse(MainDispatcherChecker.isMainDispatcherThread())
+        }
+    }
+
+    @OptIn(ExperimentalCoroutinesApi::class)
+    @Test
+    fun checkMainDispatcherChanged() {
+        try {
+            Dispatchers.setMain(ThreadChangingMainDispatcher)
+            runBlocking(Dispatchers.Main) {
+                assertTrue(MainDispatcherChecker.isMainDispatcherThread())
+            }
+            ThreadChangingMainDispatcher.changeThread()
+            runBlocking(Dispatchers.Main) {
+                assertTrue(MainDispatcherChecker.isMainDispatcherThread())
+            }
+        } finally {
+            Dispatchers.resetMain()
+        }
+    }
+
+    private object ThreadChangingMainDispatcher : MainCoroutineDispatcher() {
+        private var thread: Thread? = null
+        private var executor = newExecutorService()
+
+        override val immediate: MainCoroutineDispatcher
+            get() = this
+
+        override fun dispatch(context: CoroutineContext, block: Runnable) {
+            // support reentrancy
+            if (Thread.currentThread() == thread) {
+                block.run()
+            } else {
+                executor.submit(block)
+            }
+        }
+
+        fun changeThread() {
+            executor.shutdown()
+            executor = newExecutorService()
+        }
+
+        private fun newExecutorService(): ExecutorService =
+            Executors.newSingleThreadExecutor {
+                thread = Thread(it)
+                thread
+            }
+    }
+}