From 810ad4463429a54d3673ff027697785b2d0e08af Mon Sep 17 00:00:00 2001
From: Pawel Wieczorkiewicz <wipawel@grsecurity.net>
Date: Fri, 21 Jul 2023 14:28:07 +0200
Subject: [PATCH] sched: add task repetition functionality

By default tasks run once and get destroyed. It might be useful to run
them in an infinite loop (task done -> task scheduled -> ...) in order
to achieve tasks interleaving (multiple tasks scheduled on the same
CPU).
It might be also useful to set an arbitrary value of repetitions for a
given task.

This commit adds simple set_task_repeat() and set_task_loop() API to
allow task repetitions.

Signed-off-by: Pawel Wieczorkiewicz <wipawel@grsecurity.net>
---
 common/sched.c    | 45 +++++++++++++++++++++++++++++++++++++++------
 include/sched.h   | 21 +++++++++++++++++++++
 tests/unittests.c |  1 +
 3 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/common/sched.c b/common/sched.c
index 4d8fa50f..6a5bc6cc 100644
--- a/common/sched.c
+++ b/common/sched.c
@@ -87,6 +87,8 @@ static task_t *create_task(void) {
     task->id = next_tid++;
     task->gid = TASK_GROUP_ALL;
     set_task_state(task, TASK_STATE_NEW);
+    atomic_set(&task->exec_count, 0);
+    set_task_once(task);
 
     return task;
 }
@@ -155,6 +157,20 @@ task_t *get_task_by_name(cpu_t *cpu, const char *name) {
     return NULL;
 }
 
+static const char *task_repeat_string(task_repeat_t repeat) {
+    static char buf[20] = {0};
+
+    switch (repeat) {
+    case TASK_REPEAT_ONCE:
+        return "ONCE";
+    case TASK_REPEAT_LOOP:
+        return "LOOP";
+    default:
+        snprintf(buf, sizeof(buf), "%u times", repeat);
+        return buf;
+    }
+}
+
 int schedule_task(task_t *task, cpu_t *cpu) {
     ASSERT(task);
 
@@ -165,7 +181,8 @@ int schedule_task(task_t *task, cpu_t *cpu) {
 
     BUG_ON(get_task_state(task) != TASK_STATE_READY);
 
-    printk("CPU[%u]: Scheduling task %s[%u]\n", cpu->id, task->name, task->id);
+    printk("CPU[%u]: Scheduling task %s[%u] (%s)\n", cpu->id, task->name, task->id,
+           task_repeat_string(task->repeat));
 
     spin_lock(&cpu->lock);
     list_add_tail(&task->list, &cpu->task_queue);
@@ -182,7 +199,8 @@ static void run_task(task_t *task) {
 
     wait_for_task_state(task, TASK_STATE_SCHEDULED);
 
-    printk("CPU[%u]: Running task %s[%u]\n", task->cpu->id, task->name, task->id);
+    if (atomic64_inc_return(&task->exec_count) == 0)
+        printk("CPU[%u]: Running task %s[%u]\n", task->cpu->id, task->name, task->id);
 
     set_task_state(task, TASK_STATE_RUNNING);
     if (task->type == TASK_TYPE_USER)
@@ -214,6 +232,24 @@ void wait_for_task_group(const cpu_t *cpu, task_group_t group) {
     } while (busy);
 }
 
+void process_task_repeat(task_t *task) {
+    switch (task->repeat) {
+    case TASK_REPEAT_ONCE:
+        printk("%s task '%s' finished on CPU[%u] with result %ld (Run: %lu times)\n",
+               task->type == TASK_TYPE_KERNEL ? "Kernel" : "User", task->name,
+               task->cpu->id, task->result, atomic_read(&task->exec_count));
+        destroy_task(task);
+        break;
+    case TASK_REPEAT_LOOP:
+        set_task_state(task, TASK_STATE_SCHEDULED);
+        break;
+    default:
+        task->repeat--;
+        set_task_state(task, TASK_STATE_SCHEDULED);
+        break;
+    }
+}
+
 void run_tasks(cpu_t *cpu) {
     task_t *task, *safe;
 
@@ -225,10 +261,7 @@ void run_tasks(cpu_t *cpu) {
         list_for_each_entry_safe (task, safe, &cpu->task_queue, list) {
             switch (task->state) {
             case TASK_STATE_DONE:
-                printk("%s task '%s' finished on CPU[%u] with result %ld\n",
-                       task->type == TASK_TYPE_KERNEL ? "Kernel" : "User", task->name,
-                       cpu->id, task->result);
-                destroy_task(task);
+                process_task_repeat(task);
                 break;
             case TASK_STATE_SCHEDULED:
                 run_task(task);
diff --git a/include/sched.h b/include/sched.h
index 2e4af358..86438ead 100644
--- a/include/sched.h
+++ b/include/sched.h
@@ -59,6 +59,11 @@ typedef enum task_type task_type_t;
 
 typedef unsigned int tid_t;
 
+typedef enum task_repeat {
+    TASK_REPEAT_LOOP = 0,
+    TASK_REPEAT_ONCE = 1,
+} task_repeat_t;
+
 struct task {
     list_head_t list;
 
@@ -66,6 +71,8 @@ struct task {
     task_type_t type;
     task_group_t gid;
     task_state_t state;
+    task_repeat_t repeat;
+    atomic64_t exec_count;
 
     cpu_t *cpu;
     void *stack;
@@ -110,4 +117,18 @@ static inline void execute_tasks(void) {
     run_tasks(get_bsp_cpu());
     wait_for_all_cpus();
 }
+
+static inline void set_task_repeat(task_t *task, task_repeat_t value) {
+    ASSERT(task);
+    task->repeat = value;
+}
+
+static inline void set_task_loop(task_t *task) {
+    set_task_repeat(task, TASK_REPEAT_LOOP);
+}
+
+static inline void set_task_once(task_t *task) {
+    set_task_repeat(task, TASK_REPEAT_ONCE);
+}
+
 #endif /* KTF_SCHED_H */
diff --git a/tests/unittests.c b/tests/unittests.c
index fd8ab052..0b40974f 100644
--- a/tests/unittests.c
+++ b/tests/unittests.c
@@ -169,6 +169,7 @@ int unit_tests(void *_unused) {
     task_user1 = new_user_task("test1 user", test_user_task_func1, NULL);
     task_user2 = new_user_task("test2 user", test_user_task_func2, NULL);
 
+    set_task_repeat(task1, 10);
     schedule_task(task1, get_bsp_cpu());
     schedule_task(task2, get_cpu(1));
     schedule_task(task_user1, get_bsp_cpu());