selftests/bpf: libarena: Add spmc queue data structure

Expand libarena with a single producer multiple consumer deque data
structure. This is a single producer, multiple consumer lockless structure
that permits efficient work stealing. The structure is a Lev-Chase queue,
so it is lock-free and wait-free.

The data structure exposes three main calls. two of them are available to
the thread owning the queue and one available to all threads in the program:

spmc_owner_push(): Push an item to the top of the queue.
spmc_owner_pop(): Pop an item from the top of the queue.
spmc_steal(): Steal a thread from the bottom of the queue from
any thread.

Note that the queue is not really FIFO for all consumers, since
non-owners of the queue can only work steal from the bottom.

Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
Link: https://lore.kernel.org/r/20260605222020.5231-3-emil@etsalapatis.com
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/tools/testing/selftests/bpf/libarena/include/libarena/spmc.h b/tools/testing/selftests/bpf/libarena/include/libarena/spmc.h
new file mode 100644 (file)
index 0000000..7561127
--- /dev/null
+++ b/tools/testing/selftests/bpf/libarena/include/libarena/spmc.h
@@ -0,0 +1,27 @@
+/* SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause */
+
+#pragma once
+
+struct spmc_arr;
+
+#define SPMC_ARR_BASESZ 128
+#define SPMC_ARR_ORDERS 10
+
+struct spmc_arr {
+       u64 __arena *data;
+       u64 order;
+};
+
+struct spmc {
+       volatile struct spmc_arr __arena *cur;
+       volatile u64 top;
+       volatile u64 bottom;
+       struct spmc_arr arr[SPMC_ARR_ORDERS];
+};
+
+int spmc_owned_add(struct spmc __arena *spmc, u64 val);
+int spmc_owned_remove(struct spmc __arena *spmc, u64 *val);
+int spmc_steal(struct spmc __arena *spmc, u64 *val);
+
+struct spmc __arena *spmc_create(void);
+int spmc_destroy(struct spmc __arena *spmc);

diff --git a/tools/testing/selftests/bpf/libarena/selftests/test_spmc.bpf.c b/tools/testing/selftests/bpf/libarena/selftests/test_spmc.bpf.c
new file mode 100644 (file)
index 0000000..4d7a520
--- /dev/null
+++ b/tools/testing/selftests/bpf/libarena/selftests/test_spmc.bpf.c
@@ -0,0 +1,194 @@
+// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
+
+#include <libarena/common.h>
+
+#include <libarena/asan.h>
+#include <libarena/spmc.h>
+
+/*
+ * NOTE: These selftests only test for the single-threaded use case, which for
+ * Lev-Chase queues is obviously the simplest one. Still, it is important to
+ * exercise the API to ensure it passes verification and basic checks.
+ */
+
+SEC("syscall")
+int test_spmc_remove_empty(void)
+{
+       u64 val;
+       int ret;
+
+       struct spmc __arena *spmc = spmc_create();
+
+       if (!spmc)
+               return 1;
+
+       ret = spmc_owned_remove(spmc, &val);
+       if (ret != -ENOENT)
+               return 1;
+
+       spmc_destroy(spmc);
+
+       return 0;
+}
+
+SEC("syscall")
+int test_spmc_steal_empty(void)
+{
+       u64 val;
+       int ret;
+
+       struct spmc __arena *spmc = spmc_create();
+
+       if (!spmc)
+               return 1;
+
+       ret = spmc_steal(spmc, &val);
+       if (ret != -ENOENT)
+               return 1;
+
+       spmc_destroy(spmc);
+
+       return 0;
+}
+
+SEC("syscall")
+int test_spmc_steal_one(void)
+{
+       u64 val, newval;
+       int ret, i;
+
+       struct spmc __arena *spmc = spmc_create();
+
+       if (!spmc)
+               return 1;
+
+       for (i = 0; i < 10 && can_loop; i++) {
+               val = i;
+
+               ret = spmc_owned_add(spmc, val);
+               if (ret)
+                       return 1;
+
+               ret = spmc_steal(spmc, &newval);
+               if (ret)
+                       return 2;
+
+               if (val != newval)
+                       return 3;
+       }
+
+       spmc_destroy(spmc);
+
+       return 0;
+}
+
+SEC("syscall")
+int test_spmc_remove_one(void)
+{
+       u64 val, newval;
+       int ret, i;
+
+       struct spmc __arena *spmc = spmc_create();
+
+       if (!spmc)
+               return 1;
+
+       for (i = 0; i < 10 && can_loop; i++) {
+               val = i;
+
+               ret = spmc_owned_add(spmc, val);
+               if (ret)
+                       return 1;
+
+               ret = spmc_owned_remove(spmc, &newval);
+               if (ret)
+                       return 2;
+
+               if (val != newval)
+                       return 3;
+       }
+
+       spmc_destroy(spmc);
+
+       return 0;
+}
+
+SEC("syscall")
+int test_spmc_remove_many(void)
+{
+       u64 val, newval;
+       int ret, i;
+       u64 expected;
+
+       struct spmc __arena *spmc = spmc_create();
+
+       if (!spmc)
+               return 1;
+
+       for (i = 0; i < 500 && can_loop; i++) {
+               val = i;
+
+               ret = spmc_owned_add(spmc, val);
+               if (ret) {
+                       arena_stderr("%s:%d error %d\n", __func__, __LINE__, ret);
+                       return 1;
+               }
+       }
+
+       for (i = 0; i < 500 && can_loop; i++) {
+               ret = spmc_owned_remove(spmc, &newval);
+               if (ret) {
+                       arena_stderr("%s:%d error %d\n", __func__, __LINE__, ret);
+                       return 1;
+               }
+
+               expected = 500 - 1 - i;
+               if (newval != expected) {
+                       arena_stderr("%s:%d expected %llu found %llu\n", __func__, __LINE__, expected, newval);
+                       return 1;
+               }
+       }
+
+       spmc_destroy(spmc);
+
+       return 0;
+}
+
+SEC("syscall")
+int test_spmc_steal_many(void)
+{
+       u64 val, newval;
+       int ret, i;
+
+       struct spmc __arena *spmc = spmc_create();
+
+       if (!spmc)
+               return 1;
+
+       for (i = 0; i < 500 && can_loop; i++) {
+               val = i;
+
+               ret = spmc_owned_add(spmc, val);
+               if (ret) {
+                       arena_stderr("%s:%d error %d\n", __func__, __LINE__, ret);
+                       return 1;
+               }
+       }
+
+       for (i = 0; i < 500 && can_loop; i++) {
+               ret = spmc_steal(spmc, &newval);
+               if (ret) {
+                       arena_stderr("%s:%d error %d\n", __func__, __LINE__, ret);
+                       return 1;
+               }
+
+               if (newval != i) {
+                       arena_stderr("%s:%d expected %d found %llu\n", __func__, __LINE__, i, newval);
+                       return 1;
+               }
+       }
+
+       spmc_destroy(spmc);
+
+       return 0;
+}

