lib: Add epoch-based SMR

2026-11-01 21:06:48 +01:00
parent 3ff068e944
commit 7c04530681
24 changed files with 1059 additions and 84 deletions
@@ -0,0 +1,252 @@
+// SPDX-License-Identifier: ISC
+
+#include <zinnia/assert.h>
+#include <zinnia/atomic.h>
+#include <zinnia/clock.h>
+#include <zinnia/epoch.h>
+#include <zinnia/irq.h>
+#include <zinnia/mutex.h>
+#include <zinnia/percpu.h>
+#include <zinnia/process.h>
+#include <zinnia/sched.h>
+#include <zinnia/spin.h>
+#include <zinnia/timer.h>
+#include <arch/cpu.h>
+
+constexpr uint64_t EPOCH_NEST_BITS = 16;
+constexpr uint64_t EPOCH_NEST_MASK = ((uint64_t)1 << EPOCH_NEST_BITS) - 1;
+constexpr size_t EPOCH_BUCKETS = 3;
+constexpr size_t EPOCH_BATCH = 64;
+constexpr size_t EPOCH_SYNC_SPINS = 64;
+constexpr uint64_t EPOCH_POLL_NS = 10'000'000;
+constexpr uint64_t EPOCH_SYNC_POLL_NS = 1'000'000;
+
+struct epoch_list {
+    struct epoch_head* head;
+    struct epoch_head* tail;
+};
+
+struct epoch_bucket {
+    uint64_t epoch;
+    struct epoch_list list;
+};
+
+struct epoch_cpu {
+    _Atomic(uint64_t) local;
+    struct spinlock lock;
+    struct epoch_bucket buckets[EPOCH_BUCKETS];
+    struct epoch_list ready;
+};
+
+[[__percpu]]
+static struct epoch_cpu epoch_data;
+
+static _Atomic(uint64_t) global_epoch = 0;
+static atomic_size_t pending = 0;
+static mutex_t reclaim_lock;
+static struct timer poll_timer;
+static _Atomic(struct task*) reclaimer = nullptr;
+static atomic_bool reclaimer_idle = false;
+
+void epoch_enter() {
+    sched_preempt_disable();
+    _Atomic(uint64_t)* local = &percpu_ref(epoch_data)->local;
+    const uint64_t value = atomic_load(local, ATOMIC_RELAXED);
+
+    // If we already entered the critical section before, just increment the nest counter.
+    if (value & EPOCH_NEST_MASK) {
+        assert((value & EPOCH_NEST_MASK) != EPOCH_NEST_MASK, "Epoch nesting overflow");
+        atomic_store(local, value + 1, ATOMIC_RELAXED);
+        return;
+    }
+
+    // Store the current epoch in the upper bits, and mark this CPU as active.
+    atomic_store(local, (atomic_load(&global_epoch, ATOMIC_RELAXED) << EPOCH_NEST_BITS) | 1, ATOMIC_RELAXED);
+    atomic_fence(ATOMIC_SEQ_CST);
+}
+
+void epoch_exit() {
+    _Atomic(uint64_t)* local = &percpu_ref(epoch_data)->local;
+    const uint64_t value = atomic_load(local, ATOMIC_RELAXED);
+    assert(value & EPOCH_NEST_MASK, "Tried to exit an epoch section without entering one");
+
+    // If we didn't nest, we can clear the local epoch.
+    // Otherwise, decrease the nest counter.
+    if ((value & EPOCH_NEST_MASK) == 1)
+        atomic_store(local, 0, ATOMIC_RELEASE);
+    else
+        atomic_store(local, value - 1, ATOMIC_RELAXED);
+
+    sched_preempt_enable();
+}
+
+bool epoch_active() {
+    sched_preempt_disable();
+    const bool active = atomic_load(&percpu_ref(epoch_data)->local, ATOMIC_RELAXED) & EPOCH_NEST_MASK;
+    sched_preempt_enable();
+    return active;
+}
+
+static void list_push(struct epoch_list* list, struct epoch_head* head) {
+    head->next = list->head;
+    if (!list->head)
+        list->tail = head;
+    list->head = head;
+}
+
+static void list_splice(struct epoch_list* dst, struct epoch_list* src) {
+    if (!src->head)
+        return;
+    src->tail->next = dst->head;
+    if (!dst->head)
+        dst->tail = src->tail;
+    dst->head = src->head;
+    *src = (struct epoch_list){};
+}
+
+static void wake_reclaimer() {
+    struct task* task = atomic_load(&reclaimer, ATOMIC_ACQUIRE);
+    if (task && atomic_exchange(&reclaimer_idle, false, ATOMIC_SEQ_CST))
+        sched_wake(task);
+}
+
+void epoch_call(struct epoch_head* head, void (*func)(struct epoch_head* head)) {
+    head->func = func;
+    const size_t old = atomic_fetch_add(&pending, 1, ATOMIC_SEQ_CST);
+    atomic_fence(ATOMIC_SEQ_CST);
+
+    irq_lock();
+    struct epoch_cpu* cpu = percpu_ref(epoch_data);
+    spin_lock(&cpu->lock);
+
+    // Store the object head in the bucket of this epoch.
+    // We only need to store 3 buckets, since only these epochs are relevant:
+    // - Current epoch, may have readers
+    // - Previous epoch, may still have readers
+    // - Safe epoch, has no readers remaining
+    const uint64_t now = atomic_load(&global_epoch, ATOMIC_RELAXED);
+    struct epoch_bucket* bucket = &cpu->buckets[now % EPOCH_BUCKETS];
+    if (bucket->epoch != now) {
+        // Combine both epoch lists.
+        list_splice(&cpu->ready, &bucket->list);
+        bucket->epoch = now;
+    }
+    list_push(&bucket->list, head);
+
+    spin_unlock(&cpu->lock);
+    irq_unlock();
+
+    if (old == 0 || old % EPOCH_BATCH == 0)
+        wake_reclaimer();
+}
+
+uint64_t epoch_current() {
+    return atomic_load(&global_epoch, ATOMIC_RELAXED);
+}
+
+bool epoch_try_advance() {
+    uint64_t now = atomic_load(&global_epoch, ATOMIC_RELAXED);
+    atomic_fence(ATOMIC_SEQ_CST);
+    const size_t count = percpu_count();
+
+    // Go through all CPUs.
+    for (size_t i = 0; i < count; i++) {
+        const uint64_t local = atomic_load(&percpu_ref_on(percpu_cpu(i), epoch_data)->local, ATOMIC_ACQUIRE);
+
+        // If any of the CPUs are in a critical section, and their local epoch
+        // is not the same as the current global epoch, we can't advance.
+        if ((local & EPOCH_NEST_MASK) && (local >> EPOCH_NEST_BITS) != now)
+            return false;
+    }
+
+    // At this point all readers agree on the epoch, it's safe to increment now.
+    atomic_compare_exchange(&global_epoch, &now, now + 1, ATOMIC_SEQ_CST, ATOMIC_RELAXED);
+    return true;
+}
+
+static struct epoch_list collect(struct epoch_cpu* cpu, uint64_t now) {
+    spin_lock(&cpu->lock);
+    for (size_t i = 0; i < EPOCH_BUCKETS; i++) {
+        struct epoch_bucket* bucket = &cpu->buckets[i];
+        if (bucket->epoch + 2 <= now)
+            list_splice(&cpu->ready, &bucket->list);
+    }
+    struct epoch_list out = cpu->ready;
+    cpu->ready = (struct epoch_list){};
+    spin_unlock(&cpu->lock);
+    return out;
+}
+
+static void run(struct epoch_head* head) {
+    size_t count = 0;
+    while (head) {
+        struct epoch_head* next = head->next;
+        head->func(head);
+        head = next;
+        count++;
+    }
+    if (count)
+        atomic_fetch_sub(&pending, count, ATOMIC_RELAXED);
+}
+
+static void reclaim() {
+    mutex_lock(&reclaim_lock);
+    const uint64_t now = atomic_load(&global_epoch, ATOMIC_ACQUIRE);
+    const size_t count = percpu_count();
+    for (size_t i = 0; i < count; i++)
+        run(collect(percpu_ref_on(percpu_cpu(i), epoch_data), now).head);
+    mutex_unlock(&reclaim_lock);
+}
+
+void epoch_synchronize() {
+    // We can't synchronize a write while being a reader ourselves.
+    assert(!epoch_active(), "Tried to synchronize inside an epoch section");
+
+    atomic_fence(ATOMIC_SEQ_CST);
+    const uint64_t target = atomic_load(&global_epoch, ATOMIC_RELAXED) + 2;
+    for (size_t spins = 0; atomic_load(&global_epoch, ATOMIC_ACQUIRE) < target; spins++) {
+        if (epoch_try_advance())
+            continue;
+        if (spins < EPOCH_SYNC_SPINS)
+            arch_cpu_pause();
+        else
+            sched_sleep_until(clock_get_elapsed_ns() + EPOCH_SYNC_POLL_NS);
+    }
+}
+
+void epoch_barrier() {
+    epoch_synchronize();
+    reclaim();
+}
+
+static void poll_tick(struct timer*) {
+    wake_reclaimer();
+}
+
+static void reclaimer_fn(uintptr_t) {
+    while (true) {
+        if (atomic_load(&pending, ATOMIC_RELAXED)) {
+            epoch_try_advance();
+            reclaim();
+        }
+
+        atomic_store(&reclaimer_idle, true, ATOMIC_SEQ_CST);
+        if (atomic_load(&pending, ATOMIC_SEQ_CST))
+            timer_arm(&poll_timer, clock_get_elapsed_ns() + EPOCH_POLL_NS);
+        sched_block(&percpu_get()->sched);
+    }
+}
+
+static void epoch_stage_init(struct boot_info*) {
+    mutex_init(&reclaim_lock);
+    timer_init(&poll_timer, poll_tick);
+
+    struct task* task;
+    errno_t status = task_create("epoch", process_kernel(), reclaimer_fn, 0, &task);
+    assert(status == 0, "Unable to create the epoch reclaimer");
+    sched_set_priority(task, SCHED_PRIO_WORK_HIGH);
+    atomic_store(&reclaimer, task, ATOMIC_RELEASE);
+    sched_add_task(&percpu_get()->sched, task);
+}
+INIT_TASK(epoch_stage, "core.epoch", epoch_stage_init);
+INIT_DEPENDS(epoch_stage, sched_stage);