From bb524c1576c60439bb4d0b2aea222b75d44e4b06 Mon Sep 17 00:00:00 2001
From: Max Wash <mbswash@gmail.com>
Date: Thu, 4 May 2023 21:42:51 +0100
Subject: [PATCH] kernel: add support for getting percpu variables that belong
 to other CPUs

---
 include/socks/cpu.h    |  4 ++++
 include/socks/percpu.h | 10 ++++++++--
 include/socks/sched.h  |  1 +
 kernel/cpu.c           | 23 +++++++++++++++++++++++
 kernel/percpu.c        |  4 ++--
 sched/runqueue.c       | 11 +++++++++++
 6 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/include/socks/cpu.h b/include/socks/cpu.h
index 4b25595..d08199e 100644
--- a/include/socks/cpu.h
+++ b/include/socks/cpu.h
@@ -30,8 +30,12 @@ struct cpu_data {
 #define this_cpu() (ml_cpu_block_get_id(ml_this_cpu()))
 
 extern struct cpu_data *get_this_cpu(void);
+extern struct cpu_data *get_cpu(unsigned int id);
 extern void put_cpu(struct cpu_data *cpu);
 
+extern bool cpu_is_available(unsigned int cpu_id);
+extern bool cpu_is_online(unsigned int cpu_id);
+
 extern void cpu_set_available(unsigned int cpu_id);
 extern void cpu_set_online(unsigned int cpu_id);
 
diff --git a/include/socks/percpu.h b/include/socks/percpu.h
index b97b3ec..fb44a5b 100644
--- a/include/socks/percpu.h
+++ b/include/socks/percpu.h
@@ -15,13 +15,19 @@ extern "C" {
 #define percpu_get(var) \
 	__extension__({ \
 	preempt_disable(); \
-	__percpu_get(var); \
+	__percpu_get(this_cpu(), var); \
+	})
+
+#define percpu_get_from(cpu, var) \
+	__extension__({ \
+	preempt_disable(); \
+	__percpu_get(cpu, var); \
 	})
 
 #define percpu_put(var) preempt_enable();
 
 extern kern_status_t init_per_cpu_areas(void);
-extern void *__percpu_get(void *var);
+extern void *__percpu_get(unsigned int cpu, void *var);
 
 #ifdef __cplusplus
 }
diff --git a/include/socks/sched.h b/include/socks/sched.h
index 60d3933..6ebbfb7 100644
--- a/include/socks/sched.h
+++ b/include/socks/sched.h
@@ -127,6 +127,7 @@ static inline void rq_unlock(struct runqueue *rq, unsigned long flags)
 {
 	spin_unlock_irqrestore(&rq->rq_lock, flags);
 }
+extern struct runqueue *cpu_rq(unsigned int cpu);
 
 extern struct task *task_alloc(void);
 static inline struct task *task_ref(struct task *task) { return (struct task *)object_data(object_ref(object_header(task))); }
diff --git a/kernel/cpu.c b/kernel/cpu.c
index 001432a..987c8ba 100644
--- a/kernel/cpu.c
+++ b/kernel/cpu.c
@@ -12,11 +12,34 @@ struct cpu_data *get_this_cpu(void)
 	return percpu_get(&cpu_data);
 }
 
+struct cpu_data *get_cpu(unsigned int id)
+{
+	return percpu_get_from(id, &cpu_data);
+}
+
 void put_cpu(struct cpu_data *cpu)
 {
 	percpu_put(cpu);
 }
 
+bool cpu_is_available(unsigned int cpu_id)
+{
+	if (cpu_id >= CPU_MAX) {
+		return false;
+	}
+
+	return bitmap_check(cpu_available, cpu_id);
+}
+
+bool cpu_is_online(unsigned int cpu_id)
+{
+	if (cpu_id >= CPU_MAX) {
+		return false;
+	}
+
+	return bitmap_check(cpu_online, cpu_id);
+}
+
 void cpu_set_available(unsigned int cpu_id)
 {
 	if (cpu_id >= CPU_MAX) {
diff --git a/kernel/percpu.c b/kernel/percpu.c
index 3b9e622..2ec9ac4 100644
--- a/kernel/percpu.c
+++ b/kernel/percpu.c
@@ -24,7 +24,7 @@ extern kern_status_t init_per_cpu_areas(void)
 	return KERN_OK;
 }
 
-extern void *__percpu_get(void *var)
+extern void *__percpu_get(unsigned int cpu, void *var)
 {
 	uintptr_t pvar = (uintptr_t)var;
 	uintptr_t percpu_start = (uintptr_t)__percpu_start;
@@ -36,5 +36,5 @@ extern void *__percpu_get(void *var)
 
 	size_t var_offset = pvar - percpu_start;
 
-	return (char *)percpu_buffer + (this_cpu() * percpu_stride) + var_offset;
+	return (char *)percpu_buffer + (cpu * percpu_stride) + var_offset;
 }
diff --git a/sched/runqueue.c b/sched/runqueue.c
index d5bd097..800cc39 100644
--- a/sched/runqueue.c
+++ b/sched/runqueue.c
@@ -1,5 +1,6 @@
 #include <socks/sched.h>
 #include <socks/percpu.h>
+#include <socks/cpu.h>
 
 #define PRIO_MASK(p) (((uint32_t)1) << (p))
 #define FIRST_PRIO(m) (m > 0 ? (PRIO_MAX - __builtin_clz(m) - 1) : -1)
@@ -42,6 +43,16 @@ void rq_enqueue(struct runqueue *rq, struct thread *thr)
 
 	struct queue *q = &rq->rq_queues[prio];
 	queue_push_back(q, &thr->tr_rqentry);
+	rq->rq_nthreads++;
 
 	rq->rq_readybits |= PRIO_MASK(thread_priority(thr));
 }
+
+struct runqueue *cpu_rq(unsigned int cpu)
+{
+	struct cpu_data *cpu_data = get_cpu(cpu);
+	struct runqueue *rq = &cpu_data->c_rq;
+	put_cpu(cpu_data);
+
+	return rq;
+}