From cec6b644ac3f81ea3f96f7f5c1ba10f3237285de Mon Sep 17 00:00:00 2001
From: Max Wash <mbswash@gmail.com>
Date: Fri, 5 May 2023 15:25:44 +0100
Subject: [PATCH] kernel: add functions to query the number of available/online
 CPUs

---
 include/socks/cpu.h |  3 +++
 kernel/cpu.c        | 10 ++++++++++
 2 files changed, 13 insertions(+)

diff --git a/include/socks/cpu.h b/include/socks/cpu.h
index d08199e..1e3c47b 100644
--- a/include/socks/cpu.h
+++ b/include/socks/cpu.h
@@ -39,6 +39,9 @@ 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);
 
+extern unsigned int cpu_nr_available(void);
+extern unsigned int cpu_nr_online(void);
+
 extern cycles_t get_cycles(void);
 static inline cycles_t cycles_diff(cycles_t then, cycles_t now)
 {
diff --git a/kernel/cpu.c b/kernel/cpu.c
index 987c8ba..ecbc4a9 100644
--- a/kernel/cpu.c
+++ b/kernel/cpu.c
@@ -58,6 +58,16 @@ void cpu_set_online(unsigned int cpu_id)
 	bitmap_set(cpu_online, cpu_id);
 }
 
+unsigned int cpu_nr_available(void)
+{
+	return bitmap_count_set(cpu_available, CPU_MAX);
+}
+
+unsigned int cpu_nr_online(void)
+{
+	return bitmap_count_set(cpu_online, CPU_MAX);
+}
+
 void preempt_disable(void)
 {
 	ml_cpu_block *ml_cpu = ml_this_cpu();