Fix 'runCurrent' executing newly-added tasks

Kotlin · Oct 15, 2021 · 50dfc7b · 50dfc7b
1 parent 201053d
commit 50dfc7b
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 55 deletions.
diff --git a/kotlinx-coroutines-test/common/src/TestCoroutineScheduler.kt b/kotlinx-coroutines-test/common/src/TestCoroutineScheduler.kt
@@ -41,6 +41,7 @@ public class TestCoroutineScheduler: AbstractCoroutineContextElement(TestCorouti
     /** The current virtual time. */
     @ExperimentalCoroutinesApi
     public var currentTime: Long = 0
+        get() = synchronized(lock) { field }
         private set
 
     /**
@@ -96,19 +97,10 @@ public class TestCoroutineScheduler: AbstractCoroutineContextElement(TestCorouti
      */
     @ExperimentalCoroutinesApi
     public fun runCurrent() {
+        val timeMark = synchronized(lock) { currentTime }
         while (true) {
             val event = synchronized(lock) {
-                val timeMark = currentTime
-                val event = events.peek() ?: return
-                when {
-                    timeMark > event.time -> currentTimeAheadOfEvents()
-                    timeMark < event.time -> return
-                    else -> {
-                        val event2 = events.removeFirstOrNull()
-                        if (event !== event2) concurrentModificationUnderLock()
-                        event2
-                    }
-                }
+                events.removeFirstIf { it.time <= timeMark } ?: return
             }
             event.dispatcher.processEvent(event.time, event.marker)
         }

diff --git a/kotlinx-coroutines-test/common/test/TestCoroutineSchedulerTest.kt b/kotlinx-coroutines-test/common/test/TestCoroutineSchedulerTest.kt
@@ -10,33 +10,26 @@ import kotlin.test.*
 class TestCoroutineSchedulerTest {
     /** Tests that `TestCoroutineScheduler` attempts to detect if there are several instances of it. */
     @Test
-    fun testContextElement() {
-        runBlockingTest {
-            assertFailsWith<IllegalStateException> {
-                withContext(TestCoroutineDispatcher()) {
-                }
+    fun testContextElement() = runBlockingTest {
+        assertFailsWith<IllegalStateException> {
+            withContext(TestCoroutineDispatcher()) {
             }
         }
     }
 
     /** Tests that, as opposed to [DelayController.advanceTimeBy] or [TestCoroutineScope.advanceTimeBy],
      * [TestCoroutineScheduler.advanceTimeBy] doesn't run the tasks scheduled at the target moment. */
     @Test
-    fun testAdvanceTimeByDoesNotRunCurrent() {
-        val dispatcher = TestCoroutineDispatcher()
-        dispatcher.runBlockingTest {
-            dispatcher.pauseDispatcher {
-                var entered = false
-                launch {
-                    delay(15)
-                    entered = true
-                }
-                testScheduler.advanceTimeBy(15)
-                assertFalse(entered)
-                testScheduler.runCurrent()
-                assertTrue(entered)
-            }
+    fun testAdvanceTimeByDoesNotRunCurrent() = runBlockingTest {
+        var entered = false
+        launch {
+            delay(15)
+            entered = true
         }
+        testScheduler.advanceTimeBy(15)
+        assertFalse(entered)
+        testScheduler.runCurrent()
+        assertTrue(entered)
     }
 
     /** Tests that [TestCoroutineScheduler.advanceTimeBy] doesn't accept negative delays. */
@@ -51,32 +44,78 @@ class TestCoroutineSchedulerTest {
     /** Tests that if [TestCoroutineScheduler.advanceTimeBy] encounters an arithmetic overflow, all the tasks scheduled
      * until the moment [Long.MAX_VALUE] get run. */
     @Test
-    fun testAdvanceTimeByEnormousDelays() {
-        val dispatcher = TestCoroutineDispatcher()
-        dispatcher.runBlockingTest {
-            dispatcher.pauseDispatcher {
-                val initialDelay = 10L
-                delay(initialDelay)
-                assertEquals(initialDelay, currentTime)
-                var enteredInfinity = false
-                launch {
-                    delay(Long.MAX_VALUE - 1) // delay(Long.MAX_VALUE) does nothing
-                    assertEquals(Long.MAX_VALUE, currentTime)
-                    enteredInfinity = true
-                }
-                var enteredNearInfinity = false
-                launch {
-                    delay(Long.MAX_VALUE - initialDelay - 1)
-                    assertEquals(Long.MAX_VALUE - 1, currentTime)
-                    enteredNearInfinity = true
-                }
-                testScheduler.advanceTimeBy(Long.MAX_VALUE)
-                assertFalse(enteredInfinity)
-                assertTrue(enteredNearInfinity)
-                assertEquals(Long.MAX_VALUE, currentTime)
-                testScheduler.runCurrent()
-                assertTrue(enteredInfinity)
+    fun testAdvanceTimeByEnormousDelays() = runBlockingTest {
+        val initialDelay = 10L
+        delay(initialDelay)
+        assertEquals(initialDelay, currentTime)
+        var enteredInfinity = false
+        launch {
+            delay(Long.MAX_VALUE - 1) // delay(Long.MAX_VALUE) does nothing
+            assertEquals(Long.MAX_VALUE, currentTime)
+            enteredInfinity = true
+        }
+        var enteredNearInfinity = false
+        launch {
+            delay(Long.MAX_VALUE - initialDelay - 1)
+            assertEquals(Long.MAX_VALUE - 1, currentTime)
+            enteredNearInfinity = true
+        }
+        testScheduler.advanceTimeBy(Long.MAX_VALUE)
+        assertFalse(enteredInfinity)
+        assertTrue(enteredNearInfinity)
+        assertEquals(Long.MAX_VALUE, currentTime)
+        testScheduler.runCurrent()
+        assertTrue(enteredInfinity)
+    }
+
+    /** Tests the basic functionality of [TestCoroutineScheduler.advanceTimeBy]. */
+    @Test
+    fun testAdvanceTimeBy() {
+        val scheduler = TestCoroutineScheduler()
+        val scope = TestCoroutineScope(scheduler)
+        var stage = 1
+        scope.launch {
+            delay(1_000)
+            assertEquals(1_000, scheduler.currentTime)
+            stage = 2
+            delay(500)
+            assertEquals(1_500, scheduler.currentTime)
+            stage = 3
+            delay(501)
+            assertEquals(2_001, scheduler.currentTime)
+            stage = 4
+        }
+        assertEquals(1, stage)
+        assertEquals(0, scheduler.currentTime)
+        scheduler.advanceTimeBy(2_000)
+        assertEquals(3, stage)
+        assertEquals(2_000, scheduler.currentTime)
+        scheduler.advanceTimeBy(2)
+        assertEquals(4, stage)
+        assertEquals(2_002, scheduler.currentTime)
+        scope.cleanupTestCoroutines()
+    }
+
+    /** Tests that [TestCoroutineScheduler.runCurrent] will not run new tasks after the current time has advanced. */
+    @Test
+    fun testRunCurrentNotDrainingQueue() {
+        val scheduler = TestCoroutineScheduler()
+        val scope = TestCoroutineScope(scheduler)
+        var stage = 1
+        scope.launch {
+            delay(1)
+            launch {
+                delay(1)
+                stage = 3
             }
+            scheduler.advanceTimeBy(1)
+            stage = 2
         }
+        scheduler.advanceTimeBy(1)
+        assertEquals(1, stage)
+        scheduler.runCurrent()
+        assertEquals(2, stage)
+        scheduler.runCurrent()
+        assertEquals(3, stage)
     }
 }