|
| 1 | +/* |
| 2 | + * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license. |
| 3 | + */ |
| 4 | + |
| 5 | +package kotlinx.coroutines |
| 6 | + |
| 7 | +import kotlin.coroutines.* |
| 8 | +import kotlin.jvm.* |
| 9 | + |
| 10 | +private class VirtualTimeDispatcher(enclosingScope: CoroutineScope) : CoroutineDispatcher(), Delay { |
| 11 | + |
| 12 | + private val originalDispatcher = enclosingScope.coroutineContext[ContinuationInterceptor] as CoroutineDispatcher |
| 13 | + private val heap = ArrayList<TimedTask>() // TODO use MPP heap/ordered set implementation |
| 14 | + private var currentTime = 0L |
| 15 | + |
| 16 | + init { |
| 17 | + /* |
| 18 | + * Launch "event-loop-owning" task on start of the virtual time event loop. |
| 19 | + * It ensures the progress of the enclosing event-loop and polls the timed queue |
| 20 | + * when the enclosing event loop is empty, emulating virtual time. |
| 21 | + */ |
| 22 | + enclosingScope.launch { |
| 23 | + while (true) { |
| 24 | + val secret = ThreadLocalEventLoop.currentOrNull()?.processNextEvent() |
| 25 | + ?: error("Event loop is missing, virtual time source works only as part of event loop") |
| 26 | + if (secret <= 0) continue |
| 27 | + if (secret > 0 && secret != Long.MAX_VALUE) error("Unexpected external delay: $secret") |
| 28 | + val nextTask = heap.minBy { it.deadline } ?: return@launch |
| 29 | + heap.remove(nextTask) |
| 30 | + currentTime = nextTask.deadline |
| 31 | + nextTask.run() |
| 32 | + } |
| 33 | + } |
| 34 | + } |
| 35 | + |
| 36 | + private inner class TimedTask( |
| 37 | + private val runnable: Runnable, |
| 38 | + @JvmField val deadline: Long |
| 39 | + ) : DisposableHandle, Runnable by runnable { |
| 40 | + |
| 41 | + override fun dispose() { |
| 42 | + heap.remove(this) |
| 43 | + } |
| 44 | + } |
| 45 | + |
| 46 | + override fun dispatch(context: CoroutineContext, block: Runnable) { |
| 47 | + originalDispatcher.dispatch(context, block) |
| 48 | + } |
| 49 | + |
| 50 | + @ExperimentalCoroutinesApi |
| 51 | + override fun isDispatchNeeded(context: CoroutineContext): Boolean = originalDispatcher.isDispatchNeeded(context) |
| 52 | + |
| 53 | + override fun invokeOnTimeout(timeMillis: Long, block: Runnable): DisposableHandle { |
| 54 | + val task = TimedTask(block, currentTime + timeMillis) |
| 55 | + heap += task |
| 56 | + return task |
| 57 | + } |
| 58 | + |
| 59 | + override fun scheduleResumeAfterDelay(timeMillis: Long, continuation: CancellableContinuation<Unit>) { |
| 60 | + val task = TimedTask(Runnable { with(continuation) { resumeUndispatched(Unit) } }, currentTime + timeMillis) |
| 61 | + heap += task |
| 62 | + continuation.invokeOnCancellation { task.dispose() } |
| 63 | + } |
| 64 | +} |
| 65 | + |
| 66 | +/** |
| 67 | + * Runs a test ([TestBase.runTest]) with a virtual time source. |
| 68 | + * This runner has the following constraints: |
| 69 | + * 1) It works only in the event-loop environment and it is relying on it. |
| 70 | + * None of the coroutines should be launched in any dispatcher different from a current |
| 71 | + * 2) Regular tasks always dominate delayed ones. It means that |
| 72 | + * `launch { while(true) yield() }` will block the progress of the delayed tasks |
| 73 | + * 3) [TestBase.finish] should always be invoked. |
| 74 | + * Given all the constraints into account, it is easy to mess up with a test and actually |
| 75 | + * return from [withVirtualTime] before the test is executed completely. |
| 76 | + * To decrease the probability of such error, additional `finish` constraint is added. |
| 77 | + */ |
| 78 | +public fun TestBase.withVirtualTime(block: suspend CoroutineScope.() -> Unit) = runTest { |
| 79 | + withContext(Dispatchers.Unconfined) { |
| 80 | + // Create a platform-independent event loop |
| 81 | + val dispatcher = VirtualTimeDispatcher(this) |
| 82 | + withContext(dispatcher) { block() } |
| 83 | + ensureFinished() |
| 84 | + } |
| 85 | +} |
0 commit comments