Add basic scheduler for SMP CPUs (#33)

* sched: add basic scheduler implementation

Signed-off-by: Pawel Wieczorkiewicz <[email protected]>

* sched,smp: consider number of cpus for scheduling

Do not schedule tasks on non-existing CPUs.

Signed-off-by: Pawel Wieczorkiewicz <[email protected]>

include/lib.h

@@ -26,6 +26,7 @@
 #define KTF_LIB_H
 
 #include <ktf.h>
+#include <console.h>
 #include <segment.h>
 #include <asm-macros.h>
 

include/sched.h

@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) 2020 Amazon.com, Inc. or its affiliates.
+ * All Rights Reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ *    list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ *    this list of conditions and the following disclaimer in the documentation
+ *    and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
+ * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+#ifndef KTF_SCHED_H
+#define KTF_SCHED_H
+
+#include <ktf.h>
+#include <lib.h>
+#include <page.h>
+#include <list.h>
+
+typedef void (*task_func_t)(void *this, void *arg);
+
+enum task_state {
+    TASK_STATE_NEW,
+    TASK_STATE_READY,
+    TASK_STATE_SCHEDULED,
+    TASK_STATE_RUNNING,
+    TASK_STATE_DONE,
+};
+typedef enum task_state task_state_t;
+
+typedef unsigned int tid_t;
+
+struct task {
+    list_head_t list;
+
+    tid_t id;
+    task_state_t state;
+
+    unsigned int cpu;
+
+    const char *name;
+    task_func_t func;
+    void *arg;
+
+    unsigned long result;
+} __aligned(PAGE_SIZE);
+typedef struct task task_t;
+
+/* External declarations */
+
+extern void init_tasks(void);
+extern task_t *get_task_by_id(tid_t id);
+extern task_t *get_task_by_name(const char *name);
+extern task_t *get_task_for_cpu(unsigned int cpu);
+extern task_t *new_task(const char *name, task_func_t func, void *arg);
+extern void schedule_task(task_t *task, unsigned int cpu);
+extern void run_tasks(unsigned int cpu);
+extern void wait_for_all_tasks(void);
+
+#endif /* KTF_SCHED_H */

include/smp/smp.h

@@ -29,9 +29,12 @@
 #include <lib.h>
 #include <processor.h>
 
+#define INVALID_CPU (~0U)
+
 /* External declarations */
 
 extern void smp_init(void);
+extern unsigned get_nr_cpus(void);
 
 /* Static declarations */
 

kernel.c

@@ -28,6 +28,7 @@
 #include <setup.h>
 #include <multiboot.h>
 #include <percpu.h>
+#include <sched.h>
 
 extern void _long_to_real(void);
 
@@ -57,6 +58,8 @@ void kernel_main(void) {
 
     test_main();
 
+    printk("All tasks done.\n");
+
     while(1)
         halt();
 }

sched.c

