Merge "Support wrapping AndroidUiDispatcher" into androidx-main

Author: Treehugger Robot

compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/platform/AndroidUiDispatcherTest.kt

@@ -32,10 +32,18 @@ import androidx.test.filters.FlakyTest
 import androidx.test.filters.LargeTest
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.uiautomator.UiDevice
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.coroutines.ContinuationInterceptor
+import kotlin.coroutines.CoroutineContext
+import kotlin.test.fail
 import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Runnable
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withContext
 import kotlinx.coroutines.withTimeoutOrNull
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertNotEquals
@@ -174,4 +182,60 @@ class AndroidUiDispatcherTest {
             }
         )
     }
+
+    /**
+     * Test that an AndroidUiDispatcher can be wrapped by another ContinuationInterceptor
+     * without breaking the MonotonicFrameClock's ability to coordinate with its
+     * original dispatcher.
+     *
+     * Construct a situation where the Choreographer contains three frame callbacks:
+     * 1) checkpoint 1
+     * 2) the AndroidUiDispatcher awaiting-frame callback
+     * 3) checkpoint 2
+     * Confirm that a call to withFrameNanos made *after* these three frame callbacks
+     * are enqueued runs *before* checkpoint 2, indicating that it ran with the original
+     * dispatcher's awaiting-frame callback, even though we wrapped the dispatcher.
+     */
+    @Test
+    fun wrappedDispatcherPostsToDispatcherFrameClock() = runBlocking(Dispatchers.Main) {
+        val uiDispatcherContext = AndroidUiDispatcher.Main
+        val uiDispatcher = uiDispatcherContext[ContinuationInterceptor] as CoroutineDispatcher
+        val wrapperDispatcher = object : CoroutineDispatcher() {
+            override fun dispatch(context: CoroutineContext, block: Runnable) {
+                uiDispatcher.dispatch(context, block)
+            }
+        }
+
+        val choreographer = Choreographer.getInstance()
+
+        val expectCount = AtomicInteger(1)
+        fun expect(value: Int) {
+            while (true) {
+                val old = expectCount.get()
+                if (old != value) fail("expected sequence $old but encountered $value")
+                if (expectCount.compareAndSet(value, value + 1)) break
+            }
+        }
+
+        choreographer.postFrameCallback {
+            expect(1)
+        }
+
+        launch(uiDispatcherContext, start = CoroutineStart.UNDISPATCHED) {
+            withFrameNanos {
+                expect(2)
+            }
+        }
+
+        choreographer.postFrameCallback {
+            expect(4)
+        }
+
+        withContext(uiDispatcherContext + wrapperDispatcher) {
+            withFrameNanos {
+                expect(3)
+            }
+            expect(5)
+        }
+    }
 }

compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/AndroidUiDispatcher.android.kt

@@ -131,7 +131,7 @@ class AndroidUiDispatcher private constructor(
      * A [MonotonicFrameClock] associated with this [AndroidUiDispatcher]'s [choreographer]
      * that may be used to await [Choreographer] frame dispatch.
      */
-    val frameClock: MonotonicFrameClock = AndroidUiFrameClock(choreographer)
+    val frameClock: MonotonicFrameClock = AndroidUiFrameClock(choreographer, this)
 
     override fun dispatch(context: CoroutineContext, block: Runnable) {
         synchronized(lock) {

compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/AndroidUiFrameClock.android.kt

@@ -21,13 +21,20 @@ import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlin.coroutines.ContinuationInterceptor
 import kotlin.coroutines.coroutineContext
 
-class AndroidUiFrameClock(
-    val choreographer: Choreographer
+class AndroidUiFrameClock internal constructor(
+    val choreographer: Choreographer,
+    private val dispatcher: AndroidUiDispatcher?
 ) : androidx.compose.runtime.MonotonicFrameClock {
+
+    constructor(
+        choreographer: Choreographer
+    ) : this(choreographer, null)
+
     override suspend fun <R> withFrameNanos(
         onFrame: (Long) -> R
     ): R {
-        val uiDispatcher = coroutineContext[ContinuationInterceptor] as? AndroidUiDispatcher
+        val uiDispatcher = dispatcher
+            ?: coroutineContext[ContinuationInterceptor] as? AndroidUiDispatcher
         return suspendCancellableCoroutine { co ->
             // Important: this callback won't throw, and AndroidUiDispatcher counts on it.
             val callback = Choreographer.FrameCallback { frameTimeNanos ->