diff --git a/tools/testing/selftests/bpf/libarena/src/spmc.bpf.c b/tools/testing/selftests/bpf/libarena/src/spmc.bpf.c
new file mode 100644 (file)
index 0000000..42732b7
--- /dev/null
+++ b/tools/testing/selftests/bpf/libarena/src/spmc.bpf.c
@@ -0,0 +1,234 @@
+// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
+/*
+ * Copyright (c) 2025-2026 Meta Platforms, Inc. and affiliates.
+ * Copyright (c) 2025-2026 Emil Tsalapatis <etsal@meta.com>
+ */
+
+#include <bpf_atomic.h>
+
+#include <libarena/common.h>
+
+#include <libarena/asan.h>
+#include <libarena/spmc.h>
+
+static inline
+u64 spmc_arr_size(volatile struct spmc_arr __arena *spmc_arr)
+{
+       return SPMC_ARR_BASESZ << spmc_arr->order;
+}
+
+static inline
+u64 spmc_arr_get(volatile struct spmc_arr __arena *spmc_arr, u64 ind)
+{
+       u64 ret = READ_ONCE(spmc_arr->data[ind % spmc_arr_size(spmc_arr)]);
+
+       return ret;
+}
+
+static inline
+void spmc_arr_put(volatile struct spmc_arr __arena *spmc_arr, u64 ind, u64 value)
+{
+       WRITE_ONCE(spmc_arr->data[ind % spmc_arr_size(spmc_arr)], value);
+}
+
+static inline
+void spmc_arr_copy(volatile struct spmc_arr __arena *dst,
+                  volatile struct spmc_arr __arena *src, u64 b, u64 t)
+{
+       u64 i;
+
+       for (i = t; i < b && can_loop; i++)
+               spmc_arr_put(dst, i, spmc_arr_get(src, i));
+}
+
+static inline
+int spmc_order_init(struct spmc __arena *spmc, int order)
+{
+       volatile struct spmc_arr __arena *arr = &spmc->arr[order];
+
+       if (unlikely(!spmc))
+               return -EINVAL;
+
+       if (order >= SPMC_ARR_ORDERS)
+               return -E2BIG;
+
+       /* Already allocated? */
+       if (arr->data)
+               return 0;
+
+       arr->data = arena_malloc((SPMC_ARR_BASESZ << order) * sizeof(*arr->data));
+       if (!arr->data)
+               return -ENOMEM;
+
+       return 0;
+}
+
+__weak
+int spmc_owned_add(struct spmc __arena *spmc, u64 val)
+{
+       volatile struct spmc_arr __arena *newarr;
+       volatile struct spmc_arr __arena *arr;
+       ssize_t sz;
+       u64 b, t;
+       int ret;
+
+       if (unlikely(!spmc))
+               return -EINVAL;
+
+       /* 
+        * Bottom must always be read first, also
+        * see spmc_steal().
+        */
+       b = smp_load_acquire(&spmc->bottom);
+       t = READ_ONCE(spmc->top);
+       arr = READ_ONCE(spmc->cur);
+
+       sz = b - t;
+       if (sz >= spmc_arr_size(arr) - 1) {
+               ret = spmc_order_init(spmc, arr->order + 1);
+               if (ret)
+                       return ret;
+
+               newarr = &spmc->arr[arr->order + 1];
+
+               spmc_arr_copy(newarr, arr, b, t);
+               smp_store_release(&spmc->cur, newarr);
+               arr = newarr;
+       }
+
+       spmc_arr_put(arr, b, val);
+       smp_store_release(&spmc->bottom, b + 1);
+
+       return 0;
+}
+
+
+__weak
+int spmc_owned_remove(struct spmc __arena *spmc, u64 *val)
+{
+       volatile struct spmc_arr __arena *arr;
+       int ret = 0;
+       ssize_t sz;
+       u64 value;
+       u64 b, t;
+
+       if (unlikely(!spmc || !val))
+               return -EINVAL;
+
+       b = READ_ONCE(spmc->bottom) - 1;
+       WRITE_ONCE(spmc->bottom, b);
+       smp_mb();
+
+       t = READ_ONCE(spmc->top);
+       arr = READ_ONCE(spmc->cur);
+
+       sz = b - t;
+       if (sz < 0) {
+               WRITE_ONCE(spmc->bottom, t);
+               return -ENOENT;
+       }
+
+       value = spmc_arr_get(arr, b);
+       if (sz > 0) {
+               *val = value;
+               return 0;
+       }
+
+       if (cmpxchg(&spmc->top, t, t + 1) != t)
+               ret = -EAGAIN;
+
+       WRITE_ONCE(spmc->bottom, t + 1);
+
+       if (ret)
+               return ret;
+
+       *val = value;
+
+       return 0;
+}
+
+__weak
+int spmc_steal(struct spmc __arena *spmc, u64 *val)
+{
+       volatile struct spmc_arr __arena *arr;
+       ssize_t sz;
+       u64 value;
+       u64 b, t;
+
+       if (unlikely(!spmc || !val))
+               return -EINVAL;
+
+       /*
+        * It is important that t is read before b for
+        * stealers to avoid racing with the owner.
+        * Races between stealers are dealt with using
+        * CAS to increment the top value below.
+        */
+       t = smp_load_acquire(&spmc->top);
+       b = smp_load_acquire(&spmc->bottom);
+
+       sz = b - t;
+       if (sz <= 0)
+               return -ENOENT;
+
+       arr = smp_load_acquire(&spmc->cur);
+       value = spmc_arr_get(arr, t);
+
+       if (cmpxchg(&spmc->top, t, t + 1) != t)
+               return -EAGAIN;
+
+       *val = value;
+
+       return 0;
+}
+
+
+__weak
+struct spmc __arena *spmc_create(void)
+{
+       /*
+        * Marked as volatile because otherwise the array
+        * reference in the internal loop gets demoted to
+        * scalar and the program fails verification.
+        */
+       struct spmc __arena *volatile spmc;
+       int ret, i;
+
+       spmc = arena_malloc(sizeof(*spmc));
+       if (!spmc)
+               return NULL;
+
+       spmc->bottom = 0;
+       spmc->top = 0;
+
+       for (i = 0; i < SPMC_ARR_ORDERS && can_loop; i++) {
+               spmc->arr[i].data = NULL;
+               spmc->arr[i].order = i;
+       }
+
+       ret = spmc_order_init((struct spmc __arena *)spmc, 0);
+       if (ret) {
+               arena_free(spmc);
+               return NULL;
+       }
+
+       spmc->cur = &spmc->arr[0];
+
+       return (struct spmc __arena *)spmc;
+}
+
+__weak
+int spmc_destroy(struct spmc __arena *spmc)
+{
+       int i;
+
+       if (unlikely(!spmc))
+               return -EINVAL;
+
+       for (i = 0; i < SPMC_ARR_ORDERS && can_loop; i++)
+               arena_free(spmc->arr[i].data);
+
+       arena_free(spmc);
+
+       return 0;
+}