Run a hart lottery on SMP

On systems with SMP multiple harts will try to run the kernel, and they
will appear at random. But in this kernel, in order to keep things
simple, we want to make sure that *only one* hart is running the show,
as it greatly simplifies things on these kinds of systems.

The solution is similar to what Linux does, which is to allow the first
hart to initialize things, but then (and different to what Linux does),
it will infinitely stall all the other harts that arrive at a random
later point in time.

In order to make this more apparent, I have also added a print message
showing which hart is being used to run the whole thing.

Signed-off-by: Miquel Sabaté Solà <mikisabate@gmail.com>

8 files changed, 89 insertions, 3 deletions

diff --git a/Makefile b/Makefile
index bd53a84..77e20f6 100644
--- a/Makefile
+++ b/Makefile
@@ -148,9 +148,11 @@ archive: all
 ##
 # Hacking
 
+CPUS ?= 4
+
 .PHONY: qemu
 qemu: clean $(IMAGE) usr
-	$(Q) $(QEMU) $(QEMU_FLAGS) -machine virt -kernel $(IMAGE) -initrd $(INIT)
+	$(Q) $(QEMU) $(QEMU_FLAGS) -machine virt -smp $(CPUS) -kernel $(IMAGE) -initrd $(INIT)
 
 .PHONY: gdb
 gdb:
diff --git a/include/fbos/compiler.h b/include/fbos/compiler.h
index 67f3194..77413e9 100644
--- a/include/fbos/compiler.h
+++ b/include/fbos/compiler.h
@@ -39,6 +39,11 @@ typedef unsigned long size_t;
 typedef unsigned long uint64_t;
 typedef unsigned long uintptr_t;
 
+// Strong type for atomic integer operations.
+typedef struct {
+	int32_t value;
+} atomic32_t;
+
 /*
  * NULL
  */
diff --git a/include/fbos/init.h b/include/fbos/init.h
index 9e0049e..dcffd02 100644
--- a/include/fbos/init.h
+++ b/include/fbos/init.h
@@ -5,6 +5,31 @@
 #include <fbos/dt.h>
 #include <fbos/sched.h>
 
+// Atomic value that holds how many harts have gone through the "hart lottery".
+// This is in the same spirit as it happens on the Linux kernel: the RISC-V
+// specification leaves open which hart will appear first into the kernel code.
+// This greatly simplifies the specification and the hardware, but for the
+// kernel this means that harts will appear randomly. In order to know which
+// hart runs first, in Linux they run a "lottery": an atomic value holds how
+// many harts hav already been seen. The first hart to appear will actually
+// initialize things before bringing the others up, while the others will simply
+// wait until the first hart frees the lock for them.
+//
+// That being said, here we only want *one* hart running. Hence, whichever hart
+// wins the lottery, it's not only going to initialize the kernel, but it will
+// also be the only one running the show. This is of course a waste of
+// resources, but it's not like running fizz/buzz needs SMP and cores running at
+// full speed. Actually, keeping things under a single hart simplifies things a
+// lot.
+//
+// Instantiated in kernel/main.c
+extern atomic32_t hart_lottery;
+
+// ID of the hart that is running the show.
+//
+// Instantiated in kernel/main.c
+extern uint32_t hart_id;
+
 // Tracks the amount of seconds that have elapsed since activating timer
 // interrupts.
 //
diff --git a/include/fbos/printk.h b/include/fbos/printk.h
index e68d1eb..d6ca71e 100644
--- a/include/fbos/printk.h
+++ b/include/fbos/printk.h
@@ -9,9 +9,19 @@
  */
 
 #ifdef __KERNEL__
+// Print the given message and loop indefinitely.
 extern void die(const char *const message);
+
+// Print the given number as a single digit.
+extern void print_digit(uint32_t digit);
+
+// Print the given message.
 extern void printk(const char *const message);
+
+// Print the given message which is exactly 'n' bytes long.
 extern void write(const char *const message, size_t n);
