refactor(simulator): Remove failure callback from FlowSource

This change removes the `onFailure` method from FlowSource. Instead, the
FlowConsumer will receive the reason for failure of the source.

2 files changed, 15 insertions, 9 deletions

diff --git a/opendc-simulator/opendc-simulator-compute/build.gradle.kts b/opendc-simulator/opendc-simulator-compute/build.gradle.kts
index e2290a14..a2bb89c2 100644
--- a/opendc-simulator/opendc-simulator-compute/build.gradle.kts
+++ b/opendc-simulator/opendc-simulator-compute/build.gradle.kts
@@ -36,4 +36,6 @@ dependencies {
     api(projects.opendcSimulator.opendcSimulatorNetwork)
     implementation(projects.opendcSimulator.opendcSimulatorCore)
     implementation(projects.opendcUtils)
+
+    testImplementation(libs.slf4j.simple)
 }
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
                 }
             }