Refactor `SimHost` for improved memory tracking, error handling, and task state management.

2 files changed, 38 insertions, 27 deletions

diff --git a/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/host/SimHost.kt b/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/host/SimHost.kt
index cb7d028c..164bbee7 100644
--- a/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/host/SimHost.kt
+++ b/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/host/SimHost.kt
@@ -229,8 +229,26 @@ public class SimHost(
         return this.guests.toList()
     }
 
+    /**
+     * Calculates the total memory used by the currently running tasks on the host.
+     *
+     * Iterates through the tasks mapped to guests in `taskToGuestMap`. For tasks that are in the
+     * `TaskState.RUNNING` state, their memory consumption is summed up.
+     *
+     * @return Total memory used by tasks currently in the RUNNING state, in bytes.
+     */
+    private fun usedMemoryByRunningTasks(): Long {
+        var usedMemory: Long = 0
+        for (vm in this.taskToGuestMap) {
+            if (vm.value.state == TaskState.RUNNING) {
+                usedMemory += vm.key.memorySize
+            }
+        }
+        return usedMemory
+    }
+
     public fun canFit(task: ServiceTask): Boolean {
-        val sufficientMemory = model.memoryCapacity >= task.memorySize
+        val sufficientMemory = (model.memoryCapacity - this.usedMemoryByRunningTasks()) >= task.memorySize
         val enoughCpus = model.coreCount >= task.cpuCoreCount
         val canFit = simMachine!!.canFit(task.toMachineModel())
 
@@ -307,17 +325,8 @@ public class SimHost(
         for (guest in guests) {
             when (guest.state) {
                 TaskState.RUNNING -> running++
-                TaskState.FAILED, TaskState.TERMINATED -> {
-                    failed++
-                    // Remove guests that have been deleted
-                    this.taskToGuestMap.remove(guest.task)
-                    guests.remove(guest)
-                }
-                TaskState.COMPLETED -> {
-                    completed++
-                    this.taskToGuestMap.remove(guest.task)
-                    guests.remove(guest)
-                }
+                TaskState.FAILED, TaskState.TERMINATED -> failed++
+                TaskState.COMPLETED -> completed++
                 TaskState.PAUSED -> {}
                 else -> invalid++
             }
@@ -337,8 +346,8 @@ public class SimHost(
         )
     }
 
-    public fun getSystemStats(task: ServiceTask): GuestSystemStats {
-        val guest = requireNotNull(taskToGuestMap[task]) { "Unknown task ${task.name} at host $name" }
+    public fun getSystemStats(task: ServiceTask): GuestSystemStats? {
+        val guest = taskToGuestMap[task] ?: return null
         return guest.getSystemStats()
     }
 
@@ -359,8 +368,8 @@ public class SimHost(
         )
     }
 
-    public fun getCpuStats(task: ServiceTask): GuestCpuStats {
-        val guest = requireNotNull(taskToGuestMap[task]) { "Unknown task ${task.name} at host $name" }
+    public fun getCpuStats(task: ServiceTask): GuestCpuStats? {
+        val guest = taskToGuestMap[task] ?: return null
         return guest.getCpuStats()
     }
 
@@ -388,7 +397,7 @@ public class SimHost(
     }
 
     public fun getGpuStats(task: ServiceTask): GuestGpuStats? {
-        val guest = requireNotNull(taskToGuestMap[task]) { "Unknown task ${task.name} at host $name" }
+        val guest = taskToGuestMap[task] ?: return null
         return guest.getGpuStats()
     }
 
diff --git a/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/telemetry/table/task/TaskTableReaderImpl.kt b/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/telemetry/table/task/TaskTableReaderImpl.kt
index be858e4f..5bf917dd 100644
--- a/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/telemetry/table/task/TaskTableReaderImpl.kt
+++ b/opendc-compute/opendc-compute-simulator/src/main/kotlin/org/opendc/compute/simulator/telemetry/table/task/TaskTableReaderImpl.kt
@@ -235,17 +235,19 @@ public class TaskTableReaderImpl(
      */
     override fun record(now: Instant) {
         val newHost = service.lookupHost(task)
-        if (newHost != null && newHost.getName() != simHost?.getName()) {
+        if (newHost != simHost) {
             simHost = newHost
-            hostInfo =
-                HostInfo(
-                    newHost.getName(),
-                    newHost.getClusterName(),
-                    "x86",
-                    newHost.getModel().coreCount,
-                    newHost.getModel().cpuCapacity,
-                    newHost.getModel().memoryCapacity,
-                )
+            if (newHost != null) {
+                hostInfo =
+                    HostInfo(
+                        newHost.getName(),
+                        newHost.getClusterName(),
+                        "x86",
+                        newHost.getModel().coreCount,
+                        newHost.getModel().cpuCapacity,
+                        newHost.getModel().memoryCapacity,
+                    )
+            }
         }
 
         val cpuStats = simHost?.getCpuStats(task)