Skip to content

Commit 8f343fb

Browse files
dkhalanskyjbdkhalanskyjb
committed
Fix VirtualTimeSource not being able to skip delays
1 parent d04fea7 commit 8f343fb

File tree

9 files changed

+60
-49
lines changed

9 files changed

+60
-49
lines changed

‎kotlinx-coroutines-core/jvm/src/AbstractTimeSource.kt

+10-7
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@ internal abstract class AbstractTimeSource {
1010
abstract fun currentTimeMillis(): Long
1111
abstract fun nanoTime(): Long
1212
abstract fun wrapTask(block: Runnable): Runnable
13-
abstract fun trackTask()
14-
abstract fun unTrackTask()
13+
abstract fun trackTask(obj: Any)
14+
abstract fun unTrackTask(obj: Any)
1515
abstract fun registerTimeLoopThread()
1616
abstract fun unregisterTimeLoopThread()
1717
abstract fun parkNanos(blocker: Any, nanos: Long) // should return immediately when nanos <= 0
@@ -80,18 +80,21 @@ internal inline fun wrapTask(block: Runnable): Runnable =
8080
*
8181
* If the second thread is not tracked, the first thread effectively enters a spin loop until the second thread
8282
* physically changes the shared state.
83+
*
84+
* Every call to [trackTask] must be accompanied by a call to [unTrackTask] with the same [obj],
85+
* but [unTrackTask] can be called even if the corresponding [trackTask] wasn't called.
8386
*/
8487
@InlineOnly
85-
internal inline fun trackTask() {
86-
timeSource?.trackTask()
88+
internal inline fun trackTask(obj: Any) {
89+
timeSource?.trackTask(obj)
8790
}
8891

8992
/**
90-
* Decrements the number of tasks not under our control. See [trackTask] for more details.
93+
* Marks the task [obj] as complete. If [obj] wasn't tracked, does nothing. See [trackTask] for more details.
9194
*/
9295
@InlineOnly
93-
internal inline fun unTrackTask() {
94-
timeSource?.unTrackTask()
96+
internal inline fun unTrackTask(obj: Any) {
97+
timeSource?.unTrackTask(obj)
9598
}
9699

97100
/**

‎kotlinx-coroutines-core/jvm/src/DefaultDelay.kt

+5-3
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ package kotlinx.coroutines
33

44
import kotlinx.atomicfu.*
55
import kotlinx.coroutines.internal.*
6+
import kotlinx.coroutines.scheduling.*
7+
import kotlinx.coroutines.scheduling.scheduleBackgroundIoTask
68
import kotlin.coroutines.*
79
import kotlin.time.Duration
810

@@ -78,7 +80,7 @@ private object DefaultDelayImpl : EventLoopImplBase(), Runnable {
7880
Thread.interrupted() // just reset interruption flag
7981
val parkNanos = processNextEvent()
8082
if (parkNanos == Long.MAX_VALUE) break // no more events
81-
parkNanos(this@DefaultDelayImpl, parkNanos)
83+
if (parkNanos > 0) parkNanos(this@DefaultDelayImpl, parkNanos)
8284
}
8385
} finally {
8486
_thread.value = null
@@ -105,15 +107,14 @@ private object DefaultDelayImpl : EventLoopImplBase(), Runnable {
105107
This means that whatever thread is going to be running by the end of this function,
106108
it's going to notice the tasks it's supposed to run.
107109
We can return `null` unconditionally. */
108-
ioView.dispatch(ioView, this)
110+
scheduleBackgroundIoTask(this)
109111
return null
110112
}
111113

112114
fun shutdownForTests(timeout: Duration) {
113115
if (_thread.value != null) {
114116
val end = System.currentTimeMillis() + timeout.inWholeMilliseconds
115117
while (true) {
116-
check(delayedQueueIsEmpty) { "There are tasks in the DefaultExecutor" }
117118
synchronized(this) {
118119
unpark(_thread.value ?: return)
119120
val toWait = end - System.currentTimeMillis()
@@ -156,6 +157,7 @@ private fun defaultDelayRunningUnconfinedLoop(): Nothing {
156157
)
157158
}
158159

160+
159161
/** A view separate from [Dispatchers.IO].
160162
* [Int.MAX_VALUE] instead of `1` to avoid needlessly using the [LimitedDispatcher] machinery. */
161163
private val ioView = Dispatchers.IO.limitedParallelism(Int.MAX_VALUE)

‎kotlinx-coroutines-core/jvm/src/Executors.kt

+1-1
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ internal class ExecutorCoroutineDispatcherImpl(override val executor: Executor)
129129
try {
130130
executor.execute(wrapTask(block))
131131
} catch (e: RejectedExecutionException) {
132-
unTrackTask()
132+
unTrackTask(block)
133133
cancelJobOnRejection(context, e)
134134
rescheduleTaskFromClosedDispatcher(block)
135135
}

‎kotlinx-coroutines-core/jvm/src/scheduling/CoroutineScheduler.kt

+5-3
Original file line numberDiff line numberDiff line change
@@ -391,9 +391,11 @@ internal class CoroutineScheduler(
391391
* - Concurrent [close] that effectively shutdowns the worker thread.
392392
* Used for [yield].
393393
*/
394-
fun dispatch(block: Runnable, taskContext: TaskContext = NonBlockingContext, fair: Boolean = false) {
395-
trackTask() // this is needed for virtual time support
394+
fun dispatch(
395+
block: Runnable, taskContext: TaskContext = NonBlockingContext, fair: Boolean = false, track: Boolean = true
396+
) {
396397
val task = createTask(block, taskContext)
398+
if (track) trackTask(task) // this is needed for virtual time support
397399
val isBlockingTask = task.isBlocking
398400
// Invariant: we increment counter **before** publishing the task
399401
// so executing thread can safely decrement the number of blocking tasks
@@ -588,7 +590,7 @@ internal class CoroutineScheduler(
588590
val thread = Thread.currentThread()
589591
thread.uncaughtExceptionHandler.uncaughtException(thread, e)
590592
} finally {
591-
unTrackTask()
593+
unTrackTask(task)
592594
}
593595
}
594596

‎kotlinx-coroutines-core/jvm/src/scheduling/Dispatcher.kt

+8-4
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,11 @@ private object UnlimitedIoScheduler : CoroutineDispatcher() {
3737

3838
@InternalCoroutinesApi
3939
override fun dispatchYield(context: CoroutineContext, block: Runnable) {
40-
DefaultScheduler.dispatchWithContext(block, BlockingContext, true)
40+
DefaultScheduler.dispatchWithContext(block, BlockingContext, fair = true, track = true)
4141
}
4242

4343
override fun dispatch(context: CoroutineContext, block: Runnable) {
44-
DefaultScheduler.dispatchWithContext(block, BlockingContext, false)
44+
DefaultScheduler.dispatchWithContext(block, BlockingContext, fair = false, track = true)
4545
}
4646

4747
override fun limitedParallelism(parallelism: Int, name: String?): CoroutineDispatcher {
@@ -58,6 +58,10 @@ private object UnlimitedIoScheduler : CoroutineDispatcher() {
5858
}
5959
}
6060

61+
internal fun scheduleBackgroundIoTask(block: Runnable) {
62+
DefaultScheduler.dispatchWithContext(block, BlockingContext, fair = false, track = false)
63+
}
64+
6165
// Dispatchers.IO
6266
internal object DefaultIoScheduler : ExecutorCoroutineDispatcher(), Executor {
6367

@@ -126,8 +130,8 @@ internal open class SchedulerCoroutineDispatcher(
126130
coroutineScheduler.dispatch(block, fair = true)
127131
}
128132

129-
internal fun dispatchWithContext(block: Runnable, context: TaskContext, fair: Boolean) {
130-
coroutineScheduler.dispatch(block, context, fair)
133+
internal fun dispatchWithContext(block: Runnable, context: TaskContext, fair: Boolean, track: Boolean) {
134+
coroutineScheduler.dispatch(block, context, fair = fair, track = track)
131135
}
132136

133137
override fun close() {

‎kotlinx-coroutines-core/jvm/test/VirtualTimeSource.kt

+9-9
Original file line numberDiff line numberDiff line change
@@ -48,30 +48,29 @@ internal class VirtualTimeSource(
4848
@Volatile
4949
private var time: Long = 0
5050

51-
private var trackedTasks = 0
51+
private val trackedTasks = HashSet<Any>()
5252

5353
private val threads = ConcurrentHashMap<Thread, ThreadStatus>()
5454

5555
override fun currentTimeMillis(): Long = TimeUnit.NANOSECONDS.toMillis(time)
5656
override fun nanoTime(): Long = time
5757

5858
override fun wrapTask(block: Runnable): Runnable {
59-
trackTask()
59+
trackTask(block)
6060
return Runnable {
6161
try { block.run() }
62-
finally { unTrackTask() }
62+
finally { unTrackTask(block) }
6363
}
6464
}
6565

6666
@Synchronized
67-
override fun trackTask() {
68-
trackedTasks++
67+
override fun trackTask(obj: Any) {
68+
trackedTasks.add(obj)
6969
}
7070

7171
@Synchronized
72-
override fun unTrackTask() {
73-
assert(trackedTasks > 0)
74-
trackedTasks--
72+
override fun unTrackTask(obj: Any) {
73+
trackedTasks.remove(obj)
7574
}
7675

7776
@Synchronized
@@ -125,7 +124,7 @@ internal class VirtualTimeSource(
125124
return
126125
}
127126
if (threads[mainThread] == null) return
128-
if (trackedTasks != 0) return
127+
if (trackedTasks.isNotEmpty()) return
129128
val minParkedTill = minParkedTill()
130129
if (minParkedTill <= time) return
131130
time = minParkedTill
@@ -145,6 +144,7 @@ internal class VirtualTimeSource(
145144
isShutdown = true
146145
wakeupAll()
147146
while (!threads.isEmpty()) (this as Object).wait()
147+
assert(trackedTasks.isEmpty()) { "There are still tracked tasks: $trackedTasks" }
148148
}
149149

150150
private fun wakeupAll() {

‎kotlinx-coroutines-core/jvm/test/channels/TickerChannelCommonTest.kt

+9-9
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,12 @@ import org.junit.runners.*
99
import kotlin.test.*
1010

1111
@RunWith(Parameterized::class)
12-
class TickerChannelCommonTest(private val channelFactory: Channel) : TestBase(disableOutCheck = true) {
12+
class TickerChannelCommonTest(private val channelFactory: Channel) : TestBase() {
1313
companion object {
1414
@Parameterized.Parameters(name = "{0}")
1515
@JvmStatic
1616
fun params(): Collection<Array<Any>> =
17-
Channel.values().map { arrayOf<Any>(it) }
17+
Channel.entries.map { arrayOf<Any>(it) }
1818
}
1919

2020
enum class Channel {
@@ -35,13 +35,13 @@ class TickerChannelCommonTest(private val channelFactory: Channel) : TestBase(di
3535
fun testDelay() = withVirtualTimeSource {
3636
runTest {
3737
val delayChannel = channelFactory(delay = 10000)
38-
delayChannel.checkNotEmpty()
38+
delayChannel.receiveSingle()
3939
delayChannel.checkEmpty()
4040

4141
delay(5000)
4242
delayChannel.checkEmpty()
4343
delay(5100)
44-
delayChannel.checkNotEmpty()
44+
delayChannel.receiveSingle()
4545

4646
delayChannel.cancel()
4747
delay(5100)
@@ -57,13 +57,13 @@ class TickerChannelCommonTest(private val channelFactory: Channel) : TestBase(di
5757
delay(500)
5858
delayChannel.checkEmpty()
5959
delay(300)
60-
delayChannel.checkNotEmpty()
60+
delayChannel.receiveSingle()
6161

6262
// Regular delay
6363
delay(750)
6464
delayChannel.checkEmpty()
6565
delay(260)
66-
delayChannel.checkNotEmpty()
66+
delayChannel.receiveSingle()
6767
delayChannel.cancel()
6868
}
6969
}
@@ -72,7 +72,7 @@ class TickerChannelCommonTest(private val channelFactory: Channel) : TestBase(di
7272
fun testReceive() = withVirtualTimeSource {
7373
runTest {
7474
val delayChannel = channelFactory(delay = 1000)
75-
delayChannel.checkNotEmpty()
75+
delayChannel.receiveSingle()
7676
var value = withTimeoutOrNull(750) {
7777
delayChannel.receive()
7878
1
@@ -158,7 +158,7 @@ class TickerChannelCommonTest(private val channelFactory: Channel) : TestBase(di
158158

159159
fun ReceiveChannel<Unit>.checkEmpty() = assertNull(tryReceive().getOrNull())
160160

161-
fun ReceiveChannel<Unit>.checkNotEmpty() {
162-
assertNotNull(tryReceive().getOrNull())
161+
suspend fun ReceiveChannel<Unit>.receiveSingle() {
162+
receive()
163163
assertNull(tryReceive().getOrNull())
164164
}

‎kotlinx-coroutines-core/jvm/test/channels/TickerChannelTest.kt

+11-11
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,15 @@ class TickerChannelTest : TestBase() {
99
fun testFixedDelayChannelBackpressure() = withVirtualTimeSource {
1010
runTest {
1111
val delayChannel = ticker(delayMillis = 1000, initialDelayMillis = 0, mode = TickerMode.FIXED_DELAY)
12-
delayChannel.checkNotEmpty()
12+
delayChannel.receiveSingle()
1313
delayChannel.checkEmpty()
1414

1515
delay(1500)
16-
delayChannel.checkNotEmpty()
16+
delayChannel.receiveSingle()
1717
delay(500)
1818
delayChannel.checkEmpty()
1919
delay(520)
20-
delayChannel.checkNotEmpty()
20+
delayChannel.receiveSingle()
2121
delayChannel.cancel()
2222
}
2323
}
@@ -26,17 +26,17 @@ class TickerChannelTest : TestBase() {
2626
fun testDelayChannelBackpressure() = withVirtualTimeSource {
2727
runTest {
2828
val delayChannel = ticker(delayMillis = 1000, initialDelayMillis = 0)
29-
delayChannel.checkNotEmpty()
29+
delayChannel.receiveSingle()
3030
delayChannel.checkEmpty()
3131

3232
delay(1500)
33-
delayChannel.checkNotEmpty()
33+
delayChannel.receiveSingle()
3434
delay(520)
35-
delayChannel.checkNotEmpty()
35+
delayChannel.receiveSingle()
3636
delay(500)
3737
delayChannel.checkEmpty()
3838
delay(520)
39-
delayChannel.checkNotEmpty()
39+
delayChannel.receiveSingle()
4040
delayChannel.cancel()
4141
}
4242
}
@@ -45,17 +45,17 @@ class TickerChannelTest : TestBase() {
4545
fun testDelayChannelBackpressure2() = withVirtualTimeSource {
4646
runTest {
4747
val delayChannel = ticker(delayMillis = 200, initialDelayMillis = 0)
48-
delayChannel.checkNotEmpty()
48+
delayChannel.receiveSingle()
4949
delayChannel.checkEmpty()
5050

5151
delay(500)
52-
delayChannel.checkNotEmpty()
52+
delayChannel.receiveSingle()
5353
delay(110)
54-
delayChannel.checkNotEmpty()
54+
delayChannel.receiveSingle()
5555
delay(110)
5656
delayChannel.checkEmpty()
5757
delay(110)
58-
delayChannel.checkNotEmpty()
58+
delayChannel.receiveSingle()
5959
delayChannel.cancel()
6060
}
6161
}

‎kotlinx-coroutines-core/jvm/test/scheduling/SchedulerTestBase.kt

+2-2
Original file line numberDiff line numberDiff line change
@@ -101,11 +101,11 @@ internal fun SchedulerCoroutineDispatcher.blocking(parallelism: Int = 16): Corou
101101

102102
@InternalCoroutinesApi
103103
override fun dispatchYield(context: CoroutineContext, block: Runnable) {
104-
this@blocking.dispatchWithContext(block, BlockingContext, true)
104+
this@blocking.dispatchWithContext(block, BlockingContext, fair = true, track = true)
105105
}
106106

107107
override fun dispatch(context: CoroutineContext, block: Runnable) {
108-
this@blocking.dispatchWithContext(block, BlockingContext, false)
108+
this@blocking.dispatchWithContext(block, BlockingContext, fair = false, track = true)
109109
}
110110
}.limitedParallelism(parallelism)
111111
}

0 commit comments

Comments
 (0)