feat: hartlocks

Signed-off-by: Autumn <auctumnus@pm.me>

Author: auctumnus

src/kernel/arch/riscv64/hart_locals.c

@@ -27,7 +27,7 @@ static u64 boothart_hart_id;
  */
 void heap_change_boothart_hart(struct mm_alloc_heap *heap, struct hart *hart);
 
-struct hart_locals *get_hart_locals(void) {
+struct hart_locals *get_hart_locals(struct hartlock *_hartlock) {
   return (struct hart_locals *)csrr(RISCV64_CSR_SSCRATCH);
 }
 

src/kernel/hartlock.c

@@ -0,0 +1,35 @@
+/*
+ * SPDX-FileCopyrightText: 2026 ukoOS Contributors
+ *
+ * SPDX-License-Identifier: GPL-3.0-or-later
+ */
+
+#include <stdatomic.h>
+
+#include <hart_locals.h>
+#include <panic.h>
+#include <hartlock.h>
+#include <types.h>
+// TODO: make generic across architectures
+#include <arch/riscv64/irq.h>
+
+// None of these functions should be called manually.
+// See doc/kernel/containers/guards.md for usage.
+
+void __hartlock_init(struct hartlock *lock) {
+  if(lock == nullptr) {
+    panic("hartlock init called with null lock");
+  }
+
+  u64 flags = local_irq_save();
+
+  lock->flags = flags;
+}
+
+void __hartlock_free(struct hartlock *lock) {
+  if (lock == nullptr) {
+    panic("hartlock free called with null lock");
+  }
+
+  local_irq_restore(lock->flags);
+}

src/kernel/include.mak

@@ -35,6 +35,7 @@ kernel-objs-c += device
 kernel-objs-c += devices/hart
 kernel-objs-c += devices/uart
 kernel-objs-c += devicetree
+kernel-objs-c += hartlock
 kernel-objs-c += init
 kernel-objs-c += main
 kernel-objs-c += mm/alloc

src/kernel/include/arch/riscv64/irq.h

@@ -11,7 +11,7 @@
  * IRQ handling.
  */
 
-#include <insns.h>
+#include <arch/riscv64/insns.h>
 #include <types.h>
 
 // TODO: move this over to insns.h?

src/kernel/include/hart_locals.h

@@ -9,6 +9,7 @@
 
 #include <crypto/subtle/random_internals.h>
 #include <devices/hart.h>
+#include <hartlock.h>
 
 /**
  * The data that is local to a hart.
@@ -41,10 +42,22 @@ static_assert(offsetof(struct hart_locals, task) == 8);
 /**
  * Returns a pointer to the current hart's locals.
  *
- * TODO: Make sure this is called under some lock that ensures the current
- * thread cannot be rescheduled. Otherwise, it'd be super-easy to get a race.
+ * Requires a hartlock to be held. You can acquire a hartlock by using
+ * `WITH_HARTLOCK` in a new scope; for example:
+ *
+ * ```c
+ * void foo() {
+ *   u8 blah = 2 + 2;
+ *   {
+ *     WITH_HARTLOCK(lock);
+ *     struct hart_locals *locals = get_hart_locals(lock);
+ *     
+ *     // access `locals->rng`, etc ...
+ *   }
+ * }
+ * ```
  */
-struct hart_locals *get_hart_locals(void);
+struct hart_locals *get_hart_locals(struct hartlock *);
 
 /**
  * Called during early boot to initialize the boothart's hart-local storage

src/kernel/include/hartlock.h

@@ -0,0 +1,39 @@
+/*
+ * SPDX-FileCopyrightText: 2026 ukoOS Contributors
+ *
+ * SPDX-License-Identifier: GPL-3.0-or-later
+ */
+
+#ifndef UKO_OS_KERNEL__HARTLOCK_H
+#define UKO_OS_KERNEL__HARTLOCK_H 1
+
+#include <types.h>
+
+struct hartlock {
+  u64 flags;
+};
+
+// None of these functions should be called manually.
+// See doc/kernel/containers/guards.md for usage.
+
+void __hartlock_init(struct hartlock *);
+
+void __hartlock_free(struct hartlock *);
+
+#define HARTLOCK_GUARDED(VARS)                                                 \
+  struct VARS __guarded_by_hartlock;                                           \
+
+#define HARTLOCK_GUARD(GUARD_NAME, OBJECT)                                     \
+  typeof((OBJECT)->__guarded_by_hartlock) *GUARD_NAME;                         \
+  [[gnu::cleanup(__hartlock_free)]]                                            \
+  struct hartlock __my_hartlock = { 0 };                                       \
+  __hartlock_init(&__my_hartlock);                                             \
+  GUARD_NAME = &((OBJECT)->__guarded_by_hartlock)
+
+#define WITH_HARTLOCK(LOCK_NAME)                                               \
+  [[gnu::cleanup(__hartlock_free)]]                                            \
+  struct hartlock __my_hartlock = { 0 };                                       \
+  __hartlock_init(&__my_hartlock);                                             \
+  struct hartlock *LOCK_NAME = &__my_hartlock
+
+#endif // UKO_OS_KERNEL__HARTLOCK_H

src/kernel/include/mm/alloc.h

@@ -8,6 +8,7 @@
 #define UKO_OS_KERNEL__MM_ALLOC_H 1
 
 #include <hart_locals.h>
+#include <hartlock.h>
 
 /**
  * Frees memory returned by a previous call to `alloc`.
@@ -21,13 +22,13 @@ void free(void *ptr);
  * Like `alloc`, but only for allocations with `0 < size && size <= 1024`.
  */
 [[gnu::alloc_size(1), gnu::malloc, gnu::malloc(free, 1), nodiscard]] void *
-alloc_small(usize size, struct mm_alloc_heap *heap);
+alloc_small(usize size, struct hartlock *hartlock, struct mm_alloc_heap *heap);
 
 /**
  * The slow path of `alloc`.
  */
 [[gnu::alloc_size(1), gnu::malloc, gnu::malloc(free, 1), nodiscard]] void *
-alloc_generic(usize size, struct mm_alloc_heap *heap);
+alloc_generic(usize size, struct hartlock *hartlock, struct mm_alloc_heap *heap);
 
 /**
  * Allocates `size` bytes of memory and returns a pointer to it. On OOM, returns
@@ -40,7 +41,8 @@ alloc_generic(usize size, struct mm_alloc_heap *heap);
  */
 [[gnu::alloc_size(1), gnu::malloc, nodiscard]] static inline void *
 alloc(usize size) {
-  struct mm_alloc_heap *heap = get_hart_locals()->heap;
+  WITH_HARTLOCK(hartlock);
+  struct mm_alloc_heap *heap = get_hart_locals(hartlock)->heap;
   // Bump up the size if it's zero, since zero isn't a valid input to
   // alloc_small.
   //
@@ -50,11 +52,11 @@ alloc(usize size) {
   // constant (or at least provably non-zero) in the caller most of the time, so
   // the branch should be eliminated.
   if (size == 0)
-    return alloc_small(1, heap);
+    return alloc_small(1, hartlock, heap);
   else if (size <= 1024)
-    return alloc_small(size, heap);
+    return alloc_small(size, hartlock, heap);
   else
-    return alloc_generic(size, heap);
+    return alloc_generic(size, hartlock, heap);
 }
 
 /**

src/kernel/include/spinlock.h

@@ -8,10 +8,11 @@
 #define UKO_OS_KERNEL__SPINLOCK_H 1
 
 #include <types.h>
+#include <hartlock.h>
 
 struct spinlock {
   atomic bool locked;
-  u64 flags;
+  struct hartlock hartlock;
 
   // these only have meaning while `locked` is true
   atomic u64 owner_hart_id;

src/kernel/main.c

@@ -16,6 +16,7 @@
 #include <random.h>
 #include <selftest.h>
 #include <symbolicate.h>
+#include <hartlock.h>
 
 [[noreturn]]
 void main(u64 hart_id, paddr devicetree_start, paddr kernel_start,
@@ -64,7 +65,10 @@ void main(u64 hart_id, paddr devicetree_start, paddr kernel_start,
   print("Running self-tests...");
   run_selftests();
 
-  struct hart_locals *hart_locals = get_hart_locals();
-  print("{uptr}", hart_locals->hart);
+  {
+    WITH_HARTLOCK(hartlock);
+    struct hart_locals *hart_locals = get_hart_locals(hartlock);
+    print("{uptr}", hart_locals->hart);
+  }
   TODO();
 }

src/kernel/mm/alloc.c

@@ -9,6 +9,7 @@
 #include <mm/alloc.h>
 #include <panic.h>
 #include <stdatomic.h>
+#include "hartlock.h"
 
 void free(void *ptr) {
   if (!ptr)
@@ -17,9 +18,12 @@ void free(void *ptr) {
   struct mm_alloc_segment *segment = segment_of_ptr((uptr)ptr);
   struct mm_alloc_page *page = page_of_ptr((uptr)ptr);
   struct mm_alloc_block *block = ptr;
-  if (segment->hart == get_hart_locals()->hart) {
+  
+  WITH_HARTLOCK(hartlock);
+
+  if (segment->hart == get_hart_locals(hartlock)->hart) {
     // This is a local free; i.e., we're running on the hart that owns the page.
-    struct mm_alloc_heap *heap = get_hart_locals()->heap;
+    struct mm_alloc_heap *heap = get_hart_locals(hartlock)->heap;
 
     page_local_free_push(page, block);
     if (page_is_empty(page)) {
@@ -42,8 +46,8 @@ void free(void *ptr) {
   }
 }
 
-void *alloc_small(usize size, struct mm_alloc_heap *heap) {
-  assert(heap->hart == get_hart_locals()->hart);
+void *alloc_small(usize size, struct hartlock *hartlock, struct mm_alloc_heap *heap) {
+  assert(heap->hart == get_hart_locals(hartlock)->hart);
   assert(0 < size && size <= 1024);
 
   // Compute the index into pages_direct.
@@ -57,7 +61,7 @@ void *alloc_small(usize size, struct mm_alloc_heap *heap) {
   // generic routine.
   struct mm_alloc_block *block = page_free_pop(page);
   if (!block)
-    return alloc_generic(size, heap);
+    return alloc_generic(size, hartlock, heap);
 
   // Increment the counter of used objects.
   page->used_blocks++;
@@ -82,19 +86,19 @@ static void *alloc_generic_from_page(struct mm_alloc_page *page,
   return block;
 }
 
-static void *alloc_huge(usize size, struct mm_alloc_heap *heap) {
+static void *alloc_huge(usize size, struct hartlock *hartlock, struct mm_alloc_heap *heap) {
   assert(size_is_huge(size));
 
   // Every huge object goes in its own page, so allocate one and use it
   // directly.
-  struct mm_alloc_page *page = page_new_huge(heap, size);
+  struct mm_alloc_page *page = page_new_huge(hartlock, heap, size);
   if (!page)
     return nullptr;
   return alloc_generic_from_page(page, heap);
 }
 
-void *alloc_generic(usize size, struct mm_alloc_heap *heap) {
-  assert(heap->hart == get_hart_locals()->hart);
+void *alloc_generic(usize size, struct hartlock *hartlock, struct mm_alloc_heap *heap) {
+  assert(heap->hart == get_hart_locals(hartlock)->hart);
   assert(size);
 
   // Go through the delayed free list to free everything.
@@ -107,7 +111,7 @@ void *alloc_generic(usize size, struct mm_alloc_heap *heap) {
 
   // Huge objects get handled separately.
   if (size_is_huge(size))
-    return alloc_huge(size, heap);
+    return alloc_huge(size, hartlock, heap);
 
   // Check every page of the size class for free objects.
   usize size_class = size_class_of_size(size);
@@ -147,9 +151,9 @@ void *alloc_generic(usize size, struct mm_alloc_heap *heap) {
   // new one.
   assert(list_is_empty(&heap->pages[size_class]));
   if (size_is_small(size))
-    page = page_new_small(heap, size_class);
+    page = page_new_small(hartlock, heap, size_class);
   else
-    page = page_new_large(heap, size_class);
+    page = page_new_large(hartlock, heap, size_class);
 
   // If we couldn't allocate a new page, we're out of memory.
   if (!page)