refactor(sim/core): Re-implement SimulationScheduler as Dispatcher

This change updates the `SimulationScheduler` class to implement the
`Dispatcher` interface from the OpenDC Common module, so that OpenDC
modules only need to depend on the common module for dispatching future
task (possibly in simulation).

4 files changed, 96 insertions, 146 deletions

diff --git a/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationBuilders.kt b/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationBuilders.kt
index 882a0fc5..6e568137 100644
--- a/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationBuilders.kt
+++ b/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationBuilders.kt
@@ -22,11 +22,18 @@
 
 package org.opendc.simulator.kotlin
 
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.NonCancellable.children
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.async
-import org.opendc.simulator.SimulationScheduler
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.runBlocking
+import org.opendc.common.DispatcherProvider
+import org.opendc.common.asCoroutineDispatcher
+import org.opendc.simulator.SimulationDispatcher
+import java.time.InstantSource
 import kotlin.coroutines.ContinuationInterceptor
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
@@ -54,16 +61,16 @@ import kotlin.coroutines.EmptyCoroutineContext
  * The simulation is run in a single thread, unless other [CoroutineDispatcher] are used for child coroutines.
  * Because of this, child coroutines are not executed in parallel to [body].
  * In order for the spawned-off asynchronous code to actually be executed, one must either [yield] or suspend the
