3 files changed, 62 insertions, 0 deletions
diff --git a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/FlowMultiplexer.java b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/FlowMultiplexer.java
index 1a99d0cf..2a23b039 100644
--- a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/FlowMultiplexer.java
+++ b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/FlowMultiplexer.java
@@ -67,4 +67,19 @@ public interface FlowMultiplexer {
      * @param outlet The outlet to release.
      */
     void releaseOutput(Outlet outlet);
+
+    /**
+     * Return the total input capacity of the {@link FlowMultiplexer}.
+     */
+    float getCapacity();
+
+    /**
+     * Return the total input demand for the {@link FlowMultiplexer}.
+     */
+    float getDemand();
+
+    /**
+     * Return the total input rate for the {@link FlowMultiplexer}.
+     */
+    float getRate();
 }
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 9ff6a4c8..6394c3fd 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
@@ -51,6 +51,9 @@ public final class ForwardingFlowMultiplexer implements FlowMultiplexer, FlowSta
     private final BitSet activeOutputs;
     private final BitSet availableOutputs;
 
+    private float capacity = 0.f;
+    private float demand = 0.f;
+
     public ForwardingFlowMultiplexer(FlowGraph graph) {
         this.stage = graph.newStage(this);
 
@@ -62,6 +65,27 @@ public final class ForwardingFlowMultiplexer implements FlowMultiplexer, FlowSta
     }
 
     @Override
+    public float getCapacity() {
+        return capacity;
+    }
+
+    @Override
+    public float getDemand() {
+        return demand;
+    }
+
+    @Override
+    public float getRate() {
+        final BitSet activeOutputs = this.activeOutputs;
+        final OutPort[] outlets = this.outlets;
+        float rate = 0.f;
+        for (int i = activeOutputs.nextSetBit(0); i != -1; i = activeOutputs.nextSetBit(i + 1)) {
+            rate += outlets[i].getRate();
+        }
+        return rate;
+    }
+
+    @Override
     public int getInputCount() {
         return activeInputs.length();
     }
@@ -176,11 +200,15 @@ public final class ForwardingFlowMultiplexer implements FlowMultiplexer, FlowSta
 
         @Override
         public void onPush(InPort port, float rate) {
+            ForwardingFlowMultiplexer.this.demand += -port.getDemand() + rate;
+
             output.push(rate);
         }
 
         @Override
         public void onUpstreamFinish(InPort port, Throwable cause) {
+            ForwardingFlowMultiplexer.this.demand -= port.getDemand();
+
             final OutPort output = this.output;
             output.push(0.f);
 
@@ -197,11 +225,15 @@ public final class ForwardingFlowMultiplexer implements FlowMultiplexer, FlowSta
 
         @Override
         public void onPull(OutPort port, float capacity) {
+            ForwardingFlowMultiplexer.this.capacity += -port.getCapacity() + capacity;
+
             input.pull(capacity);
         }
 
         @Override
         public void onDownstreamFinish(OutPort port, Throwable cause) {
+            ForwardingFlowMultiplexer.this.capacity -= port.getCapacity();
+
             input.cancel(cause);
 
             releaseOutput(port);
diff --git a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/MaxMinFlowMultiplexer.java b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/MaxMinFlowMultiplexer.java
index ca0639f5..77922066 100644
--- a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/MaxMinFlowMultiplexer.java
+++ b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/flow2/mux/MaxMinFlowMultiplexer.java
@@ -76,6 +76,21 @@ public final class MaxMinFlowMultiplexer implements FlowMultiplexer, FlowStageLo
     }
 
     @Override
+    public float getCapacity() {
+        return capacity;
+    }
+
+    @Override
+    public float getDemand() {
+        return demand;
+    }
+
+    @Override
+    public float getRate() {
+        return rate;
+    }
+
+    @Override
     public long onUpdate(FlowStage ctx, long now) {
         float capacity = this.capacity;
         float demand = this.demand;