implemented automatic load balancing in the scheduler

kernel/src/entry/tests/multithreading.test.cpp

@@ -48,7 +48,7 @@ DECLARE_UNIT_TEST("Multithreading Test - Kernel Task Creation", mtTaskCreationUn
 DECLARE_UNIT_TEST("Multithreading Test - Single Core", mtSingleCoreUnitTest) {
     const size_t taskCount = 1000;
     const int targetCpu = BSP_CPU_ID;
-    const int taskExecutionTimeout = 1200;
+    const int taskExecutionTimeout = 2000;
     auto& sched = Scheduler::get();
 
     // Allocate a buffer to store the tasks
@@ -91,10 +91,10 @@ DECLARE_UNIT_TEST("Multithreading Test - Single Core", mtSingleCoreUnitTest) {
     return UNIT_TEST_SUCCESS;
 }
 
-DECLARE_UNIT_TEST("Multithreading Test - Multi Core", mtMultiCoreUnitTest) {
+DECLARE_UNIT_TEST("Multithreading Test - Multi Core (Automatic Load Balancing)", mtMultiCoreUnitTest) {
     const size_t systemCpus = AcpiController::get().getApicTable()->getCpuCount();
-    const size_t taskCount = 1000 * systemCpus;
-    const uint32_t taskExecutionTimeout = 1200;
+    const size_t taskCount = 600 * systemCpus;
+    const uint32_t taskExecutionTimeout = 2000;
     auto& sched = Scheduler::get();
 
     // Allocate a buffer to store the tasks
@@ -115,10 +115,7 @@ DECLARE_UNIT_TEST("Multithreading Test - Multi Core", mtMultiCoreUnitTest) {
 
     // Schedule all the tasks
     for (size_t i = 0; i < taskCount; i++) {
-        size_t targetCpu = i % systemCpus;
-
-        sched.addTaskToCpu(taskArray[i], targetCpu);
-        // ASSERT_TRUE(ret, "Failed to schedule a task");
+        sched.addTask(taskArray[i]);
     }
 
     // Wait for all tasks to finish

kernel/src/sched/sched.cpp

@@ -15,6 +15,7 @@ Scheduler g_primaryScheduler;
 Task g_kernelSwapperTasks[MAX_CPUS] = {};
 
 DECLARE_SPINLOCK(__sched_pid_alloc_lock);
+DECLARE_SPINLOCK(__sched_load_balancing_lock);
 
 size_t g_availableTaskPid = 10;
 size_t _allocateTaskPid() {
@@ -104,6 +105,11 @@ void Scheduler::addTaskToCpu(Task* task, int cpu) {
     releaseSpinlock(&m_runQueues[cpu]->lock);
 }
 
+void Scheduler::addTask(Task* task) {
+    size_t targetCpu = _findLeastLoadedCpu();
+    addTaskToCpu(task, targetCpu);
+}
+
 void Scheduler::removeTaskFromCpu(int pid, int cpu) {
     acquireSpinlock(&m_runQueues[cpu]->lock);
 
@@ -118,6 +124,34 @@ void Scheduler::removeTaskFromCpu(int pid, int cpu) {
     releaseSpinlock(&m_runQueues[cpu]->lock);
 }
 
+int Scheduler::_findLeastLoadedCpu() {
+    int cpu = 0;
+    size_t leastTasksFound = static_cast<size_t>(-1);
+
+    acquireSpinlock(&__sched_load_balancing_lock);
+
+    for (size_t i = 0; i < m_runQueues.size(); i++) {
+        size_t cpuTasks = _getRunQueueSize(i);
+        if (cpuTasks < leastTasksFound) {
+            leastTasksFound = cpuTasks;
+            cpu = static_cast<int>(i);
+        }
+    }
+
+    releaseSpinlock(&__sched_load_balancing_lock);
+    return cpu;
+}
+
+size_t Scheduler::_getRunQueueSize(size_t cpu) {
+    size_t ret = 0;
+
+    acquireSpinlock(&m_runQueues[cpu]->lock);
+    ret = m_runQueues[cpu]->tasks.size();
+
+    releaseSpinlock(&m_runQueues[cpu]->lock);
+    return ret;
+}
+
 Task* createKernelTask(void (*taskEntry)(), int priority) {
     Task* task = (Task*)kmalloc(sizeof(Task));
     if (!task) {

kernel/src/sched/sched.h

@@ -22,6 +22,9 @@ class Scheduler {
 
     static Scheduler& get();
 
+    void init();
+    void registerCoreForScheduling(int cpu);
+
     // Get the current running task
     Task* getCurrentTask(int cpu);
 
@@ -31,17 +34,20 @@ class Scheduler {
     // Schedule the next task to run
     void scheduleNextTask(int cpu);
 
-    // Add a task
+    // Add a task to a specified cpu core
     void addTaskToCpu(Task* task, int cpu);
 
+    // Adds a task to the next least loaded available cpu core
+    void addTask(Task* task);
+
     // Remove a task
     void removeTaskFromCpu(int pid, int cpu);
 
-    void init();
-    void registerCoreForScheduling(int cpu);
-
 private:
     kstl::vector<kstl::SharedPtr<SchedulerRunQueue>> m_runQueues;
+
+    int _findLeastLoadedCpu();
+    size_t _getRunQueueSize(size_t cpu);
 };
 
 //