feat(sim/compute): Add completion parameter to startWorkload

This change updates the interface of `SimMachine#startWorkload` to
introduce a parameter `completion` that is invoked when the workload
completes either succesfully or due to failure.

This functionality has often been implemented by wrapping a
`SimWorkload` and catching its exceptions. However, since this
functionality is used in all usages of `SimMachine#startWorkload` we
instead embed it into `SimMachine` itself.

5 files changed, 60 insertions, 44 deletions

diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimAbstractMachine.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimAbstractMachine.java
index cf5aed03..d90a7d6f 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimAbstractMachine.java
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimAbstractMachine.java
@@ -25,6 +25,7 @@ package org.opendc.simulator.compute;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 import org.opendc.simulator.compute.device.SimNetworkAdapter;
 import org.opendc.simulator.compute.model.MachineModel;
 import org.opendc.simulator.compute.model.MemoryUnit;
@@ -62,12 +63,13 @@ public abstract class SimAbstractMachine implements SimMachine {
     }
 
     @Override
-    public final SimMachineContext startWorkload(SimWorkload workload, Map<String, Object> meta) {
+    public final SimMachineContext startWorkload(
+            SimWorkload workload, Map<String, Object> meta, Consumer<Exception> completion) {
         if (activeContext != null) {
             throw new IllegalStateException("A machine cannot run multiple workloads concurrently");
         }
 
-        final Context ctx = createContext(workload, new HashMap<>(meta));
+        final Context ctx = createContext(workload, new HashMap<>(meta), completion);
         ctx.start();
         return ctx;
     }
@@ -83,10 +85,12 @@ public abstract class SimAbstractMachine implements SimMachine {
     /**
      * Construct a new {@link Context} instance representing the active execution.
      *
-     * @param workload The workload to start on the machine.
-     * @param meta The metadata to pass to the workload.
+     * @param workload   The workload to start on the machine.
+     * @param meta       The metadata to pass to the workload.
+     * @param completion A block that is invoked when the workload completes carrying an exception if thrown by the workload.
      */
-    protected abstract Context createContext(SimWorkload workload, Map<String, Object> meta);
+    protected abstract Context createContext(
+            SimWorkload workload, Map<String, Object> meta, Consumer<Exception> completion);
 
     /**
      * Return the active {@link Context} instance (if any).
@@ -102,7 +106,9 @@ public abstract class SimAbstractMachine implements SimMachine {
         private final SimAbstractMachine machine;
         private final SimWorkload workload;
         private final Map<String, Object> meta;
+        private final Consumer<Exception> completion;
         private boolean isClosed;
+        private Exception cause;
 
         /**
          * Construct a new {@link Context} instance.
@@ -110,11 +116,17 @@ public abstract class SimAbstractMachine implements SimMachine {
          * @param machine The {@link SimAbstractMachine} to which the context belongs.
          * @param workload The {@link SimWorkload} to which the context belongs.
          * @param meta The metadata passed to the context.
+         * @param completion A block that is invoked when the workload completes carrying an exception if thrown by the workload.
          */
-        public Context(SimAbstractMachine machine, SimWorkload workload, Map<String, Object> meta) {
+        public Context(
+                SimAbstractMachine machine,
+                SimWorkload workload,
+                Map<String, Object> meta,
+                Consumer<Exception> completion) {
             this.machine = machine;
             this.workload = workload;
             this.meta = meta;
+            this.completion = completion;
         }
 
         @Override
@@ -136,11 +148,19 @@ public abstract class SimAbstractMachine implements SimMachine {
             // Cancel all the resources associated with the machine
             doCancel();
 
+            Exception e = this.cause;
+
             try {
                 workload.onStop(this);
             } catch (Exception cause) {
-                LOGGER.warn("Workload failed during onStop callback", cause);
+                if (e != null) {
+                    e.addSuppressed(cause);
+                } else {
+                    e = cause;
+                }
             }
+
+            completion.accept(e);
         }
 
         /**
@@ -151,7 +171,7 @@ public abstract class SimAbstractMachine implements SimMachine {
                 machine.activeContext = this;
                 workload.onStart(this);
             } catch (Exception cause) {
-                LOGGER.warn("Workload failed during onStart callback", cause);
+                this.cause = cause;
                 shutdown();
             }
         }
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimBareMetalMachine.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimBareMetalMachine.java
index aa7502d6..11356eb2 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimBareMetalMachine.java
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimBareMetalMachine.java
@@ -26,6 +26,7 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 import org.opendc.simulator.compute.device.SimPeripheral;
 import org.opendc.simulator.compute.model.MachineModel;
 import org.opendc.simulator.compute.model.ProcessingUnit;
@@ -39,7 +40,7 @@ import org.opendc.simulator.flow2.Inlet;
  *
  * <p>
  * A {@link SimBareMetalMachine} is a stateful object, and you should be careful when operating this object concurrently. For
- * example, the class expects only a single concurrent call to {@link #startWorkload(SimWorkload, Map)}.
+ * example, the class expects only a single concurrent call to {@link #startWorkload(SimWorkload, Map, Consumer)} )}.
  */
 public final class SimBareMetalMachine extends SimAbstractMachine {
     /**
@@ -192,8 +193,9 @@ public final class SimBareMetalMachine extends SimAbstractMachine {
     }
 
     @Override
-    protected SimAbstractMachine.Context createContext(SimWorkload workload, Map<String, Object> meta) {
-        return new Context(this, workload, meta);
+    protected SimAbstractMachine.Context createContext(
+            SimWorkload workload, Map<String, Object> meta, Consumer<Exception> completion) {
+        return new Context(this, workload, meta, completion);
     }
 
     /**
@@ -206,8 +208,12 @@ public final class SimBareMetalMachine extends SimAbstractMachine {
         private final List<NetworkAdapter> net;
         private final List<StorageDevice> disk;
 
-        private Context(SimBareMetalMachine machine, SimWorkload workload, Map<String, Object> meta) {
-            super(machine, workload, meta);
+        private Context(
+                SimBareMetalMachine machine,
+                SimWorkload workload,
+                Map<String, Object> meta,
+                Consumer<Exception> completion) {
+            super(machine, workload, meta, completion);
 
             this.graph = machine.graph;
             this.cpus = machine.cpus;
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimMachine.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimMachine.java
index 59599875..1f86aa02 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimMachine.java
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/SimMachine.java
@@ -24,6 +24,7 @@ package org.opendc.simulator.compute;
 
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 import org.opendc.simulator.compute.device.SimPeripheral;
 import org.opendc.simulator.compute.model.MachineModel;
 import org.opendc.simulator.compute.workload.SimWorkload;
@@ -47,10 +48,11 @@ public interface SimMachine {
      *
      * @param workload The workload to start on the machine.
      * @param meta The metadata to pass to the workload.
+     * @param completion A block that is invoked when the workload completes carrying an exception if thrown by the workload.
      * @return A {@link SimMachineContext} that represents the execution context for the workload.
      * @throws IllegalStateException if a workload is already active on the machine or if the machine is closed.
      */
-    SimMachineContext startWorkload(SimWorkload workload, Map<String, Object> meta);
+    SimMachineContext startWorkload(SimWorkload workload, Map<String, Object> meta, Consumer<Exception> completion);
 
     /**
      * Cancel the active workload on this machine (if any).
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/kernel/SimHypervisor.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/kernel/SimHypervisor.java
index 6e295837..f03a0c20 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/kernel/SimHypervisor.java
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/kernel/SimHypervisor.java
@@ -28,6 +28,7 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.SplittableRandom;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import org.opendc.simulator.compute.SimAbstractMachine;
 import org.opendc.simulator.compute.SimMachine;
@@ -471,7 +472,8 @@ public final class SimHypervisor implements SimWorkload {
         }
 
         @Override
-        protected Context createContext(SimWorkload workload, Map<String, Object> meta) {
+        protected Context createContext(
+                SimWorkload workload, Map<String, Object> meta, Consumer<Exception> completion) {
             if (isClosed) {
                 throw new IllegalStateException("Virtual machine does not exist anymore");
             }
@@ -482,7 +484,15 @@ public final class SimHypervisor implements SimWorkload {
             }
 
             return new VmContext(
-                    context, this, random, interferenceDomain, counters, SimHypervisor.this.counters, workload, meta);
+                    context,
+                    this,
+                    random,
+                    interferenceDomain,
+                    counters,
+                    SimHypervisor.this.counters,
+                    workload,
+                    meta,
+                    completion);
         }
 
         @Override
@@ -538,8 +548,9 @@ public final class SimHypervisor implements SimWorkload {
                 VmCounters vmCounters,
                 HvCounters hvCounters,
                 SimWorkload workload,
-                Map<String, Object> meta) {
-            super(machine, workload, meta);
+                Map<String, Object> meta,
+                Consumer<Exception> completion) {
+            super(machine, workload, meta, completion);
 
             this.context = context;
             this.random = random;
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/kotlin/org/opendc/simulator/compute/Coroutines.kt b/opendc-simulator/opendc-simulator-compute/src/main/kotlin/org/opendc/simulator/compute/Coroutines.kt
index c23f48dc..b354caff 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/kotlin/org/opendc/simulator/compute/Coroutines.kt
+++ b/opendc-simulator/opendc-simulator-compute/src/main/kotlin/org/opendc/simulator/compute/Coroutines.kt
@@ -39,31 +39,8 @@ public suspend fun SimMachine.runWorkload(workload: SimWorkload, meta: Map<Strin
     return suspendCancellableCoroutine { cont ->
         cont.invokeOnCancellation { this@runWorkload.cancel() }
 
-        startWorkload(
-            object : SimWorkload {
-                override fun onStart(ctx: SimMachineContext) {
-                    try {
-                        workload.onStart(ctx)
-                    } catch (cause: Throwable) {
-                        cont.resumeWithException(cause)
-                        throw cause
-                    }
-                }
-
-                override fun onStop(ctx: SimMachineContext) {
-                    try {
-                        workload.onStop(ctx)
-
-                        if (!cont.isCompleted) {
-                            cont.resume(Unit)
-                        }
-                    } catch (cause: Throwable) {
-                        cont.resumeWithException(cause)
-                        throw cause
-                    }
-                }
-            },
-            meta
-        )
+        startWorkload(workload, meta) { cause ->
+            if (cause != null) cont.resumeWithException(cause) else cont.resume(Unit)
+        }
     }
 }