From 6776f44e994b5a0cdc3831a0a1de57fe8b4529d9 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  | 23 ++++++++++++++++++++++-
 include/sched.h | 15 +++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/common/sched.c b/common/sched.c
index 1da14b5e..376489cd 100644
--- a/common/sched.c
+++ b/common/sched.c
@@ -87,6 +87,7 @@ static task_t *create_task(void) {
     task->id = next_tid++;
     task->gid = TASK_GROUP_ALL;
     set_task_state(task, TASK_STATE_NEW);
+    set_task_once(task);
 
     return task;
 }
@@ -176,6 +177,11 @@ int schedule_task(task_t *task, cpu_t *cpu) {
     return 0;
 }
 
+void set_task_repeat(task_t *task, task_repeat_t value) {
+    ASSERT(task);
+    task->repeat = value;
+}
+
 static void run_task(task_t *task) {
     if (!task)
         return;
@@ -214,6 +220,21 @@ 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:
+        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;
 
@@ -226,7 +247,7 @@ void run_tasks(cpu_t *cpu) {
                 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 d40af72e..b645f547 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,7 @@ struct task {
     task_type_t type;
     task_group_t gid;
     task_state_t state;
+    task_repeat_t repeat;
 
     cpu_t *cpu;
     void *stack;
@@ -84,6 +90,7 @@ extern void init_tasks(void);
 extern task_t *get_task_by_name(cpu_t *cpu, const char *name);
 extern task_t *new_task(const char *name, task_func_t func, void *arg, task_type_t type);
 extern int schedule_task(task_t *task, cpu_t *cpu);
+extern void set_task_repeat(task_t *task, task_repeat_t value);
 extern void run_tasks(cpu_t *cpu);
 extern void wait_for_task_group(const cpu_t *cpu, task_group_t group);
 
@@ -105,4 +112,12 @@ static inline task_t *new_user_task(const char *name, task_func_t func, void *ar
     return new_task(name, func, arg, TASK_TYPE_USER);
 }
 
+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 */