@@ -0,0 +1,212 @@
+/*
+ * Copyright (c) 2020 Amazon.com, Inc. or its affiliates.
+ * All Rights Reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ *    list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ *    this list of conditions and the following disclaimer in the documentation
+ *    and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
+ * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+#include <ktf.h>
+#include <lib.h>
+#include <list.h>
+#include <sched.h>
+#include <string.h>
+#include <console.h>
+#include <spinlock.h>
+#include <setup.h>
+
+#include <smp/smp.h>
+
+#include <mm/vmm.h>
+
+static list_head_t tasks;
+static tid_t next_tid;
+
+static spinlock_t lock = SPINLOCK_INIT;
+
+static bool terminate;
+
+void init_tasks(void) {
+    printk("Initializing tasks\n");
+
+    list_init(&tasks);
+}
+
+static const char *task_state_names[] = {
+    [TASK_STATE_NEW]       = "NEW",
+    [TASK_STATE_READY]     = "READY",
+    [TASK_STATE_SCHEDULED] = "SCHEDULED",
+    [TASK_STATE_RUNNING]   = "RUNNING",
+    [TASK_STATE_DONE]      = "DONE",
+};
+
+static inline void set_task_state(task_t *task, task_state_t state) {
+    ASSERT(task);
+
+    dprintk("CPU[%u]: state transition %s -> %s\n", task->cpu,
+            task_state_names[task->state], task_state_names[state]);
+
+    ACCESS_ONCE(task->state) = state;
+    smp_mb();
+}
+
+static inline task_state_t get_task_state(task_t *task) {
+    task_state_t state;
+    ASSERT(task);
+
+    state = ACCESS_ONCE(task->state);
+    smp_rmb();
+
+    return state;
+}
+
+static task_t *create_task(void) {
+    task_t *task = get_free_page(GFP_KERNEL);
+
+    if (!task)
+        return NULL;
+
+    memset(task, 0, sizeof(*task));
+    task->id = next_tid++;
+    task->cpu = INVALID_CPU;
+    set_task_state(task, TASK_STATE_NEW);
+
+    spin_lock(&lock);
+    list_add(&task->list, &tasks);
+    spin_unlock(&lock);
+
+    return task;
+}
+
+static void prepare_task(task_t *task, const char *name, task_func_t func, void *arg) {
+    if (!task)
+        return;
+
+    BUG_ON(get_task_state(task) > TASK_STATE_READY);
+
+    task->name = name;
+    task->func = func;
+    task->arg = arg;
+    set_task_state(task, TASK_STATE_READY);
+}
+
+static void wait_for_task_state(task_t *task, task_state_t state) {
+    if (!task)
+        return;
+
+    while (get_task_state(task) != state)
+        cpu_relax();
+}
+
+task_t *new_task(const char *name, task_func_t func, void *arg) {
+    task_t *task = create_task();
+
+    if (!task)
+        return NULL;
+
+    prepare_task(task, name, func, arg);
+    return task;
+}
+
+task_t *get_task_by_id(tid_t id) {
+    task_t *task;
+
+    list_for_each_entry(task, &tasks, list) {
+        if (task->id == id)
+            return task;
+    }
+
+    return NULL;
+}
+
+task_t *get_task_by_name(const char *name) {
+    task_t *task;
+
+    list_for_each_entry(task, &tasks, list) {
+        if (!strcmp(task->name, name))
+            return task;
+    }
+
+    return NULL;
+}
+
+task_t *get_task_for_cpu(unsigned int cpu) {
+    task_t *task;
+
+    list_for_each_entry(task, &tasks, list) {
+        if (task->cpu == cpu)
+            return task;
+    }
+
+    return NULL;
+}
+
+void schedule_task(task_t *task, unsigned int cpu) {
+    ASSERT(task);
+
+    if (cpu > get_nr_cpus() - 1)
+        panic("CPU[%u] does not exist.\n", cpu);
+
+    BUG_ON(get_task_state(task) != TASK_STATE_READY);
+
+    printk("CPU[%u]: Scheduling task %s[%u]\n",
+           cpu, task->name, task->id);
+
+    task->cpu = cpu;
+    set_task_state(task, TASK_STATE_SCHEDULED);
+}
+
+static void run_task(task_t *task) {
+    if (!task)
+        return;
+
+    wait_for_task_state(task, TASK_STATE_SCHEDULED);
+
+    printk("CPU[%u]: Running task %s[%u]\n",
+           task->cpu, task->name, task->id);
+
+    set_task_state(task, TASK_STATE_RUNNING);
+    task->func(task, task->arg);
+    set_task_state(task, TASK_STATE_DONE);
+}
+
+void wait_for_all_tasks(void) {
+    task_t *task;
+    bool busy;
+
+    do {
+        busy = false;
+
+        list_for_each_entry(task, &tasks, list) {
+            if (get_task_state(task) != TASK_STATE_DONE) {
+                busy = true;
+                wait_for_task_state(task, TASK_STATE_DONE);
+            }
+        }
+        cpu_relax();
+    } while(busy && !terminate);
+}
+
+void run_tasks(unsigned int cpu) {
+    do {
+        run_task(get_task_for_cpu(cpu));
+        cpu_relax();
+    } while(!terminate);
+}
+

setup.c

@@ -34,6 +34,7 @@
 #include <multiboot.h>
 #include <apic.h>
 #include <percpu.h>
+#include <sched.h>
 
 #include <mm/pmm.h>
 #include <mm/vmm.h>
@@ -107,6 +108,8 @@ void __noreturn __text_init kernel_start(uint32_t multiboot_magic, multiboot_inf
 
     init_apic(APIC_MODE_XAPIC);
 
+    init_tasks();
+
     smp_init();
 
     /* Jump from .text.init section to .text */

smp/smp.c

@@ -29,6 +29,7 @@
 #include <apic.h>
 #include <traps.h>
 #include <percpu.h>
+#include <sched.h>
 
 #include <mm/vmm.h>
 
@@ -50,6 +51,8 @@ void __noreturn ap_startup(void) {
     ap_callin = true;
     smp_wmb();
 
+    run_tasks(smp_processor_id());
+
     while(true)
         halt();
 
@@ -100,3 +103,7 @@ void smp_init(void) {
     for (int i = 0; i < nr_cpus; i++)
         boot_cpu(i);
 }
+
+unsigned get_nr_cpus(void) {
+    return nr_cpus;
+}

test.c

@@ -24,6 +24,7 @@
  */
 #include <ktf.h>
 #include <console.h>
+#include <sched.h>
 
 static int __user_text func(void *arg) {
     return 0;
@@ -34,5 +35,7 @@ void test_main(void) {
 
     usermode_call(func, NULL);
 
+    wait_for_all_tasks();
+
     printk("Test done\n");
 }