From e2131b08ac757fc901553bc05e5b81b6fab32f1b Mon Sep 17 00:00:00 2001
From: Max Wash <mbswash@gmail.com>
Date: Fri, 28 Apr 2023 21:05:48 +0100
Subject: [PATCH] sched: track CPU clock cycles used by threads

---
 include/socks/sched.h |  8 ++++++++
 sched/core.c          | 44 +++++++++++++++++++++++++++++++++++++++++++
 sched/task.c          |  1 +
 3 files changed, 53 insertions(+)

diff --git a/include/socks/sched.h b/include/socks/sched.h
index 6f9ce0f..d0db98b 100644
--- a/include/socks/sched.h
+++ b/include/socks/sched.h
@@ -61,6 +61,10 @@ struct thread {
 	unsigned int tr_id;
 	unsigned int tr_prio;
 
+	cycles_t tr_charge_period_start;
+	cycles_t tr_quantum_cycles, tr_quantum_target;
+	cycles_t tr_total_cycles;
+
 	struct queue_entry tr_threads;
 	struct queue_entry tr_rqentry;
 	void *tr_kstack;
@@ -84,11 +88,15 @@ static inline struct task *task_ref(struct task *task) { return (struct task *)o
 static inline void task_deref(struct task *task) { object_deref(object_header(task)); }
 extern struct task *task_from_pid(unsigned int pid);
 extern struct task *kernel_task(void);
+extern cycles_t default_quantum(void);
 
 extern bool need_resched(void);
 extern struct task *current_task(void);
 extern struct thread *current_thread(void);
 
+extern void start_charge_period(void);
+extern void end_charge_period(void);
+
 static inline void task_lock_irqsave(struct task *task, unsigned long *flags)
 {
 	object_lock(object_header(task), flags);
diff --git a/sched/core.c b/sched/core.c
index ccedb98..cc70eb4 100644
--- a/sched/core.c
+++ b/sched/core.c
@@ -7,6 +7,8 @@ extern kern_status_t setup_kernel_task(void);
 extern kern_status_t task_object_type_init(void);
 extern kern_status_t thread_object_type_init(void);
 
+static cycles_t __default_quantum = 0;
+
 kern_status_t sched_init(void)
 {
 	kern_status_t status = KERN_OK;
@@ -33,5 +35,47 @@ kern_status_t sched_init(void)
 	this_cpu->c_current_thread = this_thread;
 	put_cpu(this_cpu);
 
+	start_charge_period();
+
 	return status;
 }
+
+void start_charge_period(void)
+{
+	struct thread *self = current_thread();
+	if (!self) {
+		return;
+	}
+
+	self->tr_charge_period_start = get_cycles();
+}
+
+void end_charge_period(void)
+{
+	preempt_disable();
+	struct thread *self = current_thread();
+	if (!self) {
+		return;
+	}
+
+	cycles_t end = get_cycles();
+	preempt_enable();
+
+	cycles_t charge = cycles_diff(self->tr_charge_period_start, end);
+
+	self->tr_quantum_cycles += charge;
+	self->tr_total_cycles += charge;
+
+	if (self->tr_quantum_cycles >= self->tr_quantum_target) {
+		self->tr_flags |= THREAD_F_NEED_RESCHED;
+	}
+
+	self->tr_charge_period_start = 0;
+
+	//printk("%llu cycles charged to %s/%u", charge, self->tr_parent->t_name, self->tr_parent->t_id);
+}
+
+cycles_t default_quantum(void)
+{
+	return __default_quantum;
+}
diff --git a/sched/task.c b/sched/task.c
index a7764a9..4354a45 100644
--- a/sched/task.c
+++ b/sched/task.c
@@ -40,6 +40,7 @@ kern_status_t setup_kernel_task(void)
 	kernel_thread->tr_prio = PRIO_NORMAL;
 	kernel_thread->tr_state = THREAD_READY;
 	kernel_thread->tr_parent = __kernel_task;
+	kernel_thread->tr_quantum_target = default_quantum();
 
 	unsigned long flags;
 	task_lock_irqsave(__kernel_task, &flags);