Added more battery policies (#312)

* some updates

* Updates

* Added comments and renamed variables

* Ran Spotless

16 files changed, 433 insertions, 42 deletions

diff --git a/opendc-compute/opendc-compute-simulator/src/main/java/org/opendc/compute/simulator/service/ComputeService.java b/opendc-compute/opendc-compute-simulator/src/main/java/org/opendc/compute/simulator/service/ComputeService.java
index aa2c13b9..1fb9c6a2 100644
--- a/opendc-compute/opendc-compute-simulator/src/main/java/org/opendc/compute/simulator/service/ComputeService.java
+++ b/opendc-compute/opendc-compute-simulator/src/main/java/org/opendc/compute/simulator/service/ComputeService.java
@@ -148,8 +148,6 @@ public final class ComputeService implements AutoCloseable {
      */
     private final Map<UUID, ServiceTask> taskById = new HashMap<>();
 
-    private final List<ServiceTask> tasks = new ArrayList<>();
-
     private final List<ServiceTask> tasksToRemove = new ArrayList<>();
 
     private ComputeMetricReader metricReader;
@@ -265,8 +263,8 @@ public final class ComputeService implements AutoCloseable {
     /**
      * Return the {@link ServiceTask}s hosted by this service.
      */
-    public List<ServiceTask> getTasks() {
-        return Collections.unmodifiableList(tasks);
+    public Map<UUID, ServiceTask> getTasks() {
+        return taskById;
     }
 
     /**
@@ -433,7 +431,6 @@ public final class ComputeService implements AutoCloseable {
     void delete(ServiceTask task) {
         completedTasks.remove(task);
         taskById.remove(task.getUid());
-        tasks.remove(task);
     }
 
     /**
@@ -452,7 +449,6 @@ public final class ComputeService implements AutoCloseable {
      * Run a single scheduling iteration.
      */
     private void doSchedule() {
-        // reorder tasks
 
         for (Iterator<SchedulingRequest> iterator = taskQueue.iterator();
                 iterator.hasNext();
@@ -660,7 +656,6 @@ public final class ComputeService implements AutoCloseable {
             ServiceTask task = new ServiceTask(service, uid, name, internalFlavor, workload, meta);
 
             service.taskById.put(uid, task);
-            service.tasks.add(task);
 
             service.tasksTotal++;
 
@@ -675,13 +670,6 @@ public final class ComputeService implements AutoCloseable {
             return service.taskById.get(id);
         }
 
-        @NotNull
-        public List<ServiceTask> queryTasks() {
-            checkOpen();
-
-            return new ArrayList<>(service.tasks);
-        }
-
         public void close() {
             isClosed = true;
         }
diff --git a/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/provisioner/HostsProvisioningStep.kt b/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/provisioner/HostsProvisioningStep.kt
index 68395fe4..211f33fe 100644
--- a/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/provisioner/HostsProvisioningStep.kt
+++ b/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/provisioner/HostsProvisioningStep.kt
@@ -27,11 +27,11 @@ import org.opendc.compute.simulator.host.SimHost
 import org.opendc.compute.simulator.service.ComputeService
 import org.opendc.compute.topology.specs.ClusterSpec
 import org.opendc.compute.topology.specs.HostSpec
+import org.opendc.compute.topology.specs.createSimBatteryPolicy
 import org.opendc.simulator.compute.power.CarbonModel
 import org.opendc.simulator.compute.power.SimPowerSource
 import org.opendc.simulator.compute.power.batteries.BatteryAggregator
 import org.opendc.simulator.compute.power.batteries.SimBattery
-import org.opendc.simulator.compute.power.batteries.policy.SingleThresholdBatteryPolicy
 import org.opendc.simulator.engine.engine.FlowEngine
 import org.opendc.simulator.engine.graph.FlowDistributor
 
@@ -99,13 +99,12 @@ public class HostsProvisioningStep internal constructor(
                 // Create Aggregator
                 val batteryAggregator = BatteryAggregator(graph, battery, batteryDistributor)
 
-                // Create BatteryPolicy
                 val batteryPolicy =
-                    SingleThresholdBatteryPolicy(
+                    createSimBatteryPolicy(
+                        cluster.battery!!.batteryPolicy,
                         graph,
                         battery,
                         batteryAggregator,
-                        cluster.battery!!.batteryPolicy.carbonThreshold,
                     )
 
                 carbonModel?.addReceiver(batteryPolicy)
diff --git a/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/telemetry/ComputeMetricReader.kt b/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/telemetry/ComputeMetricReader.kt
index 91748454..c88b98ea 100644
--- a/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/telemetry/ComputeMetricReader.kt
+++ b/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/telemetry/ComputeMetricReader.kt
@@ -140,7 +140,7 @@ public class ComputeMetricReader(
             }
 
             if (toMonitor[OutputFiles.TASK] == true) {
-                for (task in this.service.tasks) {
+                for (task in this.service.tasks.values) {
                     val reader =
                         this.taskTableReaders.computeIfAbsent(task) {
                             TaskTableReaderImpl(
@@ -197,7 +197,7 @@ public class ComputeMetricReader(
                 monitor.record(this.serviceTableReader.copy())
             }
 
-            if (loggCounter >= 100) {
+            if (loggCounter >= 24) {
                 var loggString = "\n\t\t\t\t\tMetrics after ${now.toEpochMilli() / 1000 / 60 / 60} hours:\n"
                 loggString += "\t\t\t\t\t\tTasks Total: ${this.serviceTableReader.tasksTotal}\n"
                 loggString += "\t\t\t\t\t\tTasks Active: ${this.serviceTableReader.tasksActive}\n"
diff --git a/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/TopologyFactories.kt b/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/TopologyFactories.kt
index b62c457c..a20bc2c2 100644
--- a/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/TopologyFactories.kt
+++ b/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/TopologyFactories.kt
@@ -24,7 +24,7 @@
 
 package org.opendc.compute.topology
 
-import org.opendc.compute.topology.specs.BatterySpec
+import org.opendc.compute.topology.specs.BatteryJSONSpec
 import org.opendc.compute.topology.specs.ClusterJSONSpec
 import org.opendc.compute.topology.specs.ClusterSpec
 import org.opendc.compute.topology.specs.HostJSONSpec
@@ -134,10 +134,10 @@ private fun ClusterJSONSpec.toClusterSpec(): ClusterSpec {
             carbonTracePath = this.powerSource.carbonTracePath,
         )
 
-    var batterySpec: BatterySpec? = null
+    var batterySpec: BatteryJSONSpec? = null
     if (this.battery != null) {
         batterySpec =
-            BatterySpec(
+            BatteryJSONSpec(
                 createUniqueName(this.battery.name, batteryNames),
                 this.battery.capacity,
                 this.battery.chargingSpeed,
diff --git a/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/specs/ClusterSpec.kt b/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/specs/ClusterSpec.kt
index 6e7c8dfa..97e12637 100644
--- a/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/specs/ClusterSpec.kt
+++ b/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/specs/ClusterSpec.kt
@@ -26,5 +26,5 @@ public data class ClusterSpec(
     val name: String,
     val hostSpecs: List<HostSpec>,
     val powerSource: PowerSourceSpec,
-    val battery: BatterySpec? = null,
+    val battery: BatteryJSONSpec? = null,
 )
diff --git a/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/specs/TopologySpecs.kt b/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/specs/TopologySpecs.kt
index 7b309bc3..3d8b63dc 100644
--- a/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/specs/TopologySpecs.kt
+++ b/opendc-compute/opendc-compute-topology/src/main/kotlin/org/opendc/compute/topology/specs/TopologySpecs.kt
@@ -22,10 +22,19 @@
 
 package org.opendc.compute.topology.specs
 
+import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 import org.opendc.common.units.DataSize
 import org.opendc.common.units.Frequency
 import org.opendc.common.units.Power
+import org.opendc.simulator.compute.power.batteries.BatteryAggregator
+import org.opendc.simulator.compute.power.batteries.SimBattery
+import org.opendc.simulator.compute.power.batteries.policy.BatteryPolicy
+import org.opendc.simulator.compute.power.batteries.policy.DoubleThresholdBatteryPolicy
+import org.opendc.simulator.compute.power.batteries.policy.RunningMeanBatteryPolicy
+import org.opendc.simulator.compute.power.batteries.policy.RunningMeanPlusBatteryPolicy
+import org.opendc.simulator.compute.power.batteries.policy.SingleThresholdBatteryPolicy
+import org.opendc.simulator.engine.graph.FlowGraph
 
 /**
  * Definition of a Topology modeled in the simulation.
@@ -182,6 +191,87 @@ public data class BatteryJSONSpec(
 )
 
 @Serializable
-public data class BatteryPolicyJSONSpec(
+public sealed interface BatteryPolicyJSONSpec
+
+@Serializable
+@SerialName("single")
+public data class SingleBatteryPolicyJSONSpec(
     val carbonThreshold: Double,
-)
+) : BatteryPolicyJSONSpec
+
+@Serializable
+@SerialName("double")
+public data class DoubleBatteryPolicyJSONSpec(
+    val lowerThreshold: Double,
+    val upperThreshold: Double,
+) : BatteryPolicyJSONSpec
+
+@Serializable
+@SerialName("runningMean")
+public data class RunningMeanPolicyJSONSpec(
+    val startingThreshold: Double,
+    val windowSize: Int,
+) : BatteryPolicyJSONSpec
+
+@Serializable
+@SerialName("runningMeanPlus")
+public data class RunningMeanPlusPolicyJSONSpec(
+    val startingThreshold: Double,
+    val windowSize: Int,
+) : BatteryPolicyJSONSpec
+
+@Serializable
+@SerialName("runningMedian")
+public data class RunningMedianPolicyJSONSpec(
+    val startingThreshold: Double,
+    val windowSize: Int,
+) : BatteryPolicyJSONSpec
+
+@Serializable
+@SerialName("runningQuartiles")
+public data class RunningQuartilesPolicyJSONSpec(
+    val startingThreshold: Double,
+    val windowSize: Int,
+) : BatteryPolicyJSONSpec
+
+public fun createSimBatteryPolicy(
+    batterySpec: BatteryPolicyJSONSpec,
+    graph: FlowGraph,
+    battery: SimBattery,
+    batteryAggregator: BatteryAggregator,
+): BatteryPolicy {
+    return when (batterySpec) {
+        is SingleBatteryPolicyJSONSpec ->
+            SingleThresholdBatteryPolicy(
+                graph,
+                battery,
+                batteryAggregator,
+                batterySpec.carbonThreshold,
+            )
+        is DoubleBatteryPolicyJSONSpec ->
+            DoubleThresholdBatteryPolicy(
+                graph,
+                battery,
+                batteryAggregator,
+                batterySpec.lowerThreshold,
+                batterySpec.upperThreshold,
+            )
+        is RunningMeanPolicyJSONSpec ->
+            RunningMeanBatteryPolicy(
+                graph,
+                battery,
+                batteryAggregator,
+                batterySpec.startingThreshold,
+                batterySpec.windowSize,
+            )
+        is RunningMeanPlusPolicyJSONSpec ->
+            RunningMeanPlusBatteryPolicy(
+                graph,
+                battery,
+                batteryAggregator,
+                batterySpec.startingThreshold,
+                batterySpec.windowSize,
+            )
+        else -> throw IllegalArgumentException("Unknown battery policy")
+    }
+}
diff --git a/opendc-compute/opendc-compute-workload/src/main/kotlin/org/opendc/compute/workload/WorkloadLoader.kt b/opendc-compute/opendc-compute-workload/src/main/kotlin/org/opendc/compute/workload/WorkloadLoader.kt
index 0b551a3c..ecb98ba2 100644
--- a/opendc-compute/opendc-compute-workload/src/main/kotlin/org/opendc/compute/workload/WorkloadLoader.kt
+++ b/opendc-compute/opendc-compute-workload/src/main/kotlin/org/opendc/compute/workload/WorkloadLoader.kt
@@ -54,9 +54,9 @@ public abstract class WorkloadLoader(private val submissionTime: String? = null)
 
         reScheduleTasks(workload)
 
-        if (fraction >= 1.0) {
-            return workload
-        }
+//        if (fraction >= 1.0) {
+//            return workload
+//        }
 
         if (fraction <= 0.0) {
             throw Error("The fraction of tasks to load cannot be 0.0 or lower")
@@ -65,21 +65,29 @@ public abstract class WorkloadLoader(private val submissionTime: String? = null)
         val res = mutableListOf<Task>()
 
         val totalLoad = workload.sumOf { it.totalLoad }
+        val desiredLoad = totalLoad * fraction
         var currentLoad = 0.0
 
-        val shuffledWorkload = workload.shuffled()
-        for (entry in shuffledWorkload) {
-            val entryLoad = entry.totalLoad
-
-            // TODO: ask Sacheen
-            if ((currentLoad + entryLoad) / totalLoad > fraction) {
-                break
-            }
-
-            currentLoad += entryLoad
+        while (currentLoad < desiredLoad) {
+            val entry = workload.random()
             res += entry
+
+            currentLoad += entry.totalLoad
         }
 
+//        val shuffledWorkload = workload.shuffled()
+//        for (entry in shuffledWorkload) {
+//            val entryLoad = entry.totalLoad
+//
+//            // TODO: ask Sacheen
+//            if ((currentLoad + entryLoad) / totalLoad > fraction) {
+//                break
+//            }
+//
+//            currentLoad += entryLoad
+//            res += entry
+//        }
+
         logger.info { "Sampled ${workload.size} VMs (fraction $fraction) into subset of ${res.size} VMs" }
 
         return res
diff --git a/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment1.json b/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment1.json
index 10ceaf87..8835faeb 100644
--- a/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment1.json
+++ b/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment1.json
@@ -31,6 +31,7 @@
                 "chargingSpeed": 1000,
                 "batteryPolicy":
                 {
+                    "type": "single",
                     "carbonThreshold": 90
                 }
             }
diff --git a/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment2.json b/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment2.json
index f89e9fa4..8882af09 100644
--- a/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment2.json
+++ b/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment2.json
@@ -31,6 +31,7 @@
                 "chargingSpeed": 1000,
                 "batteryPolicy":
                 {
+                    "type": "single",
                     "carbonThreshold": 120
                 }
             }
diff --git a/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment3.json b/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment3.json
index 920e09df..d78626f1 100644
--- a/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment3.json
+++ b/opendc-experiments/opendc-experiments-base/src/test/resources/topologies/batteries/experiment3.json
@@ -31,6 +31,7 @@
                 "chargingSpeed": 1000,
                 "batteryPolicy":
                 {
+                    "type": "single",
                     "carbonThreshold": 100
                 }
             }
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/BatteryPolicy.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/BatteryPolicy.java
index 5abbe861..a50f7e73 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/BatteryPolicy.java
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/BatteryPolicy.java
@@ -31,12 +31,16 @@ import org.opendc.simulator.compute.power.batteries.SimBattery;
 import org.opendc.simulator.engine.graph.FlowGraph;
 import org.opendc.simulator.engine.graph.FlowNode;
 
+/**
+ * An abstract class representing a battery policy.
+ * A battery policy is used by a {@link SimBattery} to determine when to charge or discharge the battery.
+ */
 public abstract class BatteryPolicy extends FlowNode implements CarbonReceiver {
 
     protected final SimBattery battery;
     protected final BatteryAggregator aggregator;
 
-    protected double carbonIntensity;
+    protected double carbonIntensity; // The current carbon Intensity of the grid
 
     protected BatteryState batteryState = BatteryState.IDLE;
 
@@ -61,6 +65,12 @@ public abstract class BatteryPolicy extends FlowNode implements CarbonReceiver {
     @Override
     public abstract long onUpdate(long now);
 
+    /**
+     * Set the battery state.
+     * Both the battery and the aggregator are updated based on the new state.
+     *
+     * @param newBatteryState The new battery state.
+     */
     public void setBatteryState(BatteryState newBatteryState) {
         if (newBatteryState == this.batteryState) {
             return;
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/DoubleThresholdBatteryPolicy.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/DoubleThresholdBatteryPolicy.java
new file mode 100644
index 00000000..3a9cb228
--- /dev/null
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/DoubleThresholdBatteryPolicy.java
@@ -0,0 +1,79 @@
+/*
+ * Copyright (c) 2025 AtLarge Research
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package org.opendc.simulator.compute.power.batteries.policy;
+
+import org.opendc.simulator.compute.power.batteries.BatteryAggregator;
+import org.opendc.simulator.compute.power.batteries.BatteryState;
+import org.opendc.simulator.compute.power.batteries.SimBattery;
+import org.opendc.simulator.engine.graph.FlowGraph;
+
+/**
+ * A battery policy that uses two thresholds to determine if a better should be charging or discharging.
+ * - If the Carbon Intensity is below the lower threshold,
+ *    the battery will start charging until full.
+ * - If the Carbon Intensity is above the upper threshold,
+ *    the battery will start discharging until empty.
+ * - If the Carbon Intensity is between the two thresholds,
+ *    The battery is idle.
+ */
+public class DoubleThresholdBatteryPolicy extends BatteryPolicy {
+    private final double lowerThreshold;
+    private final double upperThreshold;
+
+    /**
+     *
+     * @param parentGraph     The {@link FlowGraph} this stage belongs to.
+     * @param battery        The {@link SimBattery} to control.
+     * @param aggregator    The {@link BatteryAggregator} to use.
+     * @param lowerThreshold The lower carbon intensity threshold to trigger charging or discharging.
+     * @param upperThreshold The upper carbon intensity threshold to trigger charging or discharging.
+     */
+    public DoubleThresholdBatteryPolicy(
+            FlowGraph parentGraph,
+            SimBattery battery,
+            BatteryAggregator aggregator,
+            double lowerThreshold,
+            double upperThreshold) {
+        super(parentGraph, battery, aggregator);
+
+        this.lowerThreshold = lowerThreshold;
+        this.upperThreshold = upperThreshold;
+    }
+
+    @Override
+    public long onUpdate(long now) {
+
+        if (this.carbonIntensity > this.upperThreshold & !this.battery.isEmpty()) {
+            this.setBatteryState(BatteryState.DISCHARGING);
+            return Long.MAX_VALUE;
+        }
+
+        if (this.carbonIntensity < this.lowerThreshold & !this.battery.isFull()) {
+            this.setBatteryState(BatteryState.CHARGING);
+            return Long.MAX_VALUE;
+        }
+
+        this.setBatteryState(BatteryState.IDLE);
+        return Long.MAX_VALUE;
+    }
+}
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/RunningMeanBatteryPolicy.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/RunningMeanBatteryPolicy.java
new file mode 100644
index 00000000..1c127abd
--- /dev/null
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/RunningMeanBatteryPolicy.java
@@ -0,0 +1,103 @@
+/*
+ * Copyright (c) 2025 AtLarge Research
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package org.opendc.simulator.compute.power.batteries.policy;
+
+import java.util.LinkedList;
+import org.opendc.simulator.compute.power.batteries.BatteryAggregator;
+import org.opendc.simulator.compute.power.batteries.BatteryState;
+import org.opendc.simulator.compute.power.batteries.SimBattery;
+import org.opendc.simulator.engine.graph.FlowGraph;
+
+/**
+ * A battery policy that uses a running mean to determine if a battery should be charging or discharging.
+ * This policy is similar to {@link SingleThresholdBatteryPolicy}, but instead of using a predifined threshold,
+ * the threshold is updated dynamically based on the running mean of the carbon intensity.
+ * - If the Carbon Intensity is below the running mean,
+ *   the battery will start charging until full.
+ * - If the Carbon Intensity is above the running mean,
+ *   the battery will start discharging until empty.
+ */
+public class RunningMeanBatteryPolicy extends BatteryPolicy {
+    private final int windowSize;
+
+    private final LinkedList<Double> pastCarbonIntensities = new LinkedList<>();
+    private double pastCarbonIntensitiesSum = 0.0;
+    private double pastCarbonIntensitiesMean = 0.0;
+
+    /**
+     *
+     * @param parentGraph     The {@link FlowGraph} this stage belongs to.
+     * @param battery        The {@link SimBattery} to control.
+     * @param aggregator    The {@link BatteryAggregator} to use.
+     */
+    public RunningMeanBatteryPolicy(
+            FlowGraph parentGraph,
+            SimBattery battery,
+            BatteryAggregator aggregator,
+            double startingThreshold,
+            int windowSize) {
+        super(parentGraph, battery, aggregator);
+
+        this.windowSize = windowSize;
+
+        this.updatePastCarbonIntensities(startingThreshold);
+    }
+
+    /**
+     * Update the past carbon intensities with the new carbon intensity.
+     *
+     * Update the current sum and mean.
+     * @param newCarbonIntensity
+     */
+    private void updatePastCarbonIntensities(double newCarbonIntensity) {
+        if (this.pastCarbonIntensities.size() == this.windowSize) {
+            this.pastCarbonIntensitiesSum -= this.pastCarbonIntensities.removeFirst();
+        }
+        this.pastCarbonIntensities.addLast(newCarbonIntensity);
+        this.pastCarbonIntensitiesSum += newCarbonIntensity;
+        this.pastCarbonIntensitiesMean = this.pastCarbonIntensitiesSum / this.pastCarbonIntensities.size();
+    }
+
+    @Override
+    public void updateCarbonIntensity(double newCarbonIntensity) {
+        this.updatePastCarbonIntensities(newCarbonIntensity);
+
+        super.updateCarbonIntensity(newCarbonIntensity);
+    }
+
+    @Override
+    public long onUpdate(long now) {
+        if (this.carbonIntensity >= this.pastCarbonIntensitiesMean & !this.battery.isEmpty()) {
+            this.setBatteryState(BatteryState.DISCHARGING);
+            return Long.MAX_VALUE;
+        }
+
+        if (this.carbonIntensity < this.pastCarbonIntensitiesMean & !this.battery.isFull()) {
+            this.setBatteryState(BatteryState.CHARGING);
+            return Long.MAX_VALUE;
+        }
+
+        this.setBatteryState(BatteryState.IDLE);
+        return Long.MAX_VALUE;
+    }
+}
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/RunningMeanPlusBatteryPolicy.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/RunningMeanPlusBatteryPolicy.java
new file mode 100644
index 00000000..25b86dde
--- /dev/null
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/RunningMeanPlusBatteryPolicy.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright (c) 2025 AtLarge Research
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package org.opendc.simulator.compute.power.batteries.policy;
+
+import java.util.LinkedList;
+import org.opendc.simulator.compute.power.batteries.BatteryAggregator;
+import org.opendc.simulator.compute.power.batteries.BatteryState;
+import org.opendc.simulator.compute.power.batteries.SimBattery;
+import org.opendc.simulator.engine.graph.FlowGraph;
+
+/**
+ * An improved version of {@link RunningMeanBatteryPolicy}.
+ * It uses the same logic, but only start charging if the carbon intensity is not decreasing anymore.
+ */
+public class RunningMeanPlusBatteryPolicy extends BatteryPolicy {
+    private final int windowSize;
+
+    private final LinkedList<Double> pastCarbonIntensities = new LinkedList<>();
+    private double previousCarbonIntensity = 0.0;
+
+    private double pastCarbonIntensitiesSum = 0.0;
+    private double pastCarbonIntensitiesMean = 0.0;
+
+    /**
+     *
+     * @param parentGraph     The {@link FlowGraph} this stage belongs to.
+     * @param battery        The {@link SimBattery} to control.
+     * @param aggregator    The {@link BatteryAggregator} to use.
+     */
+    public RunningMeanPlusBatteryPolicy(
+            FlowGraph parentGraph,
+            SimBattery battery,
+            BatteryAggregator aggregator,
+            double startingThreshold,
+            int windowSize) {
+        super(parentGraph, battery, aggregator);
+
+        this.windowSize = windowSize;
+
+        this.updatePastCarbonIntensities(startingThreshold);
+    }
+
+    private void updatePastCarbonIntensities(double newCarbonIntensity) {
+        if (this.pastCarbonIntensities.size() == this.windowSize) {
+            this.pastCarbonIntensitiesSum -= this.pastCarbonIntensities.removeFirst();
+        }
+
+        if (this.pastCarbonIntensities.size() > 0) {
+            this.previousCarbonIntensity = this.pastCarbonIntensities.getLast();
+        }
+
+        this.pastCarbonIntensities.addLast(newCarbonIntensity);
+        this.pastCarbonIntensitiesSum += newCarbonIntensity;
+        this.pastCarbonIntensitiesMean = this.pastCarbonIntensitiesSum / this.pastCarbonIntensities.size();
+    }
+
+    @Override
+    public void updateCarbonIntensity(double newCarbonIntensity) {
+        this.updatePastCarbonIntensities(newCarbonIntensity);
+
+        super.updateCarbonIntensity(newCarbonIntensity);
+    }
+
+    private boolean isCharging = false;
+
+    @Override
+    public long onUpdate(long now) {
+        if (this.carbonIntensity >= this.pastCarbonIntensitiesMean & !this.battery.isEmpty()) {
+            this.isCharging = false;
+            this.setBatteryState(BatteryState.DISCHARGING);
+            return Long.MAX_VALUE;
+        }
+
+        if (this.carbonIntensity < this.pastCarbonIntensitiesMean & !this.battery.isFull()) {
+            if (this.carbonIntensity >= this.previousCarbonIntensity || this.isCharging) {
+                this.setBatteryState(BatteryState.CHARGING);
+                return Long.MAX_VALUE;
+            }
+        }
+
+        this.isCharging = false;
+        this.setBatteryState(BatteryState.IDLE);
+        return Long.MAX_VALUE;
+    }
+}
diff --git a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/SingleThresholdBatteryPolicy.java b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/SingleThresholdBatteryPolicy.java
index e917a26f..26d85958 100644
--- a/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/SingleThresholdBatteryPolicy.java
+++ b/opendc-simulator/opendc-simulator-compute/src/main/java/org/opendc/simulator/compute/power/batteries/policy/SingleThresholdBatteryPolicy.java
@@ -27,11 +27,17 @@ import org.opendc.simulator.compute.power.batteries.BatteryState;
 import org.opendc.simulator.compute.power.batteries.SimBattery;
 import org.opendc.simulator.engine.graph.FlowGraph;
 
+/**
+ * A battery policy that uses a single threshold to determine if a better should be charging or discharging.
+ * - If the Carbon Intensity is below the give thresholds,
+ *   the battery will start charging until full.
+ * - If the Carbon Intensity is above the give thresholds,
+ *   the battery will start discharging until empty.
+ */
 public class SingleThresholdBatteryPolicy extends BatteryPolicy {
     private final double carbonThreshold;
 
     /**
-     *
      * @param parentGraph     The {@link FlowGraph} this stage belongs to.
      * @param battery        The {@link SimBattery} to control.
      * @param aggregator    The {@link BatteryAggregator} to use.
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 dcbd79bb..c094560e 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
@@ -208,7 +208,7 @@ public class FlowDistributor extends FlowNode implements FlowSupplier, FlowConsu
             other.setConsumerIndex(other.getConsumerIndex() - 1);
         }
 
-        HashSet newUpdatedDemands = new HashSet<>();
+        HashSet<Integer> newUpdatedDemands = new HashSet<>();
 
         for (int idx_other : this.updatedDemands) {
             if (idx_other > idx) {