beautification

1 files changed, 11 insertions, 0 deletions

diff --git a/arch/riscv64/kernel/cpu.c b/arch/riscv64/kernel/cpu.c
index 4273d02..a730b87 100644
--- a/arch/riscv64/kernel/cpu.c
+++ b/arch/riscv64/kernel/cpu.c
@@ -7,10 +7,21 @@
  */
 
 #include <apos/tcb.h>
+#include <apos/atomic.h>
+
 #include <arch/cpu.h>
 
+static atomic_long cpus = 0;
+
 void cpu_assign(struct tcb *t)
 {
+	/* bounce cpu id forward, or if first time here, get a cpu id */
+	struct tcb *c = cur_tcb();
+	if (likely(c))
+		t->cpu_id = c->cpu_id;
+	else
+		t->cpu_id = cpus++;
+
 	__asm__ volatile ("mv tp, %0\n" : : "r" (t) :);
 }