2 files changed, 15 insertions, 3 deletions
diff --git a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/ForwardingFlowMultiplexer.java b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/ForwardingFlowMultiplexer.java
index abe3510b..e0564cd2 100644
--- a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/ForwardingFlowMultiplexer.java
+++ b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/ForwardingFlowMultiplexer.java
@@ -182,6 +182,9 @@ public final class ForwardingFlowMultiplexer implements FlowMultiplexer, FlowSta
 
         activeOutputs.set(slot);
         availableOutputs.set(slot);
+
+        port.setHandler(IDLE_OUT_HANDLER);
+
         return port;
     }
 
@@ -270,11 +273,15 @@ public final class ForwardingFlowMultiplexer implements FlowMultiplexer, FlowSta
         public void onUpstreamFinish(InPort port, Throwable cause) {}
     }
 
-    private static class IdleOutHandler implements OutHandler {
+    private class IdleOutHandler implements OutHandler {
         @Override
-        public void onPull(OutPort port, float capacity) {}
+        public void onPull(OutPort port, float capacity) {
+            ForwardingFlowMultiplexer.this.capacity += -port.getCapacity() + capacity;
+        }
 
         @Override
-        public void onDownstreamFinish(OutPort port, Throwable cause) {}
+        public void onDownstreamFinish(OutPort port, Throwable cause) {
+            ForwardingFlowMultiplexer.this.capacity -= port.getCapacity();
+        }
     }
 }
diff --git a/opendc-simulator/opendc-simulator-flow/src/test/kotlin/org/opendc/simulator/flow2/mux/ForwardingFlowMultiplexerTest.kt b/opendc-simulator/opendc-simulator-flow/src/test/kotlin/org/opendc/simulator/flow2/mux/ForwardingFlowMultiplexerTest.kt
index fef49786..d1795841 100644
--- a/opendc-simulator/opendc-simulator-flow/src/test/kotlin/org/opendc/simulator/flow2/mux/ForwardingFlowMultiplexerTest.kt
+++ b/opendc-simulator/opendc-simulator-flow/src/test/kotlin/org/opendc/simulator/flow2/mux/ForwardingFlowMultiplexerTest.kt
@@ -22,6 +22,7 @@
 
 package org.opendc.simulator.flow2.mux
 
+import kotlinx.coroutines.yield
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.assertAll
@@ -46,6 +47,10 @@ class ForwardingFlowMultiplexerTest {
         val sink = SimpleFlowSink(graph, 3200.0f)
         graph.connect(switch.newOutput(), sink.input)
 
+        yield()
+
+        assertEquals(sink.capacity, switch.capacity) { "Capacity is not detected" }
+
         val workload =
             TraceFlowSource(
                 graph,