1 files changed, 13 insertions, 9 deletions
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/kotlin/org/opendc/simulator/compute/workload/SimWorkloadLifecycle.kt b/opendc-simulator/opendc-simulator-compute/src/main/kotlin/org/opendc/simulator/compute/workload/SimWorkloadLifecycle.kt
index dabe60e0..b85be39d 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/kotlin/org/opendc/simulator/compute/workload/SimWorkloadLifecycle.kt
+++ b/opendc-simulator/opendc-simulator-compute/src/main/kotlin/org/opendc/simulator/compute/workload/SimWorkloadLifecycle.kt
@@ -41,22 +41,26 @@ public class SimWorkloadLifecycle(private val ctx: SimMachineContext) {
      */
     public fun waitFor(consumer: FlowSource): FlowSource {
         waiting.add(consumer)
-        return object : FlowSource by consumer {
+        return object : FlowSource {
+            override fun onPull(conn: FlowConnection, now: Long, delta: Long): Long {
+                return try {
+                    consumer.onPull(conn, now, delta)
+                } catch (cause: Throwable) {
+                    complete(consumer)
+                    throw cause
+                }
+            }
+
             override fun onEvent(conn: FlowConnection, now: Long, event: FlowEvent) {
                 try {
                     consumer.onEvent(conn, now, event)
-                } finally {
+
                     if (event == FlowEvent.Exit) {
                         complete(consumer)
                     }
-                }
-            }
-
-            override fun onFailure(conn: FlowConnection, cause: Throwable) {
-                try {
-                    consumer.onFailure(conn, cause)
-                } finally {
+                } catch (cause: Throwable) {
                     complete(consumer)
+                    throw cause
                 }
             }