multi gpu support (#351)

9 files changed, 238 insertions, 132 deletions

diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/machine/SimMachine.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/machine/SimMachine.java
index 8792552e..c9e3ab8c 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/machine/SimMachine.java
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/machine/SimMachine.java
@@ -211,27 +211,43 @@ public class SimMachine {
                 ResourceType.CPU,
                 new ArrayList<>(List.of(new SimCpu(engine, this.machineModel.getCpuModel(), cpuPowerModel, 0))));
 
-        new FlowEdge((FlowConsumer) this.computeResources.get(ResourceType.CPU).getFirst(), this.psu);
+        // Connect the CPU to the PSU
+        new FlowEdge(
+                (FlowConsumer) this.computeResources.get(ResourceType.CPU).getFirst(),
+                this.psu,
+                ResourceType.POWER,
+                0,
+                -1);
 
         // Create a FlowDistributor and add the cpu as supplier
         this.distributors.put(ResourceType.CPU, new FlowDistributor(engine));
-        new FlowEdge(this.distributors.get(ResourceType.CPU), (FlowSupplier)
-                this.computeResources.get(ResourceType.CPU).getFirst());
+        new FlowEdge(
+                this.distributors.get(ResourceType.CPU),
+                (FlowSupplier) this.computeResources.get(ResourceType.CPU).getFirst(),
+                ResourceType.CPU,
+                -1,
+                0);
 
         // TODO: include memory as flow node
         this.memory = new Memory(engine, this.machineModel.getMemory());
 
         if (this.availableResources.contains(ResourceType.GPU)) {
             this.distributors.put(ResourceType.GPU, new FlowDistributor(engine));
-            short i = 0;
             ArrayList<ComputeResource> gpus = new ArrayList<>();
 
             for (GpuModel gpuModel : machineModel.getGpuModels()) {
-                SimGpu gpu = new SimGpu(engine, gpuModel, gpuPowerModel, i);
+                // create a new GPU
+                SimGpu gpu = new SimGpu(engine, gpuModel, gpuPowerModel, gpuModel.getId());
                 gpus.add(gpu);
-                // suspends here without the distributor
-                new FlowEdge(this.distributors.get(ResourceType.GPU), gpu);
-                new FlowEdge(gpu, this.psu);
+                // Connect the GPU to the distributor
+                new FlowEdge(
+                        this.distributors.get(ResourceType.GPU),
+                        gpu,
+                        ResourceType.GPU,
+                        gpuModel.getId(),
+                        gpuModel.getId());
+                // Connect the GPU to the PSU
+                new FlowEdge(gpu, this.psu, ResourceType.POWER, gpuModel.getId(), gpuModel.getId());
             }
             this.computeResources.put(ResourceType.GPU, gpus);
         }
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/models/MachineModel.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/models/MachineModel.java
index 874194f6..e11d9cf2 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/models/MachineModel.java
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/models/MachineModel.java
@@ -35,7 +35,6 @@ import org.opendc.simulator.engine.graph.distributionPolicies.DistributionPolicy
 public final class MachineModel {
     private final CpuModel cpuModel;
     private final MemoryUnit memory;
-    //    private final List<GpuModel> gpuModels = new ArrayList<>(); // TODO: Implement multi GPU support
     private final List<GpuModel> gpuModels;
     private final DistributionPolicy cpuDistributionStrategy;
     private final DistributionPolicy gpuDistributionPolicy;
@@ -58,19 +57,8 @@ public final class MachineModel {
         this.gpuDistributionPolicy = gpuDistributionPolicy;
         this.availableResources.add(ResourceType.CPU);
         // TODO: Add Memory
-        //        this.usedResources.add(ResourceType.Memory);
         if (gpuModels != null && !gpuModels.isEmpty()) {
-            //            this.gpuModels = gpuModels;
-            this.gpuModels = new ArrayList<>();
-            this.gpuModels.add(new GpuModel(
-                    0,
-                    gpuModels.getFirst().getCoreCount() * gpuModels.size(), // merges multiple GPUs into one
-                    gpuModels.getFirst().getCoreSpeed(),
-                    gpuModels.getFirst().getMemoryBandwidth(),
-                    gpuModels.getFirst().getMemorySize() * gpuModels.size(), // merges multiple GPUs into one
-                    gpuModels.getFirst().getVendor(),
-                    gpuModels.getFirst().getModelName(),
-                    gpuModels.getFirst().getArchitecture()));
+            this.gpuModels = gpuModels;
             this.availableResources.add(ResourceType.GPU);
         } else {
             this.gpuModels = new ArrayList<>();
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/SimPsu.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/SimPsu.java
index 1ea7c570..f40f4fec 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/SimPsu.java
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/SimPsu.java
@@ -43,11 +43,14 @@ public final class SimPsu extends FlowNode implements FlowSupplier, FlowConsumer
     private static final Logger LOGGER = LoggerFactory.getLogger(SimPsu.class);
     private long lastUpdate;
 
-    private final HashMap<ResourceType, ArrayList<Double>> powerDemandsPerResource = new HashMap<>();
-    private final HashMap<ResourceType, ArrayList<Double>> powerSuppliedPerResource = new HashMap<>();
+    private final HashMap<ResourceType, HashMap<Integer, Double>> powerDemandsPerResource = new HashMap<>();
+    private final HashMap<ResourceType, HashMap<Integer, Double>> powerSuppliedPerResource = new HashMap<>();
+
+    private double totalPowerDemand = 0.0;
+    private double totalPowerSupplied = 0.0;
     private double totalEnergyUsage = 0.0;
 
-    private final HashMap<ResourceType, ArrayList<FlowEdge>> resourceEdges = new HashMap<>();
+    private final HashMap<ResourceType, HashMap<Integer, FlowEdge>> resourceEdges = new HashMap<>();
     private FlowEdge powerSupplyEdge;
 
     private final double capacity = Long.MAX_VALUE;
@@ -72,28 +75,51 @@ public final class SimPsu extends FlowNode implements FlowSupplier, FlowConsumer
      * This method provides access to the power consumption of the machine before PSU losses are applied.
      */
     public double getPowerDemand() {
-        return this.powerDemandsPerResource.values().stream()
-                .flatMap(List::stream)
-                .findFirst()
-                .orElse(0.0);
+        return this.totalPowerDemand;
     }
 
+    /**
+     * Return the power demand of the machine (in W) measured in the PSU for a specific resource type.
+     * <p>
+     * This method provides access to the power consumption of the machine before PSU losses are applied.
+     */
     public double getPowerDemand(ResourceType resourceType) {
-        return this.powerDemandsPerResource.get(resourceType).getFirst();
+        //        return this.powerDemandsPerResource.get(resourceType).stream().mapToDouble(Double::doubleValue).sum();
+        return this.powerDemandsPerResource.get(resourceType).values().stream()
+                .mapToDouble(Double::doubleValue)
+                .sum();
+    }
+
+    /**
+     * Return the power demand of the machine (in W) measured in the PSU for a specific resource type for a specific resource.
+     * <p>
+     * This method provides access to the power consumption of the machine before PSU losses are applied.
+     */
+    public double getPowerDemand(ResourceType resourceType, int id) {
+        return this.powerDemandsPerResource.get(resourceType).get(id);
     }
 
     /**
      * Return the instantaneous power usage of the machine (in W) measured at the InPort of the power supply.
      */
     public double getPowerDraw() {
-        return this.powerSuppliedPerResource.values().stream()
-                .flatMap(List::stream)
-                .findFirst()
-                .orElse(0.0);
+        return this.totalPowerSupplied;
     }
 
+    /**
+     * Return the instantaneous power usage of the machine (in W) measured at the InPort of the power supply for a specific resource type.
+     */
     public double getPowerDraw(ResourceType resourceType) {
-        return this.powerSuppliedPerResource.get(resourceType).getFirst();
+        return this.powerSuppliedPerResource.get(resourceType).values().stream()
+                .mapToDouble(Double::doubleValue)
+                .sum();
+    }
+
+    /**
+     * Return the instantaneous power usage of the machine (in W) measured at the InPort of the power supply for a specific resource type for a specific resource.
+     */
+    public double getPowerDraw(ResourceType resourceType, int id) {
+        return this.powerSuppliedPerResource.get(resourceType).get(id);
     }
 
     /**
@@ -127,16 +153,22 @@ public final class SimPsu extends FlowNode implements FlowSupplier, FlowConsumer
     public long onUpdate(long now) {
         updateCounters();
         for (ResourceType resourceType : this.resourceEdges.keySet()) {
-            ArrayList<FlowEdge> edges = this.resourceEdges.get(resourceType);
+            HashMap<Integer, FlowEdge> edges = this.resourceEdges.get(resourceType);
             if (edges != null && !edges.isEmpty()) {
-                double powerSupply =
-                        this.powerDemandsPerResource.get(resourceType).getFirst();
-                double powerSupplied =
-                        this.powerSuppliedPerResource.get(resourceType).getFirst();
-
-                if (powerSupply != powerSupplied) {
-                    for (FlowEdge edge : edges) {
-                        edge.pushSupply(powerSupply);
+                for (FlowEdge edge : edges.values()) {
+                    // If the edge is null, it means that the edge has been removed -> no update is needed
+                    if (edge == null) {
+                        continue;
+                    }
+
+                    int consumerIndex = edge.getConsumerIndex() == -1 ? 0 : edge.getConsumerIndex();
+                    double powerDemand =
+                            this.powerDemandsPerResource.get(resourceType).get(consumerIndex);
+                    double powerSupplied =
+                            this.powerSuppliedPerResource.get(resourceType).get(consumerIndex);
+
+                    if (powerDemand != powerSupplied) {
+                        edge.pushSupply(powerDemand);
                     }
                 }
             }
@@ -159,7 +191,8 @@ public final class SimPsu extends FlowNode implements FlowSupplier, FlowConsumer
         long duration = now - lastUpdate;
         if (duration > 0) {
             for (ResourceType resourceType : this.powerSuppliedPerResource.keySet()) {
-                for (double powerSupplied : this.powerSuppliedPerResource.get(resourceType)) {
+                for (double powerSupplied :
+                        this.powerSuppliedPerResource.get(resourceType).values()) {
                     this.totalEnergyUsage += (powerSupplied * duration * 0.001);
                 }
             }
@@ -171,17 +204,8 @@ public final class SimPsu extends FlowNode implements FlowSupplier, FlowConsumer
     ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
     @Override
-    public void pushOutgoingDemand(FlowEdge supplierEdge, double newDemand, ResourceType resourceType) {
-        this.powerDemandsPerResource.put(resourceType, new ArrayList<>(List.of(newDemand)));
-        powerSupplyEdge.pushDemand(newDemand);
-    }
-
-    @Override
     public void pushOutgoingDemand(FlowEdge supplierEdge, double newDemand) {
-        double totalDemand = this.powerDemandsPerResource.values().stream()
-                .flatMap(List::stream)
-                .reduce(0.0, Double::sum);
-        this.powerSupplyEdge.pushDemand(totalDemand);
+        this.powerSupplyEdge.pushDemand(newDemand);
     }
 
     @Override
@@ -191,7 +215,13 @@ public final class SimPsu extends FlowNode implements FlowSupplier, FlowConsumer
 
     @Override
     public void pushOutgoingSupply(FlowEdge consumerEdge, double newSupply, ResourceType resourceType) {
-        this.powerSuppliedPerResource.put(resourceType, new ArrayList<>(List.of(newSupply)));
+        int consumerIndex = consumerEdge.getConsumerIndex() == -1 ? 0 : consumerEdge.getConsumerIndex();
+
+        double previousSupply = this.powerSuppliedPerResource.get(resourceType).get(consumerIndex);
+        this.totalPowerSupplied += newSupply - previousSupply;
+
+        this.powerSuppliedPerResource.get(resourceType).put(consumerIndex, newSupply);
+
         consumerEdge.pushSupply(newSupply, false, resourceType);
     }
 
@@ -203,18 +233,29 @@ public final class SimPsu extends FlowNode implements FlowSupplier, FlowConsumer
     @Override
     public void handleIncomingDemand(FlowEdge consumerEdge, double newPowerDemand, ResourceType resourceType) {
         updateCounters();
-        this.powerDemandsPerResource.put(resourceType, new ArrayList<>(List.of(newPowerDemand)));
+        int consumerIndex = consumerEdge.getConsumerIndex() == -1 ? 0 : consumerEdge.getConsumerIndex();
+
+        double previousPowerDemand =
+                this.powerDemandsPerResource.get(resourceType).get(consumerIndex);
+        this.totalPowerDemand += newPowerDemand - previousPowerDemand;
 
-        pushOutgoingDemand(this.powerSupplyEdge, newPowerDemand);
+        this.powerDemandsPerResource.get(resourceType).put(consumerIndex, newPowerDemand);
+
+        pushOutgoingDemand(this.powerSupplyEdge, totalPowerDemand);
     }
 
     @Override
     public void handleIncomingSupply(FlowEdge supplierEdge, double newSupply) {
         updateCounters();
         for (ResourceType resourceType : this.resourceEdges.keySet()) {
-            for (FlowEdge edge : this.resourceEdges.get(resourceType)) {
+            for (FlowEdge edge : this.resourceEdges.get(resourceType).values()) {
+                // If the edge is null, it means that the edge has been removed -> no update is needed
+                if (edge == null) {
+                    continue;
+                }
+                int consumerIndex = edge.getConsumerIndex() == -1 ? 0 : edge.getConsumerIndex();
                 double outgoingSupply =
-                        Math.min(this.powerDemandsPerResource.get(resourceType).getFirst(), newSupply);
+                        Math.min(this.powerDemandsPerResource.get(resourceType).get(consumerIndex), newSupply);
                 pushOutgoingSupply(edge, outgoingSupply, resourceType);
             }
         }
@@ -222,10 +263,19 @@ public final class SimPsu extends FlowNode implements FlowSupplier, FlowConsumer
 
     @Override
     public void addConsumerEdge(FlowEdge consumerEdge) {
+
         ResourceType consumerResourceType = consumerEdge.getConsumerResourceType();
-        this.resourceEdges.put(consumerResourceType, new ArrayList<>(List.of(consumerEdge)));
-        this.powerDemandsPerResource.put(consumerResourceType, new ArrayList<>(List.of(0.0)));
-        this.powerSuppliedPerResource.put(consumerResourceType, new ArrayList<>(List.of(0.0)));
+        int consumerIndex = consumerEdge.getConsumerIndex() == -1 ? 0 : consumerEdge.getConsumerIndex();
+
+        if (!this.resourceEdges.containsKey(consumerResourceType)) {
+            this.resourceEdges.put(consumerResourceType, new HashMap<>());
+            this.powerDemandsPerResource.put(consumerResourceType, new HashMap<>());
+            this.powerSuppliedPerResource.put(consumerResourceType, new HashMap<>());
+        }
+
+        this.resourceEdges.get(consumerResourceType).put(consumerIndex, consumerEdge);
+        this.powerDemandsPerResource.get(consumerResourceType).put(consumerIndex, 0.0);
+        this.powerSuppliedPerResource.get(consumerResourceType).put(consumerIndex, 0.0);
     }
 
     @Override
@@ -236,10 +286,18 @@ public final class SimPsu extends FlowNode implements FlowSupplier, FlowConsumer
     @Override
     public void removeConsumerEdge(FlowEdge consumerEdge) {
         ResourceType resourceType = consumerEdge.getConsumerResourceType();
+        int consumerIndex = consumerEdge.getConsumerIndex() == -1 ? 0 : consumerEdge.getConsumerIndex();
+
         if (this.resourceEdges.containsKey(resourceType)) {
-            this.resourceEdges.remove(resourceType);
-            this.powerDemandsPerResource.remove(resourceType);
-            this.powerSuppliedPerResource.remove(resourceType);
+            this.resourceEdges.get(resourceType).put(consumerIndex, null);
+
+            this.totalPowerDemand -=
+                    this.powerDemandsPerResource.get(resourceType).get(consumerIndex);
+            this.powerDemandsPerResource.get(resourceType).put(consumerIndex, 0.0);
+
+            this.totalPowerSupplied -=
+                    this.powerSuppliedPerResource.get(resourceType).get(consumerIndex);
+            this.powerSuppliedPerResource.get(resourceType).put(consumerIndex, 0.0);
         }
     }
 
@@ -252,7 +310,8 @@ public final class SimPsu extends FlowNode implements FlowSupplier, FlowConsumer
     public Map<FlowEdge.NodeType, List<FlowEdge>> getConnectedEdges() {
         List<FlowEdge> supplyingEdges = new ArrayList<>();
         for (ResourceType resourceType : this.resourceEdges.keySet()) {
-            List<FlowEdge> edges = this.resourceEdges.get(resourceType);
+            List<FlowEdge> edges =
+                    this.resourceEdges.get(resourceType).values().stream().toList();
             if (edges != null && !edges.isEmpty()) {
                 supplyingEdges.addAll(edges);
             }
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/workload/VirtualMachine.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/workload/VirtualMachine.java
index 622d2b89..8922a97d 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/workload/VirtualMachine.java
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/workload/VirtualMachine.java
@@ -56,7 +56,7 @@ public final class VirtualMachine extends SimWorkload implements FlowSupplier {
     private final Hashtable<ResourceType, Double> resourceCapacities = new Hashtable<>();
     private final Hashtable<ResourceType, Double> resourceTimeScalingFactor = new Hashtable<>(); // formerly known as d
     private final Hashtable<ResourceType, FlowEdge> distributorEdges = new Hashtable<>();
-    private final Hashtable<ResourceType, List<PerformanceCounters>> resourcePerformanceCounters = new Hashtable<>();
+    private final Hashtable<ResourceType, PerformanceCounters> resourcePerformanceCounters = new Hashtable<>();
 
     private final long checkpointInterval;
     private final long checkpointDuration;
@@ -108,21 +108,11 @@ public final class VirtualMachine extends SimWorkload implements FlowSupplier {
     }
 
     public PerformanceCounters getCpuPerformanceCounters() {
-        return this.resourcePerformanceCounters.get(ResourceType.CPU).getFirst();
+        return this.resourcePerformanceCounters.get(ResourceType.CPU);
     }
 
-    public List<PerformanceCounters> getGpuPerformanceCounters() {
-        return this.resourcePerformanceCounters.get(ResourceType.GPU) != null
-                ? this.resourcePerformanceCounters.get(ResourceType.GPU)
-                : new ArrayList<>();
-    }
-
-    public PerformanceCounters getGpuPerformanceCounters(int gpuId) {
-        List<PerformanceCounters> gpuPerformanceCounters = this.resourcePerformanceCounters.get(ResourceType.GPU);
-        if (gpuId < 0 || gpuId >= gpuPerformanceCounters.size()) {
-            throw new IndexOutOfBoundsException("No such GPU id: " + gpuId);
-        }
-        return gpuPerformanceCounters.get(gpuId);
+    public PerformanceCounters getGpuPerformanceCounters() {
+        return this.resourcePerformanceCounters.get(ResourceType.GPU);
     }
 
     ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -171,13 +161,9 @@ public final class VirtualMachine extends SimWorkload implements FlowSupplier {
 
             this.resourceCapacities.put(resourceType, resources.getFirst().getCapacity());
 
-            ArrayList<PerformanceCounters> performanceCounters = new ArrayList<>();
-
-            for (ComputeResource resource : resources) {
-                performanceCounters.add(new PerformanceCounters());
-                this.resourceTimeScalingFactor.put(resourceType, 1.0 / resource.getCapacity());
-            }
-            this.resourcePerformanceCounters.put(resourceType, performanceCounters);
+            this.resourceTimeScalingFactor.put(
+                    resourceType, 1.0 / resources.getFirst().getCapacity());
+            this.resourcePerformanceCounters.put(resourceType, new PerformanceCounters());
             this.resourceDemands.put(resourceType, 0.0);
             this.resourceSupplies.put(resourceType, 0.0);
         }
@@ -225,21 +211,24 @@ public final class VirtualMachine extends SimWorkload implements FlowSupplier {
         for (ResourceType resourceType : this.availableResources) {
             int i = 0;
             final double factor = this.resourceTimeScalingFactor.get(resourceType) * delta;
-            for (PerformanceCounters performanceCounter : this.resourcePerformanceCounters.get(resourceType)) {
-                if (delta > 0) {
-                    performanceCounter.addActiveTime(Math.round(this.resourceSupplies.get(resourceType) * factor));
-                    performanceCounter.setIdleTime(Math.round(
-                            (this.resourceCapacities.get(resourceType) - this.resourceSupplies.get(resourceType))
-                                    * factor));
-                    performanceCounter.addStealTime(Math.round(
-                            (this.resourceDemands.get(resourceType) - this.resourceSupplies.get(resourceType))
-                                    * factor));
-                }
-                performanceCounter.setDemand(this.resourceDemands.get(resourceType));
-                performanceCounter.setSupply(this.resourceSupplies.get(resourceType));
-                performanceCounter.setCapacity(this.resourceCapacities.get(resourceType));
-                i++;
+            if (delta > 0) {
+                this.resourcePerformanceCounters
+                        .get(resourceType)
+                        .addActiveTime(Math.round(this.resourceSupplies.get(resourceType) * factor));
+                this.resourcePerformanceCounters
+                        .get(resourceType)
+                        .setIdleTime(Math.round(
+                                (this.resourceCapacities.get(resourceType) - this.resourceSupplies.get(resourceType))
+                                        * factor));
+                this.resourcePerformanceCounters
+                        .get(resourceType)
+                        .addStealTime(Math.round(
+                                (this.resourceDemands.get(resourceType) - this.resourceSupplies.get(resourceType))
+                                        * factor));
             }
+            this.resourcePerformanceCounters.get(resourceType).setDemand(this.resourceDemands.get(resourceType));
+            this.resourcePerformanceCounters.get(resourceType).setSupply(this.resourceSupplies.get(resourceType));
+            this.resourcePerformanceCounters.get(resourceType).setCapacity(this.resourceCapacities.get(resourceType));
         }
     }
 
@@ -317,7 +306,12 @@ public final class VirtualMachine extends SimWorkload implements FlowSupplier {
     @Override
     public void addSupplierEdge(FlowEdge supplierEdge) {
         ResourceType resourceType = supplierEdge.getSupplierResourceType();
-        this.resourceCapacities.put(resourceType, supplierEdge.getCapacity());
+        if (this.resourceCapacities.containsKey(resourceType) && this.resourceCapacities.get(resourceType) > 0) {
+            this.resourceCapacities.put(
+                    resourceType, this.resourceCapacities.get(resourceType) + supplierEdge.getCapacity());
+        } else {
+            this.resourceCapacities.put(resourceType, supplierEdge.getCapacity());
+        }
         this.distributorEdges.put(resourceType, supplierEdge);
     }
 
diff --git a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/FlowDistributor.java b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/FlowDistributor.java
index 674db8ca..f7fc2728 100644
--- a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/FlowDistributor.java
+++ b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/FlowDistributor.java
@@ -23,6 +23,7 @@
 package org.opendc.simulator.engine.graph;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -32,20 +33,30 @@ import org.opendc.common.ResourceType;
 import org.opendc.simulator.engine.engine.FlowEngine;
 import org.opendc.simulator.engine.graph.distributionPolicies.DistributionPolicy;
 import org.opendc.simulator.engine.graph.distributionPolicies.MaxMinFairnessPolicy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsumer {
+    private static final Logger LOGGER = LoggerFactory.getLogger(FlowDistributor.class);
     private final ArrayList<FlowEdge> consumerEdges = new ArrayList<>();
-    private FlowEdge supplierEdge;
+    private HashMap<Integer, FlowEdge> supplierEdges =
+            new HashMap<>(); // The suppliers that provide supply to this distributor
 
     private final ArrayList<Double> incomingDemands = new ArrayList<>(); // What is demanded by the consumers
     private final ArrayList<Double> outgoingSupplies = new ArrayList<>(); // What is supplied to the consumers
 
     private double totalIncomingDemand; // The total demand of all the consumers
-    private double currentIncomingSupply; // The current supply provided by the supplier
+    // AS index is based on the supplierIndex of the FlowEdge, ids of entries need to be stable
+    private HashMap<Integer, Double> currentIncomingSupplies =
+            new HashMap<>(); // The current supply provided by the suppliers
+    private Double totalIncomingSupply = 0.0; // The total supply provided by the suppliers
 
     private boolean outgoingDemandUpdateNeeded = false;
     private Set<Integer> updatedDemands = new HashSet<>(); // Array of consumers that updated their demand in this cycle
 
+    private ResourceType supplierResourceType;
+    private ResourceType consumerResourceType;
+
     private boolean overloaded = false;
 
     private double capacity; // What is the max capacity. Can probably be removed
@@ -66,7 +77,7 @@ public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsu
     }
 
     public double getCurrentIncomingSupply() {
-        return currentIncomingSupply;
+        return this.totalIncomingSupply;
     }
 
     public double getCapacity() {
@@ -90,7 +101,13 @@ public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsu
     }
 
     private void updateOutgoingDemand() {
-        this.pushOutgoingDemand(this.supplierEdge, this.totalIncomingDemand);
+        // equally distribute the demand to all suppliers
+        for (FlowEdge supplierEdge : this.supplierEdges.values()) {
+            this.pushOutgoingDemand(supplierEdge, this.totalIncomingDemand / this.supplierEdges.size());
+            // alternatively a relative share could be used, based on capacity minus current incoming supply
+            //            this.pushOutgoingDemand(supplierEdge, this.totalIncomingDemand * (supplierEdge.getCapacity() -
+            // currentIncomingSupplies.get(idx) / supplierEdges.size()));
+        }
 
         this.outgoingDemandUpdateNeeded = false;
 
@@ -102,11 +119,13 @@ public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsu
 
         // If the demand is higher than the current supply, the system is overloaded.
         // The available supply is distributed based on the current distribution function.
-        if (this.totalIncomingDemand > this.currentIncomingSupply) {
+        if (this.totalIncomingDemand > this.totalIncomingSupply) {
             this.overloaded = true;
 
-            double[] supplies =
-                    this.distributionPolicy.distributeSupply(this.incomingDemands, this.currentIncomingSupply);
+            double[] supplies = this.distributionPolicy.distributeSupply(
+                    this.incomingDemands,
+                    new ArrayList<>(this.currentIncomingSupplies.values()),
+                    this.totalIncomingSupply);
 
             for (int idx = 0; idx < this.consumerEdges.size(); idx++) {
                 this.pushOutgoingSupply(this.consumerEdges.get(idx), supplies[idx], this.getConsumerResourceType());
@@ -151,13 +170,18 @@ public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsu
         this.consumerEdges.add(consumerEdge);
         this.incomingDemands.add(0.0);
         this.outgoingSupplies.add(0.0);
+        this.consumerResourceType = consumerEdge.getConsumerResourceType();
     }
 
     @Override
     public void addSupplierEdge(FlowEdge supplierEdge) {
-        this.supplierEdge = supplierEdge;
-        this.capacity = supplierEdge.getCapacity();
-        this.currentIncomingSupply = 0;
+        // supplierIndex not always set, so we use 0 as default to avoid index out of bounds
+        int idx = supplierEdge.getSupplierIndex() == -1 ? 0 : supplierEdge.getSupplierIndex();
+
+        this.supplierEdges.put(idx, supplierEdge);
+        this.capacity += supplierEdge.getCapacity();
+        this.currentIncomingSupplies.put(idx, 0.0);
+        this.supplierResourceType = supplierEdge.getSupplierResourceType();
     }
 
     @Override
@@ -202,13 +226,16 @@ public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsu
 
     @Override
     public void removeSupplierEdge(FlowEdge supplierEdge) {
-        this.supplierEdge = null;
-        this.capacity = 0;
-        this.currentIncomingSupply = 0;
-
-        this.updatedDemands.clear();
-
-        this.closeNode();
+        // supplierIndex not always set, so we use 0 as default to avoid index out of bounds
+        int idx = supplierEdge.getSupplierIndex() == -1 ? 0 : supplierEdge.getSupplierIndex();
+        // to keep index consistent, entries are neutralized instead of removed
+        this.supplierEdges.put(idx, null);
+        this.capacity -= supplierEdge.getCapacity();
+        this.currentIncomingSupplies.put(idx, 0.0);
+
+        if (this.supplierEdges.isEmpty()) {
+            this.updatedDemands.clear();
+        }
     }
 
     @Override
@@ -216,7 +243,7 @@ public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsu
         int idx = consumerEdge.getConsumerIndex();
 
         if (idx == -1) {
-            System.out.println("Error (FlowDistributor): Demand pushed by an unknown consumer");
+            LOGGER.warn("Demand {} pushed by an unknown consumer", newDemand);
             return;
         }
 
@@ -224,6 +251,7 @@ public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsu
         double prevDemand = incomingDemands.get(idx);
 
         incomingDemands.set(idx, newDemand);
+        // only update the total supply if the new supply is different from the previous one
         this.totalIncomingDemand += (newDemand - prevDemand);
 
         this.updatedDemands.add(idx);
@@ -243,14 +271,20 @@ public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsu
 
     @Override
     public void handleIncomingSupply(FlowEdge supplierEdge, double newSupply) {
-        this.currentIncomingSupply = newSupply;
+        // supplierIndex not always set, so we use 0 as default to avoid index out of bounds
+        int idx = supplierEdge.getSupplierIndex() == -1 ? 0 : supplierEdge.getSupplierIndex();
+        double prevSupply = currentIncomingSupplies.get(idx);
+
+        currentIncomingSupplies.put(idx, newSupply);
+        // only update the total supply if the new supply is different from the previous one
+        this.totalIncomingSupply += (newSupply - prevSupply);
 
         this.invalidate();
     }
 
     @Override
     public void pushOutgoingDemand(FlowEdge supplierEdge, double newDemand) {
-        this.supplierEdge.pushDemand(newDemand, false, this.getSupplierResourceType());
+        supplierEdge.pushDemand(newDemand, false, this.getSupplierResourceType());
     }
 
     @Override
@@ -267,23 +301,25 @@ public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsu
 
         outgoingSupplies.set(idx, newSupply);
         consumerEdge.pushSupply(newSupply, false, this.getSupplierResourceType());
-        consumerEdge.pushSupply(newSupply, false, this.getSupplierResourceType());
     }
 
     @Override
     public Map<FlowEdge.NodeType, List<FlowEdge>> getConnectedEdges() {
-        List<FlowEdge> supplyingEdges = (this.supplierEdge != null) ? List.of(this.supplierEdge) : List.of();
-
-        return Map.of(FlowEdge.NodeType.CONSUMING, supplyingEdges, FlowEdge.NodeType.SUPPLYING, this.consumerEdges);
+        return Map.of(
+                FlowEdge.NodeType.CONSUMING,
+                this.consumerEdges,
+                FlowEdge.NodeType.SUPPLYING,
+                new ArrayList<>(this.supplierEdges.values()));
     }
 
     @Override
     public ResourceType getSupplierResourceType() {
-        return this.supplierEdge.getSupplierResourceType();
+        //        return this.supplierEdge.getSupplierResourceType();
+        return this.supplierResourceType;
     }
 
     @Override
     public ResourceType getConsumerResourceType() {
-        return this.consumerEdges.getFirst().getConsumerResourceType();
+        return this.consumerResourceType;
     }
 }
diff --git a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/FlowEdge.java b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/FlowEdge.java
index aa3894c1..db2a2944 100644
--- a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/FlowEdge.java
+++ b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/FlowEdge.java
@@ -54,6 +54,15 @@ public class FlowEdge {
     }
 
     public FlowEdge(FlowConsumer consumer, FlowSupplier supplier, ResourceType resourceType) {
+        this(consumer, supplier, resourceType, -1, -1);
+    }
+
+    public FlowEdge(
+            FlowConsumer consumer,
+            FlowSupplier supplier,
+            ResourceType resourceType,
+            int consumerIndex,
+            int supplierIndex) {
         if (!(consumer instanceof FlowNode)) {
             throw new IllegalArgumentException("Flow consumer is not a FlowNode");
         }
@@ -67,6 +76,10 @@ public class FlowEdge {
 
         this.capacity = supplier.getCapacity(resourceType);
 
+        // to avoid race condition of setting indices and requiring them in the PSU
+        this.supplierIndex = supplierIndex;
+        this.consumerIndex = consumerIndex;
+
         this.consumer.addSupplierEdge(this);
         this.supplier.addConsumerEdge(this);
     }
diff --git a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/DistributionPolicy.java b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/DistributionPolicy.java
index 9d2246cd..3a8bebbc 100644
--- a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/DistributionPolicy.java
+++ b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/DistributionPolicy.java
@@ -25,5 +25,5 @@ package org.opendc.simulator.engine.graph.distributionPolicies;
 import java.util.ArrayList;
 
 public interface DistributionPolicy {
-    double[] distributeSupply(ArrayList<Double> supply, double currentSupply);
+    double[] distributeSupply(ArrayList<Double> supply, ArrayList<Double> currentSupply, double totalSupply);
 }
diff --git a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/FixedShare.java b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/FixedShare.java
index 40d70b5e..baa04975 100644
--- a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/FixedShare.java
+++ b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/FixedShare.java
@@ -42,7 +42,7 @@ public class FixedShare implements DistributionPolicy {
     }
 
     @Override
-    public double[] distributeSupply(ArrayList<Double> supply, double currentSupply) {
+    public double[] distributeSupply(ArrayList<Double> supply, ArrayList<Double> currentSupply, double totalSupply) {
         return new double[0];
     }
 }
diff --git a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/MaxMinFairnessPolicy.java b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/MaxMinFairnessPolicy.java
index 1d387349..484e7fe4 100644
--- a/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/MaxMinFairnessPolicy.java
+++ b/opendc-simulator/opendc-simulator-flow/src/main/java/org/opendc/simulator/engine/graph/distributionPolicies/MaxMinFairnessPolicy.java
@@ -34,7 +34,7 @@ public class MaxMinFairnessPolicy implements DistributionPolicy {
     private record Demand(int idx, double value) {}
 
     @Override
-    public double[] distributeSupply(ArrayList<Double> demands, double currentSupply) {
+    public double[] distributeSupply(ArrayList<Double> demands, ArrayList<Double> currentSupply, double totalSupply) {
         int inputSize = demands.size();
 
         final double[] supplies = new double[inputSize];
@@ -50,7 +50,7 @@ public class MaxMinFairnessPolicy implements DistributionPolicy {
             return i1.compareTo(i2);
         });
 
-        double availableCapacity = currentSupply; // totalSupply
+        double availableCapacity = totalSupply;
 
         for (int i = 0; i < inputSize; i++) {
             double d = tempDemands[i].value;