Fix VirtualTimeSource not being able to skip delays

Author: dkhalanskyjb

Diff for: kotlinx-coroutines-core/jvm/src/AbstractTimeSource.kt

@@ -10,8 +10,8 @@ internal abstract class AbstractTimeSource {
     abstract fun currentTimeMillis(): Long
     abstract fun nanoTime(): Long
     abstract fun wrapTask(block: Runnable): Runnable
-    abstract fun trackTask()
-    abstract fun unTrackTask()
+    abstract fun trackTask(obj: Any)
+    abstract fun unTrackTask(obj: Any)
     abstract fun registerTimeLoopThread()
     abstract fun unregisterTimeLoopThread()
     abstract fun parkNanos(blocker: Any, nanos: Long) // should return immediately when nanos <= 0
@@ -80,18 +80,21 @@ internal inline fun wrapTask(block: Runnable): Runnable =
  *
  * If the second thread is not tracked, the first thread effectively enters a spin loop until the second thread
  * physically changes the shared state.
+ *
+ * Every call to [trackTask] must be accompanied by a call to [unTrackTask] with the same [obj],
+ * but [unTrackTask] can be called even if the corresponding [trackTask] wasn't called.
  */
 @InlineOnly
-internal inline fun trackTask() {
-    timeSource?.trackTask()
+internal inline fun trackTask(obj: Any) {
+    timeSource?.trackTask(obj)
 }
 
 /**
- * Decrements the number of tasks not under our control. See [trackTask] for more details.
+ * Marks the task [obj] as complete. If [obj] wasn't tracked, does nothing. See [trackTask] for more details.
  */
 @InlineOnly
-internal inline fun unTrackTask() {
-    timeSource?.unTrackTask()
+internal inline fun unTrackTask(obj: Any) {
+    timeSource?.unTrackTask(obj)
 }
 
 /**

Diff for: kotlinx-coroutines-core/jvm/src/DefaultDelay.kt

@@ -3,6 +3,8 @@ package kotlinx.coroutines
 
 import kotlinx.atomicfu.*
 import kotlinx.coroutines.internal.*
+import kotlinx.coroutines.scheduling.*
+import kotlinx.coroutines.scheduling.scheduleBackgroundIoTask
 import kotlin.coroutines.*
 import kotlin.time.Duration
 
@@ -105,7 +107,7 @@ private object DefaultDelayImpl : EventLoopImplBase(), Runnable {
         This means that whatever thread is going to be running by the end of this function,
         it's going to notice the tasks it's supposed to run.
         We can return `null` unconditionally. */
-        ioView.dispatch(ioView, this)
+        scheduleBackgroundIoTask(this)
         return null
     }
 
@@ -156,6 +158,7 @@ private fun defaultDelayRunningUnconfinedLoop(): Nothing {
     )
 }
 
+
 /** A view separate from [Dispatchers.IO].
  * [Int.MAX_VALUE] instead of `1` to avoid needlessly using the [LimitedDispatcher] machinery. */
 private val ioView = Dispatchers.IO.limitedParallelism(Int.MAX_VALUE)

Diff for: kotlinx-coroutines-core/jvm/src/Executors.kt

@@ -129,7 +129,7 @@ internal class ExecutorCoroutineDispatcherImpl(override val executor: Executor)
         try {
             executor.execute(wrapTask(block))
         } catch (e: RejectedExecutionException) {
-            unTrackTask()
+            unTrackTask(block)
             cancelJobOnRejection(context, e)
             rescheduleTaskFromClosedDispatcher(block)
         }

Diff for: kotlinx-coroutines-core/jvm/src/scheduling/CoroutineScheduler.kt

@@ -391,9 +391,11 @@ internal class CoroutineScheduler(
      *   - Concurrent [close] that effectively shutdowns the worker thread.
      * Used for [yield].
      */
-    fun dispatch(block: Runnable, taskContext: TaskContext = NonBlockingContext, fair: Boolean = false) {
-        trackTask() // this is needed for virtual time support
+    fun dispatch(
+        block: Runnable, taskContext: TaskContext = NonBlockingContext, fair: Boolean = false, track: Boolean = true
+    ) {
         val task = createTask(block, taskContext)
+        if (track) trackTask(task) // this is needed for virtual time support
         val isBlockingTask = task.isBlocking
         // Invariant: we increment counter **before** publishing the task
         // so executing thread can safely decrement the number of blocking tasks
@@ -588,7 +590,7 @@ internal class CoroutineScheduler(
             val thread = Thread.currentThread()
             thread.uncaughtExceptionHandler.uncaughtException(thread, e)
         } finally {
-            unTrackTask()
+            unTrackTask(task)
         }
     }
 

Diff for: kotlinx-coroutines-core/jvm/src/scheduling/Dispatcher.kt

@@ -37,11 +37,11 @@ private object UnlimitedIoScheduler : CoroutineDispatcher() {
 
     @InternalCoroutinesApi
     override fun dispatchYield(context: CoroutineContext, block: Runnable) {
-        DefaultScheduler.dispatchWithContext(block, BlockingContext, true)
+        DefaultScheduler.dispatchWithContext(block, BlockingContext, fair = true, track = true)
     }
 
     override fun dispatch(context: CoroutineContext, block: Runnable) {
-        DefaultScheduler.dispatchWithContext(block, BlockingContext, false)
+        DefaultScheduler.dispatchWithContext(block, BlockingContext, fair = false, track = true)
     }
 
     override fun limitedParallelism(parallelism: Int, name: String?): CoroutineDispatcher {
@@ -58,6 +58,10 @@ private object UnlimitedIoScheduler : CoroutineDispatcher() {
     }
 }
 
+internal fun scheduleBackgroundIoTask(block: Runnable) {
+    DefaultScheduler.dispatchWithContext(block, BlockingContext, fair = false, track = false)
+}
+
 // Dispatchers.IO
 internal object DefaultIoScheduler : ExecutorCoroutineDispatcher(), Executor {
 
@@ -126,8 +130,8 @@ internal open class SchedulerCoroutineDispatcher(
         coroutineScheduler.dispatch(block, fair = true)
     }
 
-    internal fun dispatchWithContext(block: Runnable, context: TaskContext, fair: Boolean) {
-        coroutineScheduler.dispatch(block, context, fair)
+    internal fun dispatchWithContext(block: Runnable, context: TaskContext, fair: Boolean, track: Boolean) {
+        coroutineScheduler.dispatch(block, context, fair = fair, track = track)
     }
 
     override fun close() {

Diff for: kotlinx-coroutines-core/jvm/test/VirtualTimeSource.kt

@@ -48,30 +48,29 @@ internal class VirtualTimeSource(
     @Volatile
     private var time: Long = 0
 
-    private var trackedTasks = 0
+    private val trackedTasks = HashSet<Any>()
 
     private val threads = ConcurrentHashMap<Thread, ThreadStatus>()
 
     override fun currentTimeMillis(): Long = TimeUnit.NANOSECONDS.toMillis(time)
     override fun nanoTime(): Long = time
 
     override fun wrapTask(block: Runnable): Runnable {
-        trackTask()
+        trackTask(block)
         return Runnable {
             try { block.run() }
-            finally { unTrackTask() }
+            finally { unTrackTask(block) }
         }
     }
 
     @Synchronized
-    override fun trackTask() {
-        trackedTasks++
+    override fun trackTask(obj: Any) {
+        trackedTasks.add(obj)
     }
 
     @Synchronized
-    override fun unTrackTask() {
-        assert(trackedTasks > 0)
-        trackedTasks--
+    override fun unTrackTask(obj: Any) {
+        trackedTasks.remove(obj)
     }
 
     @Synchronized
@@ -125,7 +124,7 @@ internal class VirtualTimeSource(
             return
         }
         if (threads[mainThread] == null) return
-        if (trackedTasks != 0) return
+        if (trackedTasks.isNotEmpty()) return
         val minParkedTill = minParkedTill()
         if (minParkedTill <= time) return
         time = minParkedTill
@@ -145,6 +144,7 @@ internal class VirtualTimeSource(
         isShutdown = true
         wakeupAll()
         while (!threads.isEmpty()) (this as Object).wait()
+        assert(trackedTasks.isEmpty()) { "There are still tracked tasks: $trackedTasks" }
     }
 
     private fun wakeupAll() {

Diff for: kotlinx-coroutines-core/jvm/test/channels/TickerChannelCommonTest.kt

@@ -9,7 +9,7 @@ import org.junit.runners.*
 import kotlin.test.*
 
 @RunWith(Parameterized::class)
-class TickerChannelCommonTest(private val channelFactory: Channel) : TestBase(disableOutCheck = true) {
+class TickerChannelCommonTest(private val channelFactory: Channel) : TestBase() {
     companion object {
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic

Diff for: kotlinx-coroutines-core/jvm/test/scheduling/SchedulerTestBase.kt

@@ -101,11 +101,11 @@ internal fun SchedulerCoroutineDispatcher.blocking(parallelism: Int = 16): Corou
 
         @InternalCoroutinesApi
         override fun dispatchYield(context: CoroutineContext, block: Runnable) {
-            this@blocking.dispatchWithContext(block, BlockingContext, true)
+            this@blocking.dispatchWithContext(block, BlockingContext, fair = true, track = true)
         }
 
         override fun dispatch(context: CoroutineContext, block: Runnable) {
-            this@blocking.dispatchWithContext(block, BlockingContext, false)
+            this@blocking.dispatchWithContext(block, BlockingContext, fair = false, track = true)
         }
     }.limitedParallelism(parallelism)
 }