- * body some other way, or use commands that control scheduling (see [SimulationScheduler]).
+ * body some other way, or use commands that control scheduling (see [SimulationDispatcher]).
  */
 @OptIn(ExperimentalCoroutinesApi::class)
 public fun runSimulation(
     context: CoroutineContext = EmptyCoroutineContext,
-    scheduler: SimulationScheduler = SimulationScheduler(),
+    scheduler: SimulationDispatcher = SimulationDispatcher(),
     body: suspend SimulationCoroutineScope.() -> Unit
 ) {
-    val (safeContext, dispatcher) = context.checkArguments(scheduler)
-    val startingJobs = safeContext.activeJobs()
+    val (safeContext, job, dispatcher) = context.checkArguments(scheduler)
+    val startingJobs = job.activeJobs()
     val scope = SimulationCoroutineScope(safeContext)
     val deferred = scope.async {
         body(scope)
@@ -72,7 +79,7 @@ public fun runSimulation(
     deferred.getCompletionExceptionOrNull()?.let {
         throw it
     }
-    val endingJobs = safeContext.activeJobs()
+    val endingJobs = job.activeJobs()
     if ((endingJobs - startingJobs).isNotEmpty()) {
         throw IllegalStateException("Test finished with active jobs: $endingJobs")
     }
@@ -82,24 +89,51 @@ public fun runSimulation(
  * Convenience method for calling [runSimulation] on an existing [SimulationCoroutineScope].
  */
 public fun SimulationCoroutineScope.runSimulation(block: suspend SimulationCoroutineScope.() -> Unit): Unit =
-    runSimulation(coroutineContext, scheduler, block)
+    runSimulation(coroutineContext, dispatcher, block)
+
+private fun CoroutineContext.checkArguments(scheduler: SimulationDispatcher): Triple<CoroutineContext, Job, SimulationController> {
+    val job = get(Job) ?: SupervisorJob()
+    val dispatcher = get(ContinuationInterceptor) ?: scheduler.asCoroutineDispatcher()
+    val simulationDispatcher = dispatcher.asSimulationDispatcher()
+    return Triple(this + dispatcher + job, job, simulationDispatcher.asController())
+}
+
+private fun Job.activeJobs(): Set<Job> {
+    return children.filter { it.isActive }.toSet()
+}
 
 /**
- * Convenience method for calling [runSimulation] on an existing [SimulationCoroutineDispatcher].
+ * Convert a [ContinuationInterceptor] into a [SimulationDispatcher] if possible.
  */
-public fun SimulationCoroutineDispatcher.runSimulation(block: suspend SimulationCoroutineScope.() -> Unit): Unit =
-    runSimulation(this, scheduler, block)
-
-private fun CoroutineContext.checkArguments(scheduler: SimulationScheduler): Pair<CoroutineContext, SimulationController> {
-    val dispatcher = get(ContinuationInterceptor).run {
-        this?.let { require(this is SimulationController) { "Dispatcher must implement SimulationController: $this" } }
-        this ?: SimulationCoroutineDispatcher(scheduler)
-    }
+internal fun ContinuationInterceptor.asSimulationDispatcher(): SimulationDispatcher {
+    val provider = this as? DispatcherProvider ?: throw IllegalArgumentException(
+        "DispatcherProvider such as SimulatorCoroutineDispatcher as the ContinuationInterceptor(Dispatcher) is required"
+    )
 
-    val job = get(Job) ?: SupervisorJob()
-    return Pair(this + dispatcher + job, dispatcher as SimulationController)
+    return provider.dispatcher as? SimulationDispatcher ?: throw IllegalArgumentException("Active dispatcher is not a SimulationDispatcher")
 }
 
-private fun CoroutineContext.activeJobs(): Set<Job> {
-    return checkNotNull(this[Job]).children.filter { it.isActive }.toSet()
+/**
+ * Helper method to convert a [SimulationDispatcher] into a [SimulationController].
+ */
+internal fun SimulationDispatcher.asController(): SimulationController {
+    return object : SimulationController {
+        override val dispatcher: SimulationDispatcher
+            get() = this@asController
+
+        override val timeSource: InstantSource
+            get() = this@asController.timeSource
+
+        override fun advanceUntilIdle() {
+            dispatcher.advanceUntilIdle()
+        }
+
+        override fun advanceBy(delayMs: Long) {
+            dispatcher.advanceBy(delayMs)
+        }
+
+        override fun runCurrent() {
+            dispatcher.runCurrent()
+        }
+    }
 }
diff --git a/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationController.kt b/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationController.kt
index f96b2326..f7470ad9 100644
--- a/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationController.kt
+++ b/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationController.kt
@@ -23,30 +23,48 @@
 package org.opendc.simulator.kotlin
 
 import kotlinx.coroutines.CoroutineDispatcher
-import org.opendc.simulator.SimulationScheduler
-import java.time.Clock
+import org.opendc.simulator.SimulationDispatcher
+import java.time.InstantSource
 
 /**
- * Control the virtual clock of a [CoroutineDispatcher].
+ * Interface to control the virtual clock of a [CoroutineDispatcher].
  */
 public interface SimulationController {
     /**
      * The current virtual clock as it is known to this Dispatcher.
      */
-    public val clock: Clock
+    public val timeSource: InstantSource
 
     /**
-     * The [SimulationScheduler] driving the simulation.
+     * The current virtual timestamp of the dispatcher (in milliseconds since epoch).
      */
-    public val scheduler: SimulationScheduler
+    public val currentTime: Long
+        get() = timeSource.millis()
+
+    /**
+     * Return the [SimulationDispatcher] driving the simulation.
+     */
+    public val dispatcher: SimulationDispatcher
 
     /**
      * Immediately execute all pending tasks and advance the virtual clock-time to the last delay.
      *
      * If new tasks are scheduled due to advancing virtual time, they will be executed before `advanceUntilIdle`
      * returns.
+     */
+    public fun advanceUntilIdle()
+
+    /**
+     * Move the virtual clock of this dispatcher forward by the specified amount, running the scheduled tasks in the
+     * meantime.
      *
-     * @return the amount of delay-time that this Dispatcher's clock has been forwarded in milliseconds.
+     * @param delayMs The amount of time to move the virtual clock forward (in milliseconds).
+     * @throws IllegalStateException if passed a negative <code>delay</code>.
+     */
+    public fun advanceBy(delayMs: Long)
+
+    /**
+     * Execute the tasks that are scheduled to execute at this moment of virtual time.
      */
-    public fun advanceUntilIdle(): Long
+    public fun runCurrent()
 }
diff --git a/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationCoroutineDispatcher.kt b/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationCoroutineDispatcher.kt
deleted file mode 100644
index cacbbbf7..00000000
--- a/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationCoroutineDispatcher.kt
+++ /dev/null
@@ -1,98 +0,0 @@
-/*
- * Copyright (c) 2021 AtLarge Research
- *
- * Permission is hereby granted, free of charge, to any person obtaining a copy
- * of this software and associated documentation files (the "Software"), to deal
- * in the Software without restriction, including without limitation the rights
- * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
- * copies of the Software, and to permit persons to whom the Software is
- * furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in all
- * copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
- * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
- * SOFTWARE.
- */
-
-package org.opendc.simulator.kotlin
-
-import kotlinx.coroutines.CancellableContinuation
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.Delay
-import kotlinx.coroutines.DisposableHandle
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.InternalCoroutinesApi
-import org.opendc.simulator.SimulationScheduler
-import java.lang.Runnable
-import java.time.Clock
-import kotlin.coroutines.CoroutineContext
-
-/**
- * A [CoroutineDispatcher] that performs both immediate execution of coroutines on the main thread and uses a virtual
- * clock for time management.
- *
- * @param scheduler The [SimulationScheduler] used to manage the execution of future tasks.
- */
-@OptIn(InternalCoroutinesApi::class)
-public class SimulationCoroutineDispatcher(
-    override val scheduler: SimulationScheduler = SimulationScheduler()
-) : CoroutineDispatcher(), SimulationController, Delay {
-    /**
-     * The virtual clock of this dispatcher.
-     */
-    override val clock: Clock = scheduler.clock
-
-    override fun dispatch(context: CoroutineContext, block: Runnable) {
-        block.run()
-    }
-
-    override fun dispatchYield(context: CoroutineContext, block: Runnable) {
-        scheduler.execute(block)
-    }
-
-    @OptIn(ExperimentalCoroutinesApi::class)
-    override fun scheduleResumeAfterDelay(timeMillis: Long, continuation: CancellableContinuation<Unit>) {
-        scheduler.schedule(timeMillis, CancellableContinuationRunnable(continuation) { resumeUndispatched(Unit) })
-    }
-
-    override fun invokeOnTimeout(timeMillis: Long, block: Runnable, context: CoroutineContext): DisposableHandle {
-        return object : DisposableHandle {
-            private val deadline = (scheduler.currentTime + timeMillis).let { if (it >= 0) it else Long.MAX_VALUE }
-            private val id = scheduler.schedule(timeMillis, block)
-
-            override fun dispose() {
-                scheduler.cancel(deadline, id)
-            }
-        }
-    }
-
-    override fun toString(): String {
-        return "SimulationCoroutineDispatcher[time=${scheduler.currentTime}ms]"
-    }
-
-    override fun advanceUntilIdle(): Long {
-        val scheduler = scheduler
-        val oldTime = scheduler.currentTime
-
-        scheduler.advanceUntilIdle()
-
-        return scheduler.currentTime - oldTime
-    }
-
-    /**
-     * This class exists to allow cleanup code to avoid throwing for cancelled continuations scheduled
-     * in the future.
-     */
-    private class CancellableContinuationRunnable<T>(
-        @JvmField val continuation: CancellableContinuation<T>,
-        private val block: CancellableContinuation<T>.() -> Unit
-    ) : Runnable {
-        override fun run() = continuation.block()
-    }
-}
diff --git a/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationCoroutineScope.kt b/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationCoroutineScope.kt
index 6be8e67a..ca49fc53 100644
--- a/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationCoroutineScope.kt
+++ b/opendc-simulator/opendc-simulator-core/src/main/kotlin/org/opendc/simulator/kotlin/SimulationCoroutineScope.kt
@@ -24,7 +24,8 @@ package org.opendc.simulator.kotlin
 
 import kotlinx.coroutines.CoroutineExceptionHandler
 import kotlinx.coroutines.CoroutineScope
-import org.opendc.simulator.SimulationScheduler
+import org.opendc.common.asCoroutineDispatcher
+import org.opendc.simulator.SimulationDispatcher
 import kotlin.coroutines.ContinuationInterceptor
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
@@ -34,33 +35,28 @@ import kotlin.coroutines.EmptyCoroutineContext
  */
 public interface SimulationCoroutineScope : CoroutineScope, SimulationController
 
-private class SimulationCoroutineScopeImpl(
-    override val coroutineContext: CoroutineContext
-) :
-    SimulationCoroutineScope,
-    SimulationController by coroutineContext.simulationController
-
 /**
  * A scope which provides detailed control over the execution of coroutines for simulations.
  *
  * If the provided context does not provide a [ContinuationInterceptor] (Dispatcher) or [CoroutineExceptionHandler], the
- * scope adds [SimulationCoroutineDispatcher] automatically.
+ * scope adds a dispatcher automatically.
  */
-@Suppress("FunctionName")
 public fun SimulationCoroutineScope(
     context: CoroutineContext = EmptyCoroutineContext,
-    scheduler: SimulationScheduler = SimulationScheduler()
+    scheduler: SimulationDispatcher = SimulationDispatcher()
 ): SimulationCoroutineScope {
     var safeContext = context
-    if (context[ContinuationInterceptor] == null) safeContext += SimulationCoroutineDispatcher(scheduler)
-    return SimulationCoroutineScopeImpl(safeContext)
-}
+    val simulationDispatcher: SimulationDispatcher
+    val interceptor = context[ContinuationInterceptor]
 
-private inline val CoroutineContext.simulationController: SimulationController
-    get() {
-        val handler = this[ContinuationInterceptor]
-        return handler as? SimulationController ?: throw IllegalArgumentException(
-            "SimulationCoroutineScope requires a SimulationController such as SimulatorCoroutineDispatcher as " +
-                "the ContinuationInterceptor (Dispatcher)"
-        )
+    if (interceptor != null) {
+        simulationDispatcher = interceptor.asSimulationDispatcher()
+    } else {
+        simulationDispatcher = scheduler
+        safeContext += scheduler.asCoroutineDispatcher()
     }
+
+    return object : SimulationCoroutineScope, SimulationController by simulationDispatcher.asController() {
+        override val coroutineContext: CoroutineContext = safeContext
+    }
+}