+
+// 'write' system call.
 extern void sys_write(const char *const message, size_t n);
 #else
 #include <stdio.h>
diff --git a/kernel/fbos.ld.S b/kernel/fbos.ld.S
index 160a5cb..3bda55b 100644
--- a/kernel/fbos.ld.S
+++ b/kernel/fbos.ld.S
@@ -41,13 +41,16 @@ SECTIONS {
 	.data : {
     	*(.data)
 	}
-
 	. = ALIGN(8);
 
 	.rodata : {
     	*(.rodata)
 	}
+	. = ALIGN(8);
 
+	.sdata : {
+		*(.sdata*)
+	}
 	. = ALIGN(8);
 
 	.bss : {
diff --git a/kernel/head.S b/kernel/head.S
index 74f83d8..b175b89 100644
--- a/kernel/head.S
+++ b/kernel/head.S
@@ -63,6 +63,23 @@ _start_kernel:
 	// Flush the instruction cache
 	fence.i
 
+	// Run the hart lottery. If this is not the first time that it happens, then
+	// stall this hart forever: on this simple kernel we only want one hart
+	// available to avoid SMP shenanigans. See explanation on fbos/init.h.
+	la a3, hart_lottery
+	li a2, 1
+	amoadd.w a3, a2, (a3)
+	beqz a3, .Lhart_proceed
+.Lhart_wait:
+	wfi
+	j .Lhart_wait
+
+.Lhart_proceed:
+	// Store the hart id as we will tamper with the 'a0' register later when
+	// performing the call to `start_kernel`.
+	la a3, hart_id
+	sw a0, 0(a3)
+
 	// Reset all registers except ra, a0, a1.
 	li sp, 0
 	li gp, 0
diff --git a/kernel/main.c b/kernel/main.c
index a2e4d13..822d638 100644
--- a/kernel/main.c
+++ b/kernel/main.c
@@ -5,6 +5,12 @@
 #include <fbos/string.h>
 #include <fbos/dt.h>
 
+/*
+ * Data related to harts and set by kernel/head.S. Defined in fbos/init.h.
+ */
+atomic32_t hart_lottery __section(".sdata");
+uint32_t hart_id;
+
 // Stack to be used by our processes, which is initialized in head.S.
 // "Blasphemy!" I hear you say. "How dare you use the same stack for kernel and
 // user space?" It's not like this is some sort of utopian system in which
@@ -52,6 +58,11 @@ __noreturn __kernel void start_kernel(void *dtb)
 		write("\n", 1);
 	}
 
+	// And print the hart ID where this is being run.
+	printk("Running on Hart ID: ");
+	print_digit(hart_id);
+	write("\n", 1);
+
 	// At this point everything has already been handled: setup the interrupt
 	// vector and enable the timer to start ticking and scheduling the three
 	// tasks at hand.
diff --git a/kernel/printk.c b/kernel/printk.c
index bd2bcab..1cc68db 100644
--- a/kernel/printk.c
+++ b/kernel/printk.c
@@ -3,7 +3,7 @@
 #include <fbos/string.h>
 #include <fbos/sbi.h>
 
-void __noreturn __kernel die(const char *const message)
+__kernel __noreturn void die(const char *const message)
 {
 	if (message) {
 		printk(message);
@@ -13,6 +13,19 @@ void __noreturn __kernel die(const char *const message)
 		;
 }
 
+__kernel void print_digit(uint32_t digit)
+{
+	char buffer[2];
+
+	if (digit > 9) {
+		die("We cannot print numbers with two or more digits :D\n");
+	}
+
+	buffer[0] = '0' + digit;
+	buffer[1] = '\0';
+	write(buffer, 2);
+}
+
 __kernel void write(const char *const message, size_t n)
 {
 	struct sbi_ret ret = sbi_ecall2(DBCN_EXT, DBCN_WRITE, n, (unsigned long)message);