Fix FlushCoroutineDispatcher.scheduleResumeAfterDelay to correctly ca…

…ncel the delayed task when the continuation is cancelled. (#596)
JetBrains · Jul 3, 2023 · 53f869f · 53f869f
1 parent 6033a53
commit 53f869f
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 13 deletions.
diff --git a/...ui/ui/src/skikoMain/kotlin/androidx/compose/ui/platform/FlushCoroutineDispatcher.skiko.kt b/...ui/ui/src/skikoMain/kotlin/androidx/compose/ui/platform/FlushCoroutineDispatcher.skiko.kt
@@ -100,16 +100,22 @@ internal class FlushCoroutineDispatcher(
         synchronized(tasksLock) {
             delayedTasks.add(block)
         }
-        scope.launch {
-            kotlinx.coroutines.delay(timeMillis)
-            performRun {
-                val isTaskAlive = synchronized(tasksLock) {
-                    delayedTasks.remove(block)
-                }
-                if (isTaskAlive) {
-                    block.run()
+        val job = scope.launch {
+            try{
+                kotlinx.coroutines.delay(timeMillis)
+            } finally {
+                performRun {
+                    val isTaskAlive = synchronized(tasksLock) {
+                        delayedTasks.remove(block)
+                    }
+                    if (isTaskAlive) {
+                        block.run()
+                    }
                 }
             }
         }
+        continuation.invokeOnCancellation {
+            job.cancel()
+        }
     }
 }
diff --git a/...e/ui/ui/src/skikoTest/kotlin/androidx/compose/ui/platform/FlushCoroutineDispatcherTest.kt b/...e/ui/ui/src/skikoTest/kotlin/androidx/compose/ui/platform/FlushCoroutineDispatcherTest.kt
@@ -16,16 +16,13 @@
 
 package androidx.compose.ui.platform
 
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.runTest
-import kotlinx.coroutines.yield
 import kotlin.random.Random
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFalse
-import kotlinx.coroutines.withContext
+import kotlin.test.assertTrue
+import kotlinx.coroutines.*
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class FlushCoroutineDispatcherTest {
@@ -106,4 +103,19 @@ class FlushCoroutineDispatcherTest {
         assertEquals((0 until 10000).toList(), actualNumbers)
         assertFalse(dispatcher.hasTasks())
     }
+
+    @Test
+    fun delayed_tasks_are_cancelled() = runTest {
+        val coroutineScope = CoroutineScope(Dispatchers.Unconfined)
+        val dispatcher = FlushCoroutineDispatcher(coroutineScope)
+        val job = launch(dispatcher){
+            delay(Long.MAX_VALUE/2)
+        }
+        assertTrue(dispatcher.hasTasks())
+        job.cancel()
+        assertTrue(
+            !dispatcher.hasTasks(),
+            "FlushCoroutineDispatcher has a delayed task that has been cancelled"
+        )
+    }
 }