refactor(#18): Provide access to process context in nested calls

This change provides a method in the standard library to access the
process context in nested suspending function calls.

5 files changed, 82 insertions, 9 deletions

diff --git a/opendc-core/src/main/kotlin/com/atlarge/opendc/simulator/Context.kt b/opendc-core/src/main/kotlin/com/atlarge/opendc/simulator/Context.kt
index 13170256..24f87eff 100644
--- a/opendc-core/src/main/kotlin/com/atlarge/opendc/simulator/Context.kt
+++ b/opendc-core/src/main/kotlin/com/atlarge/opendc/simulator/Context.kt
@@ -25,6 +25,7 @@
 package com.atlarge.opendc.simulator
 
 import java.util.*
+import kotlin.coroutines.experimental.CoroutineContext
 
 /**
  * This interface provides a context for simulation of [Entity] instances, by defining the environment in which the
@@ -33,7 +34,7 @@ import java.util.*
  *
  * @author Fabian Mastenbroek (f.s.mastenbroek@student.tudelft.nl)
  */
-interface Context<S, out M> {
+interface Context<S, M> : CoroutineContext.Element {
     /**
      * The model of simulation in which the entity exists.
      */
@@ -51,6 +52,11 @@ interface Context<S, out M> {
     val delta: Duration
 
     /**
+     * The [Entity] associated with this context.
+     */
+    val self: Entity<S, M>
+
+    /**
      * The sender of the last received message or `null` in case the process has not received any messages yet.
      *
      * Note that this property is only guaranteed to be correct when accessing after a single suspending call. Methods
@@ -71,12 +77,22 @@ interface Context<S, out M> {
     /**
      * Interrupt an [Entity] process in simulation.
      *
+     * @see [Entity.interrupt(Interrupt)]
+     * @param reason The reason for interrupting the entity.
+     */
+    suspend fun Entity<*, *>.interrupt(reason: String) = interrupt(Interrupt(reason))
+
+    /**
+     * Interrupt an [Entity] process in simulation.
+     *
      * If an [Entity] process has been suspended, the suspending call will throw an [Interrupt] object as a result of
      * this call.
      * Make sure the [Entity] process actually has error handling in place, so it won't take down the whole [Entity]
      * process as result of the interrupt.
+     *
+     * @param interrupt The interrupt to throw at the entity.
      */
-    suspend fun Entity<*, *>.interrupt()
+    suspend fun Entity<*, *>.interrupt(interrupt: Interrupt)
 
     /**
      * Suspend the [Context] of the [Entity] in simulation for the given duration of simulation time before resuming
@@ -159,6 +175,11 @@ interface Context<S, out M> {
      * @param delay The amount of time to wait before the message should be received by the entity.
      */
     suspend fun Entity<*, *>.send(msg: Any, sender: Entity<*, *>, delay: Duration = 0)
+
+    /**
+     * This key provides users access to an untyped process context in case the coroutine runs inside a simulation.
+     */
+    companion object Key : CoroutineContext.Key<Context<*, *>>
 }
 
 /**
@@ -185,8 +206,9 @@ interface Envelope<out T : Any> {
 }
 
 /**
- * An [Interrupt] message is sent to a [Entity] process in order to interrupt its suspended state.
+ * An [Interrupt] message is sent to an [Entity] process in order to interrupt its suspended state.
  *
+ * @param reason The reason for the interruption of the process.
  * @author Fabian Mastenbroek (f.s.mastenbroek@student.tudelft.nl)
  */
-object Interrupt : Throwable("The entity process has been interrupted by another entity")
+open class Interrupt(reason: String) : Throwable(reason)
diff --git a/opendc-core/src/main/kotlin/com/atlarge/opendc/simulator/Process.kt b/opendc-core/src/main/kotlin/com/atlarge/opendc/simulator/Process.kt
index e8b4d988..f2b8a52b 100644
--- a/opendc-core/src/main/kotlin/com/atlarge/opendc/simulator/Process.kt
+++ b/opendc-core/src/main/kotlin/com/atlarge/opendc/simulator/Process.kt
@@ -8,7 +8,7 @@ package com.atlarge.opendc.simulator
  * @param M The shape of the model in which the process exists.
  * @author Fabian Mastenbroek (f.s.mastenbroek@student.tudelft.nl)
  */
-interface Process<S, in M> : Entity<S, M> {
+interface Process<S, M> : Entity<S, M> {
     /**
      * This method is invoked to start the simulation a process.
      *
diff --git a/opendc-kernel-omega/src/main/kotlin/com/atlarge/opendc/omega/Messages.kt b/opendc-kernel-omega/src/main/kotlin/com/atlarge/opendc/omega/Messages.kt
index 73c3676f..d63a53c8 100644
--- a/opendc-kernel-omega/src/main/kotlin/com/atlarge/opendc/omega/Messages.kt
+++ b/opendc-kernel-omega/src/main/kotlin/com/atlarge/opendc/omega/Messages.kt
@@ -27,4 +27,4 @@ object Timeout
  *
  * @author Fabian Mastenbroek (f.s.mastenbroek@student.tudelft.nl)
  */
-data class Launch<in M>(val process: Process<*, M>)
+data class Launch<M>(val process: Process<*, M>)
diff --git a/opendc-kernel-omega/src/main/kotlin/com/atlarge/opendc/omega/OmegaSimulation.kt b/opendc-kernel-omega/src/main/kotlin/com/atlarge/opendc/omega/OmegaSimulation.kt
index bd3f4529..22382ccd 100644
--- a/opendc-kernel-omega/src/main/kotlin/com/atlarge/opendc/omega/OmegaSimulation.kt
+++ b/opendc-kernel-omega/src/main/kotlin/com/atlarge/opendc/omega/OmegaSimulation.kt
@@ -242,7 +242,8 @@ internal class OmegaSimulation<M>(bootstrap: Bootstrap<M>) : Simulation<M>, Boot
     /**
      * This internal class provides the default implementation for the [Context] interface for this simulator.
      */
-    private inner class OmegaContext<S>(val process: Process<S, M>) : Context<S, M>, Continuation<Unit> {
+    private inner class OmegaContext<S>(val process: Process<S, M>) : Context<S, M>, Continuation<Unit>,
+        AbstractCoroutineContextElement(Context) {
         /**
          * The model in which the process exists.
          */
@@ -256,6 +257,12 @@ internal class OmegaSimulation<M>(bootstrap: Bootstrap<M>) : Simulation<M>, Boot
             get() = this@OmegaSimulation.time
 
         /**
+         * The [Entity] associated with this context.
+         */
+        override val self: Entity<S, M>
+            get() = process
+
+        /**
          * The duration between the current point in simulation time and the last point in simulation time where the
          * [Context] has executed some work.
          */
@@ -281,7 +288,7 @@ internal class OmegaSimulation<M>(bootstrap: Bootstrap<M>) : Simulation<M>, Boot
         /**
          * The [CoroutineContext] for a [Context].
          */
-        override val context: CoroutineContext = EmptyCoroutineContext
+        override val context: CoroutineContext = this
 
         /**
          * The continuation to resume the execution of the process.
@@ -321,7 +328,7 @@ internal class OmegaSimulation<M>(bootstrap: Bootstrap<M>) : Simulation<M>, Boot
         override suspend fun Entity<*, *>.send(msg: Any, sender: Entity<*, *>, delay: Duration) =
             schedule(prepare(msg, this, sender, delay))
 
-        override suspend fun Entity<*, *>.interrupt() = send(Interrupt)
+        override suspend fun Entity<*, *>.interrupt(interrupt: Interrupt) = send(interrupt)
 
         override suspend fun hold(duration: Duration) {
             require(duration >= 0) { "The amount of time to hold must be a positive number" }
diff --git a/opendc-stdlib/src/main/kotlin/com/atlarge/opendc/simulator/Helpers.kt b/opendc-stdlib/src/main/kotlin/com/atlarge/opendc/simulator/Helpers.kt
new file mode 100644
index 00000000..acf3fe41
--- /dev/null
+++ b/opendc-stdlib/src/main/kotlin/com/atlarge/opendc/simulator/Helpers.kt
@@ -0,0 +1,44 @@
+package com.atlarge.opendc.simulator
+
+import kotlin.coroutines.experimental.intrinsics.*
+
+/**
+ * Try to find the [Context] instance associated with the [Process] in the call chain which has (indirectly) invoked the
+ * caller of this method.
+ *
+ * Note however that this method does not guarantee type-safety as this method allows the user to cast to a context
+ * with different generic type arguments.
+ *
+ * @return The context that has been found or `null` if this method is not called in a simulation context.
+ */
+suspend fun <S, M> contextOrNull(): Context<S, M>? = suspendCoroutineOrReturn { it.context[Context] }
+
+/**
+ * Find the [Context] instance associated with the [Process] in the call chain which has (indirectly) invoked the
+ * caller of this method.
+ *
+ * Note however that this method does not guarantee type-safety as this method allows the user to cast to a context
+ * with different generic type arguments.
+ *
+ * @throws IllegalStateException if the context cannot be found.
+ * @return The context that has been found.
+ */
+suspend fun <S, M> context(): Context<S, M> =
+    contextOrNull() ?: throw IllegalStateException("The suspending call does not have an associated process context")
+
+/**
+ * Try to find the untyped [Context] instance associated with the [Process] in the call chain which has (indirectly)
+ * invoked the caller of this method.
+ *
+ * @return The untyped context that has been found or `null` if this method is not called in a simulation context.
+ */
+suspend fun untypedContextOrNull(): Context<*, *>? = contextOrNull<Any?, Any?>()
+
+/**
+ * Find the [Context] instance associated with the [Process] in the call chain which has (indirectly) invoked the
+ * caller of this method.
+ *
+ * @throws IllegalStateException if the context cannot be found.
+ * @return The untyped context that has been found.
+ */
+suspend fun untypedContext(): Context<*, *> = context<Any?, Any?>()