Resguardo inicial

This commit is contained in:
lastdragon
2026-09-06 12:53:21 -06:00
parent 662efe3c2a
commit ad2088d178
94853 changed files with 43200404 additions and 0 deletions
@@ -0,0 +1,92 @@
# SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
# Copyright (c) 2026 Meta Platforms, Inc. and affiliates.
.PHONY: clean
# Defaults for standalone builds
CLANG ?= clang
BPFTOOL ?= bpftool
LDLIBS ?= -lbpf -lelf -lz -lrt -lpthread -lzstd
ifeq ($(V),1)
Q =
msg =
else
Q ?= @
msg = @printf ' %-8s%s %s%s\n' "$(1)" "$(if $(2), [$(2)])" "$(notdir $(3))" "$(if $(4), $(4))";
endif
IS_LITTLE_ENDIAN = $(shell $(CC) -dM -E - </dev/null | \
grep 'define __BYTE_ORDER__ __ORDER_LITTLE_ENDIAN__')
BPF_TARGET_ENDIAN ?= $(if $(IS_LITTLE_ENDIAN),--target=bpfel,--target=bpfeb)
LIBARENA=$(abspath .)
BPFDIR=$(abspath $(LIBARENA)/..)
INCLUDE_DIR ?= $(BPFDIR)/tools/include
LIBBPF_INCLUDE ?= $(INCLUDE_DIR)
# Scan src/ and selftests/ to generate the final binaries
LIBARENA_SOURCES = $(wildcard $(LIBARENA)/src/*.bpf.c) $(wildcard $(LIBARENA)/selftests/*.bpf.c)
LIBARENA_OBJECTS = $(notdir $(LIBARENA_SOURCES:.bpf.c=.bpf.o))
LIBARENA_OBJECTS_ASAN = $(notdir $(LIBARENA_SOURCES:.bpf.c=_asan.bpf.o))
INCLUDES = -I$(LIBARENA)/include -I$(BPFDIR)
ifneq ($(INCLUDE_DIR),)
INCLUDES += -I$(INCLUDE_DIR)
endif
ifneq ($(LIBBPF_INCLUDE),)
INCLUDES += -I$(LIBBPF_INCLUDE)
endif
ASAN_FLAGS = -fsanitize=kernel-address -fno-stack-protector -fno-builtin
ASAN_FLAGS += -mllvm -asan-instrument-address-spaces=1 -mllvm -asan-shadow-addr-space=1
ASAN_FLAGS += -mllvm -asan-use-stack-safety=0 -mllvm -asan-stack=0
ASAN_FLAGS += -mllvm -asan-kernel=1
ASAN_FLAGS += -mllvm -asan-constructor-kind=none
ASAN_FLAGS += -mllvm -asan-destructor-kind=none
# ENABLE_ATOMICS_TESTS required because we use arena spinlocks
override BPF_CFLAGS += -DENABLE_ATOMICS_TESTS
override BPF_CFLAGS += -O2 -g
override BPF_CFLAGS += -Wno-incompatible-pointer-types-discards-qualifiers
# Required for suppressing harmless vmlinux.h-related warnings.
override BPF_CFLAGS += -Wno-missing-declarations
override BPF_CFLAGS += $(INCLUDES)
CFLAGS = -O2 -no-pie
CFLAGS += $(INCLUDES)
vpath %.bpf.c $(LIBARENA)/src $(LIBARENA)/selftests
vpath %.c $(LIBARENA)/src $(LIBARENA)/selftests
skeletons: libarena.skel.h libarena_asan.skel.h
.PHONY: skeletons
libarena_asan.skel.h: libarena_asan.bpf.o
$(call msg,GEN-SKEL,libarena,$@)
$(Q)$(BPFTOOL) gen skeleton $< name "libarena_asan" > $@
libarena.skel.h: libarena.bpf.o
$(call msg,GEN-SKEL,libarena,$@)
$(Q)$(BPFTOOL) gen skeleton $< name "libarena" > $@
libarena_asan.bpf.o: $(LIBARENA_OBJECTS_ASAN)
$(call msg,GEN-OBJ,libarena,$@)
$(Q)$(BPFTOOL) gen object $@ $^
libarena.bpf.o: $(LIBARENA_OBJECTS)
$(call msg,GEN-OBJ,libarena,$@)
$(Q)$(BPFTOOL) gen object $@ $^
%_asan.bpf.o: %.bpf.c
$(call msg,CLNG-BPF,libarena,$@)
$(Q)$(CLANG) $(BPF_CFLAGS) $(ASAN_FLAGS) -DBPF_ARENA_ASAN $(BPF_TARGET_ENDIAN) -c $< -o $@
%.bpf.o: %.bpf.c
$(call msg,CLNG-BPF,libarena,$@)
$(Q)$(CLANG) $(BPF_CFLAGS) $(BPF_TARGET_ENDIAN) -c $< -o $@
clean:
$(Q)rm -f *.skel.h *.bpf.o *.linked*.o
@@ -0,0 +1,74 @@
/* SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause) */
/* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */
#pragma once
#ifndef WRITE_ONCE
#define WRITE_ONCE(x, val) ((*(volatile typeof(x) *) &(x)) = (val))
#endif
#ifndef NUMA_NO_NODE
#define NUMA_NO_NODE (-1)
#endif
#ifndef arena_container_of
#define arena_container_of(ptr, type, member) \
({ \
void __arena *__mptr = (void __arena *)(ptr); \
((type *)(__mptr - offsetof(type, member))); \
})
#endif
#ifdef __BPF__ /* when compiled as bpf program */
#ifndef PAGE_SIZE
#define PAGE_SIZE __PAGE_SIZE
/*
* for older kernels try sizeof(struct genradix_node)
* or flexible:
* static inline long __bpf_page_size(void) {
* return bpf_core_enum_value(enum page_size_enum___l, __PAGE_SIZE___l) ?: sizeof(struct genradix_node);
* }
* but generated code is not great.
*/
#endif
#if defined(__BPF_FEATURE_ADDR_SPACE_CAST) && !defined(BPF_ARENA_FORCE_ASM)
#define __arena __attribute__((address_space(1))) __attribute__((btf_type_tag("arena")))
#define __arena_global __attribute__((address_space(1)))
#define cast_kern(ptr) /* nop for bpf prog. emitted by LLVM */
#define cast_user(ptr) /* nop for bpf prog. emitted by LLVM */
#else
#define __arena __attribute__((btf_type_tag("arena")))
#define __arena_global SEC(".addr_space.1")
#define cast_kern(ptr) bpf_addr_space_cast(ptr, 0, 1)
#define cast_user(ptr) bpf_addr_space_cast(ptr, 1, 0)
#endif
void __arena* bpf_arena_alloc_pages(void *map, void __arena *addr, __u32 page_cnt,
int node_id, __u64 flags) __ksym __weak;
int bpf_arena_reserve_pages(void *map, void __arena *addr, __u32 page_cnt) __ksym __weak;
void bpf_arena_free_pages(void *map, void __arena *ptr, __u32 page_cnt) __ksym __weak;
#define arena_base(map) ((void __arena *)((struct bpf_arena *)(map))->user_vm_start)
#else /* when compiled as user space code */
#define __arena
#define cast_kern(ptr) /* nop for user space */
#define cast_user(ptr) /* nop for user space */
__weak char arena[1];
#ifndef offsetof
#define offsetof(type, member) ((unsigned long)&((type *)0)->member)
#endif
static inline void __arena* bpf_arena_alloc_pages(void *map, void *addr, __u32 page_cnt,
int node_id, __u64 flags)
{
return NULL;
}
static inline void bpf_arena_free_pages(void *map, void __arena *ptr, __u32 page_cnt)
{
}
#endif
@@ -0,0 +1,543 @@
// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2025 Meta Platforms, Inc. and affiliates. */
#ifndef BPF_ARENA_SPIN_LOCK_H
#define BPF_ARENA_SPIN_LOCK_H
#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include <bpf_atomic.h>
#define arch_mcs_spin_lock_contended_label(l, label) smp_cond_load_acquire_label(l, VAL, label)
#define arch_mcs_spin_unlock_contended(l) smp_store_release((l), 1)
#if defined(ENABLE_ATOMICS_TESTS) && defined(__BPF_FEATURE_ADDR_SPACE_CAST)
#define EBUSY 16
#define EOPNOTSUPP 95
#define ETIMEDOUT 110
extern unsigned long CONFIG_NR_CPUS __kconfig;
/*
* Typically, we'd just rely on the definition in vmlinux.h for qspinlock, but
* PowerPC overrides the definition to define lock->val as u32 instead of
* atomic_t, leading to compilation errors. Import a local definition below so
* that we don't depend on the vmlinux.h version.
*/
struct __qspinlock {
union {
atomic_t val;
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
struct {
u8 locked;
u8 pending;
};
struct {
u16 locked_pending;
u16 tail;
};
#else
struct {
u16 tail;
u16 locked_pending;
};
struct {
u8 reserved[2];
u8 pending;
u8 locked;
};
#endif
};
};
#define arena_spinlock_t struct __qspinlock
/* FIXME: Using typedef causes CO-RE relocation error */
/* typedef struct qspinlock arena_spinlock_t; */
struct arena_mcs_spinlock {
struct arena_mcs_spinlock __arena *next;
int locked;
int count;
};
struct arena_qnode {
struct arena_mcs_spinlock mcs;
};
#define _Q_MAX_NODES 4
#define _Q_PENDING_LOOPS 1
/*
* Bitfields in the atomic value:
*
* 0- 7: locked byte
* 8: pending
* 9-15: not used
* 16-17: tail index
* 18-31: tail cpu (+1)
*/
#define _Q_MAX_CPUS 1024
#define _Q_SET_MASK(type) (((1U << _Q_ ## type ## _BITS) - 1)\
<< _Q_ ## type ## _OFFSET)
#define _Q_LOCKED_OFFSET 0
#define _Q_LOCKED_BITS 8
#define _Q_LOCKED_MASK _Q_SET_MASK(LOCKED)
#define _Q_PENDING_OFFSET (_Q_LOCKED_OFFSET + _Q_LOCKED_BITS)
#define _Q_PENDING_BITS 8
#define _Q_PENDING_MASK _Q_SET_MASK(PENDING)
#define _Q_TAIL_IDX_OFFSET (_Q_PENDING_OFFSET + _Q_PENDING_BITS)
#define _Q_TAIL_IDX_BITS 2
#define _Q_TAIL_IDX_MASK _Q_SET_MASK(TAIL_IDX)
#define _Q_TAIL_CPU_OFFSET (_Q_TAIL_IDX_OFFSET + _Q_TAIL_IDX_BITS)
#define _Q_TAIL_CPU_BITS (32 - _Q_TAIL_CPU_OFFSET)
#define _Q_TAIL_CPU_MASK _Q_SET_MASK(TAIL_CPU)
#define _Q_TAIL_OFFSET _Q_TAIL_IDX_OFFSET
#define _Q_TAIL_MASK (_Q_TAIL_IDX_MASK | _Q_TAIL_CPU_MASK)
#define _Q_LOCKED_VAL (1U << _Q_LOCKED_OFFSET)
#define _Q_PENDING_VAL (1U << _Q_PENDING_OFFSET)
/*
* The qnodes are marked __weak so we can define them in the header
* while still ensuring all compilation units use the same struct
* instance.
*/
struct arena_qnode __weak __arena __hidden qnodes[_Q_MAX_CPUS][_Q_MAX_NODES];
static inline u32 encode_tail(int cpu, int idx)
{
u32 tail;
tail = (cpu + 1) << _Q_TAIL_CPU_OFFSET;
tail |= idx << _Q_TAIL_IDX_OFFSET; /* assume < 4 */
return tail;
}
static inline struct arena_mcs_spinlock __arena *decode_tail(u32 tail)
{
u32 cpu = (tail >> _Q_TAIL_CPU_OFFSET) - 1;
u32 idx = (tail & _Q_TAIL_IDX_MASK) >> _Q_TAIL_IDX_OFFSET;
return &qnodes[cpu][idx].mcs;
}
static inline
struct arena_mcs_spinlock __arena *grab_mcs_node(struct arena_mcs_spinlock __arena *base, int idx)
{
return &((struct arena_qnode __arena *)base + idx)->mcs;
}
#define _Q_LOCKED_PENDING_MASK (_Q_LOCKED_MASK | _Q_PENDING_MASK)
/**
* xchg_tail - Put in the new queue tail code word & retrieve previous one
* @lock : Pointer to queued spinlock structure
* @tail : The new queue tail code word
* Return: The previous queue tail code word
*
* xchg(lock, tail)
*
* p,*,* -> n,*,* ; prev = xchg(lock, node)
*/
static __always_inline u32 xchg_tail(arena_spinlock_t __arena *lock, u32 tail)
{
u32 old, new;
old = atomic_read(&lock->val);
do {
new = (old & _Q_LOCKED_PENDING_MASK) | tail;
/*
* We can use relaxed semantics since the caller ensures that
* the MCS node is properly initialized before updating the
* tail.
*/
/* These loops are not expected to stall, but we still need to
* prove to the verifier they will terminate eventually.
*/
cond_break_label(out);
} while (!atomic_try_cmpxchg_relaxed(&lock->val, &old, new));
return old;
out:
bpf_printk("RUNTIME ERROR: %s unexpected cond_break exit!!!", __func__);
return old;
}
/**
* clear_pending - clear the pending bit.
* @lock: Pointer to queued spinlock structure
*
* *,1,* -> *,0,*
*/
static __always_inline void clear_pending(arena_spinlock_t __arena *lock)
{
WRITE_ONCE(lock->pending, 0);
}
/**
* clear_pending_set_locked - take ownership and clear the pending bit.
* @lock: Pointer to queued spinlock structure
*
* *,1,0 -> *,0,1
*
* Lock stealing is not allowed if this function is used.
*/
static __always_inline void clear_pending_set_locked(arena_spinlock_t __arena *lock)
{
WRITE_ONCE(lock->locked_pending, _Q_LOCKED_VAL);
}
/**
* set_locked - Set the lock bit and own the lock
* @lock: Pointer to queued spinlock structure
*
* *,*,0 -> *,0,1
*/
static __always_inline void set_locked(arena_spinlock_t __arena *lock)
{
WRITE_ONCE(lock->locked, _Q_LOCKED_VAL);
}
static __always_inline
u32 arena_fetch_set_pending_acquire(arena_spinlock_t __arena *lock)
{
u32 old, new;
old = atomic_read(&lock->val);
do {
new = old | _Q_PENDING_VAL;
/*
* These loops are not expected to stall, but we still need to
* prove to the verifier they will terminate eventually.
*/
cond_break_label(out);
} while (!atomic_try_cmpxchg_acquire(&lock->val, &old, new));
return old;
out:
bpf_printk("RUNTIME ERROR: %s unexpected cond_break exit!!!", __func__);
return old;
}
/**
* arena_spin_trylock - try to acquire the queued spinlock
* @lock : Pointer to queued spinlock structure
* Return: 1 if lock acquired, 0 if failed
*/
static __always_inline int arena_spin_trylock(arena_spinlock_t __arena *lock)
{
int val = atomic_read(&lock->val);
if (unlikely(val))
return 0;
return likely(atomic_try_cmpxchg_acquire(&lock->val, &val, _Q_LOCKED_VAL));
}
__noinline __weak
int arena_spin_lock_slowpath(arena_spinlock_t __arena *lock, u32 val)
{
struct arena_mcs_spinlock __arena *prev, *next, *node0, *node;
int ret = -ETIMEDOUT;
u32 old, tail;
int idx;
/*
* Wait for in-progress pending->locked hand-overs with a bounded
* number of spins so that we guarantee forward progress.
*
* 0,1,0 -> 0,0,1
*/
if (val == _Q_PENDING_VAL) {
int cnt = _Q_PENDING_LOOPS;
val = atomic_cond_read_relaxed_label(&lock->val,
(VAL != _Q_PENDING_VAL) || !cnt--,
release_err);
}
/*
* If we observe any contention; queue.
*/
if (val & ~_Q_LOCKED_MASK)
goto queue;
/*
* trylock || pending
*
* 0,0,* -> 0,1,* -> 0,0,1 pending, trylock
*/
val = arena_fetch_set_pending_acquire(lock);
/*
* If we observe contention, there is a concurrent locker.
*
* Undo and queue; our setting of PENDING might have made the
* n,0,0 -> 0,0,0 transition fail and it will now be waiting
* on @next to become !NULL.
*/
if (unlikely(val & ~_Q_LOCKED_MASK)) {
/* Undo PENDING if we set it. */
if (!(val & _Q_PENDING_MASK))
clear_pending(lock);
goto queue;
}
/*
* We're pending, wait for the owner to go away.
*
* 0,1,1 -> *,1,0
*
* this wait loop must be a load-acquire such that we match the
* store-release that clears the locked bit and create lock
* sequentiality; this is because not all
* clear_pending_set_locked() implementations imply full
* barriers.
*/
if (val & _Q_LOCKED_MASK)
(void)smp_cond_load_acquire_label(&lock->locked, !VAL, release_err);
/*
* take ownership and clear the pending bit.
*
* 0,1,0 -> 0,0,1
*/
clear_pending_set_locked(lock);
return 0;
/*
* End of pending bit optimistic spinning and beginning of MCS
* queuing.
*/
queue:
node0 = &(qnodes[bpf_get_smp_processor_id()])[0].mcs;
idx = node0->count++;
tail = encode_tail(bpf_get_smp_processor_id(), idx);
/*
* 4 nodes are allocated based on the assumption that there will not be
* nested NMIs taking spinlocks. That may not be true in some
* architectures even though the chance of needing more than 4 nodes
* will still be extremely unlikely. When that happens, we simply return
* an error. Original qspinlock has a trylock fallback in this case.
*/
if (unlikely(idx >= _Q_MAX_NODES)) {
ret = -EBUSY;
goto release_node_err;
}
node = grab_mcs_node(node0, idx);
/*
* Ensure that we increment the head node->count before initialising
* the actual node. If the compiler is kind enough to reorder these
* stores, then an IRQ could overwrite our assignments.
*/
barrier();
node->locked = 0;
node->next = NULL;
/*
* We touched a (possibly) cold cacheline in the per-cpu queue node;
* attempt the trylock once more in the hope someone let go while we
* weren't watching.
*/
if (arena_spin_trylock(lock))
goto release;
/*
* Ensure that the initialisation of @node is complete before we
* publish the updated tail via xchg_tail() and potentially link
* @node into the waitqueue via WRITE_ONCE(prev->next, node) below.
*/
smp_wmb();
/*
* Publish the updated tail.
* We have already touched the queueing cacheline; don't bother with
* pending stuff.
*
* p,*,* -> n,*,*
*/
old = xchg_tail(lock, tail);
next = NULL;
/*
* if there was a previous node; link it and wait until reaching the
* head of the waitqueue.
*/
if (old & _Q_TAIL_MASK) {
prev = decode_tail(old);
/* Link @node into the waitqueue. */
WRITE_ONCE(prev->next, node);
(void)arch_mcs_spin_lock_contended_label(&node->locked, release_node_err);
/*
* While waiting for the MCS lock, the next pointer may have
* been set by another lock waiter. We cannot prefetch here
* due to lack of equivalent instruction in BPF ISA.
*/
next = READ_ONCE(node->next);
}
/*
* we're at the head of the waitqueue, wait for the owner & pending to
* go away.
*
* *,x,y -> *,0,0
*
* this wait loop must use a load-acquire such that we match the
* store-release that clears the locked bit and create lock
* sequentiality; this is because the set_locked() function below
* does not imply a full barrier.
*/
val = atomic_cond_read_acquire_label(&lock->val, !(VAL & _Q_LOCKED_PENDING_MASK),
release_node_err);
/*
* claim the lock:
*
* n,0,0 -> 0,0,1 : lock, uncontended
* *,*,0 -> *,*,1 : lock, contended
*
* If the queue head is the only one in the queue (lock value == tail)
* and nobody is pending, clear the tail code and grab the lock.
* Otherwise, we only need to grab the lock.
*/
/*
* In the PV case we might already have _Q_LOCKED_VAL set, because
* of lock stealing; therefore we must also allow:
*
* n,0,1 -> 0,0,1
*
* Note: at this point: (val & _Q_PENDING_MASK) == 0, because of the
* above wait condition, therefore any concurrent setting of
* PENDING will make the uncontended transition fail.
*/
if ((val & _Q_TAIL_MASK) == tail) {
if (atomic_try_cmpxchg_relaxed(&lock->val, &val, _Q_LOCKED_VAL))
goto release; /* No contention */
}
/*
* Either somebody is queued behind us or _Q_PENDING_VAL got set
* which will then detect the remaining tail and queue behind us
* ensuring we'll see a @next.
*/
set_locked(lock);
/*
* contended path; wait for next if not observed yet, release.
*/
if (!next)
next = smp_cond_load_relaxed_label(&node->next, (VAL), release_node_err);
arch_mcs_spin_unlock_contended(&next->locked);
release:;
/*
* release the node
*
* Doing a normal dec vs this_cpu_dec is fine. An upper context always
* decrements count it incremented before returning, thus we're fine.
* For contexts interrupting us, they either observe our dec or not.
* Just ensure the compiler doesn't reorder this statement, as a
* this_cpu_dec implicitly implied that.
*/
barrier();
node0->count--;
return 0;
release_node_err:
barrier();
node0->count--;
goto release_err;
release_err:
return ret;
}
/**
* arena_spin_lock - acquire a queued spinlock
* @lock: Pointer to queued spinlock structure
*
* On error, returned value will be negative.
* On success, zero is returned.
*
* The return value _must_ be tested against zero for success,
* instead of checking it against negative, for passing the
* BPF verifier.
*
* The user should do:
* if (arena_spin_lock(...) != 0) // failure
* or
* if (arena_spin_lock(...) == 0) // success
* or
* if (arena_spin_lock(...)) // failure
* or
* if (!arena_spin_lock(...)) // success
* instead of:
* if (arena_spin_lock(...) < 0) // failure
*
* The return value can still be inspected later.
*/
static __always_inline int arena_spin_lock(arena_spinlock_t __arena *lock)
{
int val = 0;
if (CONFIG_NR_CPUS > 1024)
return -EOPNOTSUPP;
bpf_preempt_disable();
if (likely(atomic_try_cmpxchg_acquire(&lock->val, &val, _Q_LOCKED_VAL)))
return 0;
val = arena_spin_lock_slowpath(lock, val);
/* FIXME: bpf_assert_range(-MAX_ERRNO, 0) once we have it working for all cases. */
if (val)
bpf_preempt_enable();
return val;
}
/**
* arena_spin_unlock - release a queued spinlock
* @lock : Pointer to queued spinlock structure
*/
static __always_inline void arena_spin_unlock(arena_spinlock_t __arena *lock)
{
/*
* unlock() needs release semantics:
*/
smp_store_release(&lock->locked, 0);
bpf_preempt_enable();
}
#define arena_spin_lock_irqsave(lock, flags) \
({ \
int __ret; \
bpf_local_irq_save(&(flags)); \
__ret = arena_spin_lock((lock)); \
if (__ret) \
bpf_local_irq_restore(&(flags)); \
(__ret); \
})
#define arena_spin_unlock_irqrestore(lock, flags) \
({ \
arena_spin_unlock((lock)); \
bpf_local_irq_restore(&(flags)); \
})
#endif
#endif /* BPF_ARENA_SPIN_LOCK_H */
@@ -0,0 +1,142 @@
// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2025 Meta Platforms, Inc. and affiliates. */
#ifndef BPF_ATOMIC_H
#define BPF_ATOMIC_H
#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include <bpf_may_goto.h>
extern bool CONFIG_X86_64 __kconfig __weak;
/*
* __unqual_typeof(x) - Declare an unqualified scalar type, leaving
* non-scalar types unchanged,
*
* Prefer C11 _Generic for better compile-times and simpler code. Note: 'char'
* is not type-compatible with 'signed char', and we define a separate case.
*
* This is copied verbatim from kernel's include/linux/compiler_types.h, but
* with default expression (for pointers) changed from (x) to (typeof(x)0).
*
* This is because LLVM has a bug where for lvalue (x), it does not get rid of
* an extra address_space qualifier, but does in case of rvalue (typeof(x)0).
* Hence, for pointers, we need to create an rvalue expression to get the
* desired type. See https://github.com/llvm/llvm-project/issues/53400.
*/
#define __scalar_type_to_expr_cases(type) \
unsigned type : (unsigned type)0, signed type : (signed type)0
#define __unqual_typeof(x) \
typeof(_Generic((x), \
char: (char)0, \
__scalar_type_to_expr_cases(char), \
__scalar_type_to_expr_cases(short), \
__scalar_type_to_expr_cases(int), \
__scalar_type_to_expr_cases(long), \
__scalar_type_to_expr_cases(long long), \
default: (typeof(x))0))
/* No-op for BPF */
#define cpu_relax() ({})
#define READ_ONCE(x) (*(volatile typeof(x) *)&(x))
#ifndef WRITE_ONCE
#define WRITE_ONCE(x, val) ((*(volatile typeof(x) *)&(x)) = (val))
#endif
#define cmpxchg(p, old, new) __sync_val_compare_and_swap((p), old, new)
#define try_cmpxchg(p, pold, new) \
({ \
__unqual_typeof(*(pold)) __o = *(pold); \
__unqual_typeof(*(p)) __r = cmpxchg(p, __o, new); \
if (__r != __o) \
*(pold) = __r; \
__r == __o; \
})
#define try_cmpxchg_relaxed(p, pold, new) try_cmpxchg(p, pold, new)
#define try_cmpxchg_acquire(p, pold, new) try_cmpxchg(p, pold, new)
#define smp_mb() \
({ \
volatile unsigned long __val; \
__sync_fetch_and_add(&__val, 0); \
})
#define smp_rmb() \
({ \
if (!CONFIG_X86_64) \
smp_mb(); \
else \
barrier(); \
})
#define smp_wmb() \
({ \
if (!CONFIG_X86_64) \
smp_mb(); \
else \
barrier(); \
})
/* Control dependency provides LOAD->STORE, provide LOAD->LOAD */
#define smp_acquire__after_ctrl_dep() ({ smp_rmb(); })
#define smp_load_acquire(p) \
({ \
__unqual_typeof(*(p)) __v = READ_ONCE(*(p)); \
if (!CONFIG_X86_64) \
smp_mb(); \
barrier(); \
__v; \
})
#define smp_store_release(p, val) \
({ \
if (!CONFIG_X86_64) \
smp_mb(); \
barrier(); \
WRITE_ONCE(*(p), val); \
})
#define smp_cond_load_relaxed_label(p, cond_expr, label) \
({ \
typeof(p) __ptr = (p); \
__unqual_typeof(*(p)) VAL; \
for (;;) { \
VAL = (__unqual_typeof(*(p)))READ_ONCE(*__ptr); \
if (cond_expr) \
break; \
cond_break_label(label); \
cpu_relax(); \
} \
(typeof(*(p)))VAL; \
})
#define smp_cond_load_acquire_label(p, cond_expr, label) \
({ \
__unqual_typeof(*p) __val = \
smp_cond_load_relaxed_label(p, cond_expr, label); \
smp_acquire__after_ctrl_dep(); \
(typeof(*(p)))__val; \
})
#define atomic_read(p) READ_ONCE((p)->counter)
#define atomic_cond_read_relaxed_label(p, cond_expr, label) \
smp_cond_load_relaxed_label(&(p)->counter, cond_expr, label)
#define atomic_cond_read_acquire_label(p, cond_expr, label) \
smp_cond_load_acquire_label(&(p)->counter, cond_expr, label)
#define atomic_try_cmpxchg_relaxed(p, pold, new) \
try_cmpxchg_relaxed(&(p)->counter, pold, new)
#define atomic_try_cmpxchg_acquire(p, pold, new) \
try_cmpxchg_acquire(&(p)->counter, pold, new)
#endif /* BPF_ATOMIC_H */
@@ -0,0 +1,84 @@
#pragma once
/*
* Note that cond_break can only be portably used in the body of a breakable
* construct, whereas can_loop can be used anywhere.
*/
#ifdef __BPF_FEATURE_MAY_GOTO
#define can_loop \
({ __label__ l_break, l_continue; \
bool ret = true; \
asm volatile goto("may_goto %l[l_break]" \
:::: l_break); \
goto l_continue; \
l_break: ret = false; \
l_continue:; \
ret; \
})
#define __cond_break(expr) \
({ __label__ l_break, l_continue; \
asm volatile goto("may_goto %l[l_break]" \
:::: l_break); \
goto l_continue; \
l_break: expr; \
l_continue:; \
})
#else
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
#define can_loop \
({ __label__ l_break, l_continue; \
bool ret = true; \
asm volatile goto("1:.byte 0xe5; \
.byte 0; \
.long ((%l[l_break] - 1b - 8) / 8) & 0xffff; \
.short 0" \
:::: l_break); \
goto l_continue; \
l_break: ret = false; \
l_continue:; \
ret; \
})
#define __cond_break(expr) \
({ __label__ l_break, l_continue; \
asm volatile goto("1:.byte 0xe5; \
.byte 0; \
.long ((%l[l_break] - 1b - 8) / 8) & 0xffff; \
.short 0" \
:::: l_break); \
goto l_continue; \
l_break: expr; \
l_continue:; \
})
#else
#define can_loop \
({ __label__ l_break, l_continue; \
bool ret = true; \
asm volatile goto("1:.byte 0xe5; \
.byte 0; \
.long (((%l[l_break] - 1b - 8) / 8) & 0xffff) << 16; \
.short 0" \
:::: l_break); \
goto l_continue; \
l_break: ret = false; \
l_continue:; \
ret; \
})
#define __cond_break(expr) \
({ __label__ l_break, l_continue; \
asm volatile goto("1:.byte 0xe5; \
.byte 0; \
.long (((%l[l_break] - 1b - 8) / 8) & 0xffff) << 16; \
.short 0" \
:::: l_break); \
goto l_continue; \
l_break: expr; \
l_continue:; \
})
#endif
#endif
#define cond_break __cond_break(break)
#define cond_break_label(label) __cond_break(goto label)
@@ -0,0 +1,101 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#pragma once
struct asan_init_args {
u64 arena_all_pages;
u64 arena_globals_pages;
};
int asan_init(struct asan_init_args *args);
extern volatile u64 __asan_shadow_memory_dynamic_address;
extern volatile u32 asan_reported;
extern volatile bool asan_inited;
extern volatile bool asan_report_once;
#ifdef __BPF__
#define ASAN_SHADOW_SHIFT 3
#define ASAN_SHADOW_SCALE (1ULL << ASAN_SHADOW_SHIFT)
#define ASAN_GRANULE_MASK ((1ULL << ASAN_SHADOW_SHIFT) - 1)
#define ASAN_GRANULE(addr) ((s8)((u32)(u64)((addr)) & ASAN_GRANULE_MASK))
#define __noasan __attribute__((no_sanitize("address")))
#ifdef BPF_ARENA_ASAN
static inline
s8 __arena *mem_to_shadow(void __arena *addr)
{
return (s8 __arena *)(((u32)(u64)addr >> ASAN_SHADOW_SHIFT) +
__asan_shadow_memory_dynamic_address);
}
__weak __noasan
bool asan_ready(void)
{
return __asan_shadow_memory_dynamic_address;
}
int asan_poison(void __arena *addr, s8 val, size_t size);
int asan_unpoison(void __arena *addr, size_t size);
bool asan_shadow_set(void __arena *addr);
/*
* Dummy calls to ensure the ASAN runtime's BTF information is present
* in every object file when compiling the runtime and local BPF code
* separately. The runtime calls are injected into the LLVM IR file
*/
#define DECLARE_ASAN_LOAD_STORE_SIZE(size) \
void __asan_store##size(intptr_t addr); \
void __asan_store##size##_noabort(intptr_t addr); \
void __asan_load##size(intptr_t addr); \
void __asan_load##size##_noabort(intptr_t addr); \
void __asan_report_store##size(intptr_t addr); \
void __asan_report_store##size##_noabort(intptr_t addr); \
void __asan_report_load##size(intptr_t addr); \
void __asan_report_load##size##_noabort(intptr_t addr);
DECLARE_ASAN_LOAD_STORE_SIZE(1);
DECLARE_ASAN_LOAD_STORE_SIZE(2);
DECLARE_ASAN_LOAD_STORE_SIZE(4);
DECLARE_ASAN_LOAD_STORE_SIZE(8);
void __asan_storeN(intptr_t addr, ssize_t size);
void __asan_storeN_noabort(intptr_t addr, ssize_t size);
void __asan_loadN(intptr_t addr, ssize_t size);
void __asan_loadN_noabort(intptr_t addr, ssize_t size);
/*
* Force LLVM to emit BTF information for the stubs,
* because the ASAN pass in LLVM by itself doesn't.
*/
#define ASAN_LOAD_STORE_SIZE(size) \
__asan_store##size, \
__asan_store##size##_noabort, \
__asan_load##size, \
__asan_load##size##_noabort, \
__asan_report_store##size, \
__asan_report_store##size##_noabort, \
__asan_report_load##size, \
__asan_report_load##size##_noabort
__attribute__((used))
static void (*__asan_btf_anchors[])(intptr_t) = {
ASAN_LOAD_STORE_SIZE(1),
ASAN_LOAD_STORE_SIZE(2),
ASAN_LOAD_STORE_SIZE(4),
ASAN_LOAD_STORE_SIZE(8),
};
#else /* BPF_ARENA_ASAN */
static inline int asan_poison(void __arena *addr, s8 val, size_t size) { return 0; }
static inline int asan_unpoison(void __arena *addr, size_t size) { return 0; }
static inline bool asan_shadow_set(void __arena *addr) { return 0; }
__weak bool asan_ready(void) { return true; }
#endif /* BPF_ARENA_ASAN */
#endif /* __BPF__ */
@@ -0,0 +1,81 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#pragma once
enum buddy_consts {
/*
* Minimum allocation is 1 << BUDDY_MIN_ALLOC_SHIFT.
* Larger sizes increase internal fragmentation, but smaller
* sizes increase the space overhead of the block metadata.
*/
BUDDY_MIN_ALLOC_SHIFT = 4,
BUDDY_MIN_ALLOC_BYTES = 1 << BUDDY_MIN_ALLOC_SHIFT,
/*
* How many orders the buddy allocator can serve. Minimum block
* size is 1 << BUDDY_MIN_ALLOC_SHIFT, maximum block size is
* 1 << (BUDDY_MIN_ALLOC_SHIFT + BUDDY_CHUNK_NUM_ORDERS - 1):
* Each block has size 1 << BUDDY_MIN_ALLOC_SHIFT, and the
* allocation orders are in [0, BUDDY_CHUNK_NUM_ORDERS).
* We keep two blocks of the maximum size to retain the
* property in the code that all blocks have a buddy.
* Higher values increase the maximum allocation size,
* but also the size of the metadata for each block.
*/
BUDDY_CHUNK_NUM_ORDERS = 1 << 4,
BUDDY_CHUNK_BYTES = BUDDY_MIN_ALLOC_BYTES << (BUDDY_CHUNK_NUM_ORDERS),
/* Offset of the buddy header within a free block, see buddy.bpf.c for details */
BUDDY_HEADER_OFF = 8,
/* The maximum number of blocks a chunk may have to track. */
BUDDY_CHUNK_ITEMS = 1 << (BUDDY_CHUNK_NUM_ORDERS),
BUDDY_CHUNK_OFFSET_MASK = BUDDY_CHUNK_BYTES - 1,
/*
* Alignment for chunk allocations based on bpf_arena_alloc_pages.
* The arena allocation kfunc does not have an alignment argument,
* but that is required for all block calculations in the chunk to
* work.
*/
BUDDY_VADDR_OFFSET = BUDDY_CHUNK_BYTES,
/* Total arena virtual address space the allocator can consume. */
BUDDY_VADDR_SIZE = BUDDY_CHUNK_BYTES << 10
};
struct buddy_header {
u32 prev_index; /* "Pointer" to the previous available allocation of the same size. */
u32 next_index; /* Same for the next allocation. */
};
/*
* We bring memory into the allocator 1 MiB at a time.
*/
struct buddy_chunk {
/* The order of the current allocation for a item. 4 bits per order. */
u8 orders[BUDDY_CHUNK_ITEMS / 2];
/*
* Bit to denote whether chunk is allocated. Size of the allocated/free
* chunk found from the orders array.
*/
u8 allocated[BUDDY_CHUNK_ITEMS / 8];
/* Freelists for O(1) allocation. */
u64 freelists[BUDDY_CHUNK_NUM_ORDERS];
struct buddy_chunk __arena *next;
};
struct buddy {
struct buddy_chunk __arena *first_chunk; /* Pointer to the chunk linked list. */
arena_spinlock_t lock; /* Allocator lock */
u64 vaddr; /* Allocation into reserved vaddr */
};
#ifdef __BPF__
int buddy_init(struct buddy __arena *buddy);
int buddy_destroy(struct buddy __arena *buddy);
int buddy_free(struct buddy __arena *buddy, void __arena *free);
void __arena *buddy_alloc(struct buddy __arena *buddy, size_t size);
#endif /* __BPF__ */
@@ -0,0 +1,93 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#pragma once
#ifdef __BPF__
#include <vmlinux.h>
#include <bpf_arena_common.h>
#include <bpf_arena_spin_lock.h>
#include <asm-generic/errno.h>
#ifndef __BPF_FEATURE_ADDR_SPACE_CAST
#error "Arena allocators require bpf_addr_space_cast feature"
#endif
#define arena_stdout(fmt, ...) bpf_stream_printk(1, (fmt), ##__VA_ARGS__)
#define arena_stderr(fmt, ...) bpf_stream_printk(2, (fmt), ##__VA_ARGS__)
#ifndef __maybe_unused
#define __maybe_unused __attribute__((__unused__))
#endif
#define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8)))
#define ARENA_PAGES (1UL << (32 - __builtin_ffs(__PAGE_SIZE) + 1))
struct {
__uint(type, BPF_MAP_TYPE_ARENA);
__uint(map_flags, BPF_F_MMAPABLE);
__uint(max_entries, ARENA_PAGES); /* number of pages */
#if defined(__TARGET_ARCH_arm64) || defined(__aarch64__)
__ulong(map_extra, (1ull << 32)); /* start of mmap() region */
#else
__ulong(map_extra, (1ull << 44)); /* start of mmap() region */
#endif
} arena __weak SEC(".maps");
/*
* This is a variable used to aid verification. The may_goto directive
* permits open-coded for loops, but requires that the index variable is
* imprecise. To force the variable to be imprecise, initialize it with
* the opaque volatile variable 0 instead of the constant 0.
*/
extern const volatile u32 zero;
extern volatile u64 asan_violated;
int arena_fls(__u64 word);
void __arena *arena_malloc(size_t size);
void arena_free(void __arena *ptr);
/*
* The verifier associates arenas with programs by checking LD.IMM
* instruction operands for an arena and populating the program state
* with the first instance it finds. This requires accessing our global
* arena variable, but subprogs do not necessarily do so while still
* using pointers from that arena. Insert an LD.IMM instruction to
* access the arena and help the verifier.
*/
#define arena_subprog_init() do { asm volatile ("" :: "r"(&arena)); } while (0)
#else /* ! __BPF__ */
#include <stdint.h>
#define __arena
typedef uint8_t u8;
typedef uint16_t u16;
typedef uint32_t u32;
typedef uint64_t u64;
typedef int8_t s8;
typedef int16_t s16;
typedef int32_t s32;
typedef int64_t s64;
/* Dummy "definition" for userspace. */
#define arena_spinlock_t int
#endif /* __BPF__ */
struct arena_get_info_args {
void __arena *arena_base;
};
struct arena_alloc_reserve_args {
u64 nr_pages;
};
/* Reasonable default number of pages reserved by arena_alloc_reserve. */
#define ARENA_RESERVE_PAGES_DFL (8)
@@ -0,0 +1,83 @@
/* SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause */
#pragma once
#define RB_MAXLVL_PRINT (16)
struct rbnode;
struct rbnode {
struct rbnode __arena *parent;
union {
struct {
struct rbnode __arena *left;
struct rbnode __arena *right;
};
struct rbnode __arena *child[2];
};
uint64_t key;
/* Used as a linked list or to store KV pairs. */
union {
struct rbnode __arena *next;
uint64_t value;
};
bool is_red;
};
/*
* Does the rbtree allocate its own nodes, or do they get
* allocated by the caller?
*/
enum rbtree_alloc {
RB_ALLOC,
RB_NOALLOC,
};
/*
* Specify the behavior of rbtree insertions when the key is
* already present in the tree.
*
* RB_DEFAULT: Default behavior, reject the new insert.
*
* RB_UPDATE: Update the existing value in the rbtree.
* This updates the node itself, not just the value in
* the existing node.
*
* RB_DUPLICATE: Allow nodes with identical keys in the rbtree.
* Finding/popping/removing a key acts on any of the nodes
* with the appropriate key - there is no ordering by time
* of insertion.
*/
enum rbtree_insert_mode {
RB_DEFAULT,
RB_UPDATE,
RB_DUPLICATE,
};
struct rbtree {
struct rbnode __arena *root;
enum rbtree_alloc alloc;
enum rbtree_insert_mode insert;
};
#ifdef __BPF__
struct rbtree __arena *rb_create(enum rbtree_alloc alloc, enum rbtree_insert_mode insert);
int rb_destroy(struct rbtree __arena *rbtree);
int rb_insert(struct rbtree __arena *rbtree, u64 key, u64 value);
int rb_remove(struct rbtree __arena *rbtree, u64 key);
int rb_find(struct rbtree __arena *rbtree, u64 key, u64 *value);
int rb_print(struct rbtree __arena *rbtree);
int rb_least(struct rbtree __arena *rbtree, u64 *key, u64 *value);
int rb_pop(struct rbtree __arena *rbtree, u64 *key, u64 *value);
int rb_insert_node(struct rbtree __arena *rbtree, struct rbnode __arena *node);
int rb_remove_node(struct rbtree __arena *rbtree, struct rbnode __arena *node);
struct rbnode __arena *rb_node_alloc(u64 key, u64 value);
void rb_node_free(struct rbnode __arena *rbnode);
int rb_integrity_check(struct rbtree __arena *rbtree);
#endif /* __BPF__ */
@@ -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);
@@ -0,0 +1,138 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#pragma once
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/mman.h>
#include <bpf/libbpf.h>
#include <bpf/bpf.h>
static inline int libarena_run_prog(int prog_fd)
{
LIBBPF_OPTS(bpf_test_run_opts, opts);
int ret;
ret = bpf_prog_test_run_opts(prog_fd, &opts);
if (ret)
return ret;
return opts.retval;
}
static inline bool libarena_is_test_prog(const char *name)
{
return strstr(name, "test_") == name;
}
static inline bool libarena_is_asan_test_prog(const char *name)
{
return strstr(name, "asan_test") == name;
}
static inline bool libarena_is_parallel_test_prog(const char *name)
{
return strstr(name, "parallel_test") == name;
}
static inline int libarena_run_prog_args(int prog_fd, void *args, size_t argsize)
{
LIBBPF_OPTS(bpf_test_run_opts, opts);
int ret;
opts.ctx_in = args;
opts.ctx_size_in = argsize;
ret = bpf_prog_test_run_opts(prog_fd, &opts);
return ret ?: opts.retval;
}
static inline int libarena_get_arena_base(int arena_get_info_fd,
void **arena_base)
{
LIBBPF_OPTS(bpf_test_run_opts, opts);
struct arena_get_info_args args = { .arena_base = NULL };
int ret;
opts.ctx_in = &args;
opts.ctx_size_in = sizeof(args);
ret = bpf_prog_test_run_opts(arena_get_info_fd, &opts);
if (ret)
return ret;
if (opts.retval)
return opts.retval;
*arena_base = args.arena_base;
return 0;
}
static inline int libarena_get_globals_pages(int arena_get_globals_fd,
size_t arena_all_pages,
u64 *globals_pages)
{
size_t pgsize = sysconf(_SC_PAGESIZE);
void *arena_base;
ssize_t i;
u8 *vec;
int ret;
ret = libarena_get_arena_base(arena_get_globals_fd, &arena_base);
if (ret)
return ret;
if (!arena_base)
return -EINVAL;
vec = calloc(arena_all_pages, sizeof(*vec));
if (!vec)
return -ENOMEM;
if (mincore(arena_base, arena_all_pages * pgsize, vec) < 0) {
ret = -errno;
free(vec);
return ret;
}
*globals_pages = 0;
for (i = arena_all_pages - 1; i >= 0; i--) {
if (!(vec[i] & 0x1))
break;
*globals_pages += 1;
}
free(vec);
return 0;
}
static inline int libarena_asan_init(int arena_asan_init_fd,
int asan_init_fd,
size_t arena_all_pages)
{
LIBBPF_OPTS(bpf_test_run_opts, opts);
struct asan_init_args args;
u64 globals_pages;
int ret;
ret = libarena_get_globals_pages(arena_asan_init_fd,
arena_all_pages, &globals_pages);
if (ret)
return ret;
args = (struct asan_init_args){
.arena_all_pages = arena_all_pages,
.arena_globals_pages = globals_pages,
};
opts.ctx_in = &args;
opts.ctx_size_in = sizeof(args);
ret = bpf_prog_test_run_opts(asan_init_fd, &opts);
if (ret)
return ret;
return opts.retval;
}
@@ -0,0 +1,258 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <libarena/common.h>
#include <libarena/asan.h>
#include <libarena/buddy.h>
/* Required for parsing the ASAN call stacks. */
#include "test_progs_compat.h"
extern struct buddy __arena buddy;
#ifdef BPF_ARENA_ASAN
#include "st_asan_common.h"
static __always_inline int asan_test_buddy_oob_single(size_t alloc_size)
{
u8 __arena *mem;
int ret, i;
ret = asan_validate();
if (ret < 0)
return ret;
mem = buddy_alloc(&buddy, alloc_size);
if (!mem) {
arena_stdout("buddy_alloc failed for size %lu", alloc_size);
return -ENOMEM;
}
ret = asan_validate();
if (ret < 0)
return ret;
for (i = zero; i < alloc_size && can_loop; i++) {
mem[i] = 0xba;
ret = asan_validate_addr(false, &mem[i]);
if (ret < 0)
return ret;
}
mem[alloc_size] = 0xba;
ret = asan_validate_addr(true, &mem[alloc_size]);
if (ret < 0)
return ret;
buddy_free(&buddy, mem);
return 0;
}
/*
* Factored out because asan_validate_addr is complex enough to cause
* verification failures if verified with the rest of asan_test_buddy_uaf_single.
*/
__weak int asan_test_buddy_byte(u8 __arena *mem, int i, bool freed)
{
int ret;
/* The header in freed blocks doesn't get poisoned. */
if (freed && BUDDY_HEADER_OFF <= i &&
i < BUDDY_HEADER_OFF + sizeof(struct buddy_header))
return 0;
mem[i] = 0xba;
ret = asan_validate_addr(freed, &mem[i]);
if (ret < 0)
return ret;
return 0;
}
__weak int asan_test_buddy_uaf_single(size_t alloc_size)
{
u8 __arena *mem;
int ret;
int i;
mem = buddy_alloc(&buddy, alloc_size);
if (!mem) {
arena_stdout("buddy_alloc failed for size %lu", alloc_size);
return -ENOMEM;
}
ret = asan_validate();
if (ret < 0)
return ret;
for (i = zero; i < alloc_size && can_loop; i++) {
ret = asan_test_buddy_byte(mem, i, false);
if (ret)
return ret;
}
ret = asan_validate();
if (ret < 0)
return ret;
buddy_free(&buddy, mem);
for (i = zero; i < alloc_size && can_loop; i++) {
ret = asan_test_buddy_byte(mem, i, true);
if (ret)
return ret;
}
return 0;
}
struct buddy_blob {
volatile u8 mem[48];
u8 oob;
};
static __always_inline int asan_test_buddy_blob_single(void)
{
volatile struct buddy_blob __arena *blob;
const size_t alloc_size = sizeof(struct buddy_blob) - 1;
int ret;
blob = buddy_alloc(&buddy, alloc_size);
if (!blob)
return -ENOMEM;
blob->mem[0] = 0xba;
ret = asan_validate_addr(false, &blob->mem[0]);
if (ret < 0)
return ret;
blob->mem[47] = 0xba;
ret = asan_validate_addr(false, &blob->mem[47]);
if (ret < 0)
return ret;
blob->oob = 0;
ret = asan_validate_addr(true, &blob->oob);
if (ret < 0)
return ret;
buddy_free(&buddy, (void __arena *)blob);
return 0;
}
SEC("syscall")
__stderr("Memory violation for address {{.*}} for write of size 1")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
__weak int asan_test_buddy_oob(void)
{
size_t sizes[] = {
7, 8, 17, 18, 64, 256, 317, 512, 1024,
};
int ret, i;
ret = buddy_init(&buddy);
if (ret) {
arena_stdout("buddy_init failed with %d", ret);
return ret;
}
for (i = zero; i < sizeof(sizes) / sizeof(sizes[0]) && can_loop; i++) {
ret = asan_test_buddy_oob_single(sizes[i]);
if (ret) {
arena_stdout("%s:%d Failed for size %lu", __func__,
__LINE__, sizes[i]);
buddy_destroy(&buddy);
return ret;
}
}
buddy_destroy(&buddy);
ret = asan_validate();
if (ret < 0)
return ret;
return 0;
}
SEC("syscall")
__stderr("Memory violation for address {{.*}} for write of size 1")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
__weak int asan_test_buddy_uaf(void)
{
size_t sizes[] = { 16, 32, 64, 128, 256, 512, 1024, 16384 };
int ret, i;
ret = buddy_init(&buddy);
if (ret) {
arena_stdout("buddy_init failed with %d", ret);
return ret;
}
for (i = zero; i < sizeof(sizes) / sizeof(sizes[0]) && can_loop; i++) {
ret = asan_test_buddy_uaf_single(sizes[i]);
if (ret) {
arena_stdout("%s:%d Failed for size %lu", __func__,
__LINE__, sizes[i]);
buddy_destroy(&buddy);
return ret;
}
}
buddy_destroy(&buddy);
ret = asan_validate();
if (ret < 0)
return ret;
return 0;
}
SEC("syscall")
__stderr("Memory violation for address {{.*}} for write of size 1")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
__weak int asan_test_buddy_blob(void)
{
const int iters = 10;
int ret, i;
ret = buddy_init(&buddy);
if (ret) {
arena_stdout("buddy_init failed with %d", ret);
return ret;
}
for (i = zero; i < iters && can_loop; i++) {
ret = asan_test_buddy_blob_single();
if (ret) {
arena_stdout("%s:%d Failed on iteration %d", __func__,
__LINE__, i);
buddy_destroy(&buddy);
return ret;
}
}
buddy_destroy(&buddy);
ret = asan_validate();
if (ret < 0)
return ret;
return 0;
}
#endif
__weak char _license[] SEC("license") = "GPL";
@@ -0,0 +1,52 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#pragma once
#define ST_PAGES 64
static inline void print_asan_map_state(void __arena *addr)
{
arena_stdout("%s:%d ASAN %p -> (val: %x gran: %x set: [%s])",
__func__, __LINE__, addr,
*(s8 __arena *)(addr), ASAN_GRANULE(addr),
asan_shadow_set(addr) ? "yes" : "no");
}
/*
* Emit an error and force the current function to exit if the ASAN
* violation state is unexpected. Reset the violation state after.
*/
static inline int asan_validate_addr(bool cond, void __arena *addr)
{
if ((asan_violated != 0) == cond) {
asan_violated = 0;
return 0;
}
arena_stdout("%s:%d ASAN asan_violated %lx", __func__, __LINE__,
(u64)asan_violated);
print_asan_map_state(addr);
asan_violated = 0;
return -EINVAL;
}
static inline int asan_validate(void)
{
if (!asan_violated)
return 0;
arena_stdout("%s:%d Found ASAN violation at %lx", __func__, __LINE__,
asan_violated);
asan_violated = 0;
return -EINVAL;
}
struct blob {
volatile u8 mem[59];
u8 oob;
};
@@ -0,0 +1,209 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <libarena/common.h>
#include <libarena/asan.h>
#include <libarena/buddy.h>
extern struct buddy __arena buddy;
struct segarr_entry {
u8 __arena *block;
size_t sz;
u8 poison;
};
#define SEGARRLEN (512)
static struct segarr_entry __arena segarr[SEGARRLEN];
static void __arena *ptrs[17];
size_t __arena alloc_sizes[] = { 3, 17, 1025, 129, 16350, 333, 9, 517 };
size_t __arena alloc_multiple_sizes[] = { 3, 17, 1025, 129, 16350, 333, 9, 517, 2099 };
size_t __arena alloc_free_sizes[] = { 3, 17, 64, 129, 256, 333, 512, 517 };
size_t __arena alignment_sizes[] = { 1, 3, 7, 8, 9, 15, 16, 17, 31,
32, 64, 100, 128, 255, 256, 512, 1000 };
SEC("syscall")
__weak int test_buddy_create(void)
{
const int iters = 10;
int ret, i;
for (i = zero; i < iters && can_loop; i++) {
ret = buddy_init(&buddy);
if (ret)
return ret;
ret = buddy_destroy(&buddy);
if (ret)
return ret;
}
return 0;
}
SEC("syscall")
__weak int test_buddy_alloc(void)
{
void __arena *mem;
int ret, i;
for (i = zero; i < 8 && can_loop; i++) {
ret = buddy_init(&buddy);
if (ret)
return ret;
mem = buddy_alloc(&buddy, alloc_sizes[i]);
if (!mem) {
buddy_destroy(&buddy);
return -ENOMEM;
}
buddy_destroy(&buddy);
}
return 0;
}
SEC("syscall")
__weak int test_buddy_alloc_free(void)
{
const int iters = 800;
void __arena *mem;
int ret, i;
ret = buddy_init(&buddy);
if (ret)
return ret;
for (i = zero; i < iters && can_loop; i++) {
mem = buddy_alloc(&buddy, alloc_free_sizes[(i * 5) % 8]);
if (!mem) {
buddy_destroy(&buddy);
return -ENOMEM;
}
buddy_free(&buddy, mem);
}
buddy_destroy(&buddy);
return 0;
}
SEC("syscall")
__weak int test_buddy_alloc_multiple(void)
{
int ret, j;
u32 i, idx;
u8 __arena *mem;
size_t sz;
u8 poison;
ret = buddy_init(&buddy);
if (ret)
return ret;
/*
* Cycle through each size, allocating an entry in the
* segarr. Continue for SEGARRLEN iterations. For every
* allocation write down the size, use the current index
* as a poison value, and log it with the pointer in the
* segarr entry. Use the poison value to poison the entire
* allocated memory according to the size given.
*/
for (i = zero; i < SEGARRLEN && can_loop; i++) {
sz = alloc_multiple_sizes[i % 9];
poison = (u8)i;
mem = buddy_alloc(&buddy, sz);
if (!mem) {
buddy_destroy(&buddy);
arena_stdout("%s:%d", __func__, __LINE__);
return -ENOMEM;
}
segarr[i].block = mem;
segarr[i].sz = sz;
segarr[i].poison = poison;
for (j = zero; j < sz && can_loop; j++) {
mem[j] = poison;
if (mem[j] != poison) {
buddy_destroy(&buddy);
return -EINVAL;
}
}
}
/*
* Go to (i * 17) % SEGARRLEN, and free the block pointed to.
* Before freeing, check all bytes have the poisoned value
* corresponding to the element. If any values are unexpected,
* return an error. Skip some elements to test destroying the
* buddy allocator while data is still allocated.
*/
for (i = 10; i < SEGARRLEN && can_loop; i++) {
idx = (i * 17) % SEGARRLEN;
mem = segarr[idx].block;
sz = segarr[idx].sz;
poison = segarr[idx].poison;
for (j = zero; j < sz && can_loop; j++) {
if (mem[j] != poison) {
buddy_destroy(&buddy);
arena_stdout("%s:%d %lx %u vs %u", __func__,
__LINE__, (uintptr_t)&mem[j],
mem[j], poison);
return -EINVAL;
}
}
buddy_free(&buddy, mem);
}
buddy_destroy(&buddy);
return 0;
}
SEC("syscall")
__weak int test_buddy_alignment(void)
{
int ret, i;
ret = buddy_init(&buddy);
if (ret)
return ret;
/* Allocate various sizes and check alignment */
for (i = zero; i < 17 && can_loop; i++) {
ptrs[i] = buddy_alloc(&buddy, alignment_sizes[i]);
if (!ptrs[i]) {
arena_stdout("alignment test: alloc failed for size %lu",
alignment_sizes[i]);
buddy_destroy(&buddy);
return -ENOMEM;
}
/* Check 8-byte alignment */
if ((u64)ptrs[i] & 0x7) {
arena_stdout(
"alignment test: ptr %llx not 8-byte aligned (size %lu)",
(u64)ptrs[i], alignment_sizes[i]);
buddy_destroy(&buddy);
return -EINVAL;
}
}
/* Free all allocations */
for (i = zero; i < 17 && can_loop; i++)
buddy_free(&buddy, ptrs[i]);
buddy_destroy(&buddy);
return 0;
}
__weak char _license[] SEC("license") = "GPL";
@@ -0,0 +1,669 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
#include <bpf_atomic.h>
#include <libarena/common.h>
#include <libarena/asan.h>
#include <libarena/spmc.h>
#define TEST_SPMC_THREADS 3
#define TEST_SPMC_STEALERS (TEST_SPMC_THREADS - 1)
/*
* The test requires the stealers/owners to sometimes quiesce
* before continuing the benchmark. Normally we'd use something
* like a condition variable, but since the benchmark is short-lived
* and operations are wait-free we just spin around the quiescence
* point instead. If we time out, we just fail the benchmark.
*/
#define TEST_SPMC_SYNC_SPINS BPF_MAX_LOOPS
/*
* We track all the values we retrieve from the queue
* to get some guarantee we're, not corrupting data,
* e.g., accidentally reusing a past value from a slot.
*/
#define TEST_SPMC_MAX_VALUES (1024)
static u64 __arena seen[TEST_SPMC_MAX_VALUES];
/* The single spmc queue for the benchmark. */
static struct spmc __arena *spmc;
/* Owner and stealer epochs. We define the , */
static volatile u64 owner_epoch;
static volatile u64 stealer_epoch;
/* Map owner epochs to stealer epochs (simply scale by # of stealers). */
#define STEALER_EPOCH(owner_epoch) ((owner_epoch) * TEST_SPMC_STEALERS)
/* Global abort switch. If any thread fails, all others exit ASAP. */
static volatile bool test_abort;
/*
* Counters useful for ensuring conservation of pushes/pops of unique values
* (we're not stealing/popping more/fewer items than were pushed).
*/
static volatile u64 expected_total;
static volatile u64 total_seen;
/* Measure how many pops and steals we've made (irrespective of retrieved value). */
static volatile u64 pops;
static volatile u64 steals;
/* Used for the resize selftest, see below. */
static volatile u64 stealers_started;
/* Used for the mixed selftest, see below. */
static volatile u64 round_steals;
/*
* We have multiple stealers and a single owner. We sometimes want the owner
* to successfully outproduce the stealers, we add a busy loop in them.
*/
#define TEST_SPMC_WASTE_ROUNDS (1UL << 12)
/*
* The spmc data structure depends on the runtime fully
* supporting acquire/release semantics, which is not
* the case for all architectures.
*/
#if defined(ENABLE_ATOMICS_TESTS) && \
(defined(__TARGET_ARCH_arm64) || defined(__TARGET_ARCH_x86) || \
(defined(__TARGET_ARCH_riscv) && __riscv_xlen == 64))
static bool spmc_tests_enabled(void)
{
return true;
}
#else
static bool spmc_tests_enabled(void)
{
return false;
}
#endif
/*
* Scaffolding for each parallel test. Each test has setup/teardown,
* a single owner thread that owns the queue, and TEST_SPMC_STEALER
* threads that try to steal.
*/
#define DEFINE_PARALLEL_SPMC_TEST(prefix, expected_total) \
SEC("syscall") int parallel_test_spmc_##prefix##__enabled(void) \
{ \
return spmc_tests_enabled() ? 0 : -EOPNOTSUPP; \
} \
SEC("syscall") int parallel_test_spmc_##prefix##__init(void) \
{ \
return spmc_common_init(expected_total); \
} \
SEC("syscall") int parallel_test_spmc_##prefix##__fini(void) \
{ \
return spmc_common_fini(); \
} \
SEC("syscall") int parallel_test_spmc_##prefix##__0(void) \
{ \
return spmc_##prefix##_owner(); \
} \
SEC("syscall") int parallel_test_spmc_##prefix##__1(void) \
{ \
return spmc_##prefix##_stealer(); \
} \
SEC("syscall") int parallel_test_spmc_##prefix##__2(void) \
{ \
return spmc_##prefix##_stealer(); \
} \
static int spmc_common_init(u64 total)
{
u64 i;
if (total > TEST_SPMC_MAX_VALUES)
return -E2BIG;
owner_epoch = 0;
stealer_epoch = 0;
test_abort = false;
expected_total = total;
total_seen = 0;
pops = 0;
steals = 0;
stealers_started = 0;
round_steals = 0;
for (i = zero; i < TEST_SPMC_MAX_VALUES && can_loop; i++)
seen[i] = 0;
spmc = spmc_create();
if (!spmc)
return -ENOMEM;
return 0;
}
static int spmc_common_fini(void)
{
int ret;
ret = spmc_destroy(spmc);
spmc = NULL;
return ret;
}
__weak
int spmc_quiesce_on_owner(u64 epoch)
{
u64 i;
bpf_for(i, 0, TEST_SPMC_SYNC_SPINS) {
if (test_abort)
return -EINTR;
if (smp_load_acquire(&owner_epoch) >= epoch)
return 0;
}
test_abort = true;
return -ETIMEDOUT;
}
__weak
int spmc_quiesce_on_stealer(u64 epoch)
{
u64 target, cur;
unsigned int i;
int err = -ETIMEDOUT;
target = STEALER_EPOCH(epoch);
bpf_for(i, 0, TEST_SPMC_SYNC_SPINS) {
if (test_abort) {
err = -EINTR;
break;
}
cur = smp_load_acquire(&stealer_epoch);
if (cur > target) {
err = -EINVAL;
test_abort = true;
break;
}
if (cur == target)
return 0;
}
test_abort = true;
return err;
}
static int spmc_update_stats(u64 val, bool owner)
{
u64 total;
total = expected_total;
if (val >= total || val >= TEST_SPMC_MAX_VALUES) {
test_abort = true;
return -EINVAL;
}
if (__sync_fetch_and_add(&seen[val], 1) != 0) {
test_abort = true;
return -EINVAL;
}
__sync_fetch_and_add(&total_seen, 1);
if (owner)
__sync_fetch_and_add(&pops, 1);
else
__sync_fetch_and_add(&steals, 1);
return 0;
}
static int spmc_validate_owner_empty(void)
{
u64 val;
int ret;
ret = spmc_owned_remove(spmc, &val);
if (ret != -ENOENT) {
test_abort = true;
/* Change a 0 return value into -EINVAL. */
return ret ?: -EINVAL;
}
return 0;
}
__weak
int spmc_validate_all_seen(void)
{
u64 i, total;
total = expected_total;
if (total_seen != total)
goto err;
if (pops + steals != total)
goto err;
for (i = zero; i < total && can_loop; i++) {
if (seen[i % TEST_SPMC_MAX_VALUES] != 1)
goto err;
}
return 0;
err:
test_abort = true;
return -EINVAL;
}
/*
* Single value benchmark. The owner adds an item then races with
* the stealers for it. This way directly race between owner and
* stealers on the same slot.
*/
#define TEST_SPMC_SINGLEVAL_ITERS (64)
__weak
int spmc_singleval_tryconsume(u64 expected, bool steal)
{
u64 val;
int ret;
while (can_loop) {
if (steal)
ret = spmc_steal(spmc, &val);
else
ret = spmc_owned_remove(spmc, &val);
/* Success. Update and validate. */
if (!ret) {
if (val != expected)
return -EINVAL;
ret = spmc_update_stats(val, !steal);
if (ret)
return ret;
return 0;
}
/*
* If we got -ENOENT, the queue is empty
* and we're good to go.
*/
if (ret != -EAGAIN)
return (ret == -ENOENT) ? 0 : ret;
}
/* Impossible. */
return -EINVAL;
}
static int spmc_singleval_owner(void)
{
int ret;
u64 i;
for (i = zero; i < TEST_SPMC_SINGLEVAL_ITERS && can_loop; i++) {
ret = spmc_quiesce_on_stealer(i);
if (ret)
goto err;
ret = spmc_owned_add(spmc, i);
if (ret)
goto err;
__sync_fetch_and_add(&owner_epoch, 1);
ret = spmc_singleval_tryconsume(i, false);
if (ret)
goto err;
ret = spmc_quiesce_on_stealer(i + 1);
if (ret)
goto err;
}
ret = spmc_validate_owner_empty();
if (ret)
return ret;
return spmc_validate_all_seen();
err:
test_abort = true;
return -EINVAL;
}
static int spmc_singleval_stealer(void)
{
int ret;
u64 i;
for (i = zero; i < TEST_SPMC_SINGLEVAL_ITERS && can_loop; i++) {
ret = spmc_quiesce_on_owner(i + 1);
if (ret)
goto err;
ret = spmc_singleval_tryconsume(i, true);
if (ret)
goto err;
__sync_fetch_and_add(&stealer_epoch, 1);
}
return 0;
err:
test_abort = true;
return -EINVAL;
}
DEFINE_PARALLEL_SPMC_TEST(singleval, TEST_SPMC_SINGLEVAL_ITERS)
/*
* The resize test. Force a resize from the owner even while the stealers
* are trying to consume. Then make sure the queue is still consistent
* after the resize.
*
* The owner _doesn't_ consume from the queue. The test makes sure that
* switching the array from underneath the stealers works.
*/
/* Force 2 resizes (since the rate of resize is logarithmic). */
#define TEST_SPMC_RESIZE_ORDER (2)
#define TEST_SPMC_RESIZE_PREFILL ((SPMC_ARR_BASESZ << TEST_SPMC_RESIZE_ORDER) - 1)
/* */
#define TEST_SPMC_RESIZE_TAIL (SPMC_ARR_BASESZ << TEST_SPMC_RESIZE_ORDER)
#define TEST_SPMC_RESIZE_TOTAL (TEST_SPMC_RESIZE_PREFILL + TEST_SPMC_RESIZE_TAIL)
__weak
int spmc_wait_for_stealers_to_start(u64 target)
{
u64 i;
bpf_for(i, 0, TEST_SPMC_SYNC_SPINS) {
if (test_abort)
return -EINTR;
if (READ_ONCE(stealers_started) >= target)
return 0;
}
test_abort = true;
return -ETIMEDOUT;
}
__weak
void spmc_waste_time(void)
{
int i;
int j;
for (i = zero; i < TEST_SPMC_WASTE_ROUNDS && can_loop; i++) {
/* Random computation. */
WRITE_ONCE(j, i * 17 + 23);
}
}
static int spmc_resize_owner(void)
{
bool resized = false;
u64 i;
int ret;
/* Get a head start vs the consumers. */
for (i = zero; i < TEST_SPMC_RESIZE_PREFILL && can_loop; i++) {
ret = spmc_owned_add(spmc, i);
if (ret) {
test_abort = true;
return ret;
}
}
__sync_fetch_and_add(&owner_epoch, 1);
/* Wait for stealers to start then start racing. */
ret = spmc_wait_for_stealers_to_start(TEST_SPMC_STEALERS);
if (ret)
return ret;
for (i = TEST_SPMC_RESIZE_PREFILL; i < TEST_SPMC_RESIZE_TOTAL && can_loop; i++) {
ret = spmc_owned_add(spmc, i);
if (ret) {
test_abort = true;
return ret;
}
if (spmc->cur->order > TEST_SPMC_RESIZE_ORDER)
resized = true;
}
/* Did we get to resize while racing? */
if (!resized) {
test_abort = true;
return -EINVAL;
}
/*
* Wait for the stealers to drain and make sure
* we didn't lose any items along the way.
*/
__sync_fetch_and_add(&owner_epoch, 1);
ret = spmc_quiesce_on_stealer(1);
if (ret)
return ret;
ret = spmc_validate_owner_empty();
if (ret)
return ret;
return spmc_validate_all_seen();
}
static int spmc_resize_stealer(void)
{
bool owner_done = false;
u64 val;
int ret;
arena_subprog_init();
ret = spmc_quiesce_on_owner(1);
if (ret)
return ret;
__sync_fetch_and_add(&stealers_started, 1);
while (can_loop) {
spmc_waste_time();
if (test_abort)
return -EINTR;
ret = spmc_steal(spmc, &val);
if (!ret) {
ret = spmc_update_stats(val, false);
if (ret)
return ret;
continue;
}
if (ret == -EAGAIN)
continue;
if (ret == -ENOENT) {
if (owner_done)
break;
owner_done = owner_epoch >= 2;
continue;
}
test_abort = true;
return ret;
}
__sync_fetch_and_add(&stealer_epoch, 1);
return 0;
}
DEFINE_PARALLEL_SPMC_TEST(resize, TEST_SPMC_RESIZE_TOTAL)
/*
* The burst benchmark. The owner generates data all at once,
* then waits for the stealers to steal half then starts removing
* items until the queue empties. The owner also makes sure the
* item order is not jumbled.
*/
#define TEST_SPMC_BURST_ROUNDS (4)
#define TEST_SPMC_BURST_BURST (64)
#define TEST_SPMC_BURST_TOTAL (TEST_SPMC_BURST_ROUNDS * TEST_SPMC_BURST_BURST)
#define TEST_SPMC_BURST_STEAL_TARGET (TEST_SPMC_BURST_BURST / 2)
static int spmc_wait_for_round_steals(u64 target)
{
u64 i;
arena_subprog_init();
bpf_for(i, 0, TEST_SPMC_SYNC_SPINS) {
if (test_abort)
return -EINTR;
if (round_steals >= target)
return 0;
}
test_abort = true;
return -ETIMEDOUT;
}
__weak int
spmc_burst_owner_round(u64 round)
{
u64 i, base, stolen, expected, val;
int ret;
base = round * TEST_SPMC_BURST_BURST;
round_steals = 0;
for (i = zero; i < TEST_SPMC_BURST_BURST && can_loop; i++) {
ret = spmc_owned_add(spmc, base + i);
if (ret)
return ret;
}
__sync_fetch_and_add(&owner_epoch, 1);
ret = spmc_wait_for_round_steals(TEST_SPMC_BURST_STEAL_TARGET);
if (ret == -EINTR || ret == -ETIMEDOUT)
return ret;
__sync_fetch_and_add(&owner_epoch, 1);
ret = spmc_quiesce_on_stealer(round + 1);
if (ret)
return ret;
stolen = round_steals;
if (stolen > TEST_SPMC_BURST_BURST)
return -EINVAL;
for (i = zero; i < TEST_SPMC_BURST_BURST - stolen && can_loop; i++) {
ret = spmc_owned_remove(spmc, &val);
if (ret)
return ret;
expected = base + TEST_SPMC_BURST_BURST - 1 - i;
if (val != expected)
return -EINVAL;
ret = spmc_update_stats(val, true);
if (ret) {
test_abort = true;
return -EINVAL;
}
}
ret = spmc_validate_owner_empty();
if (ret)
return ret;
return 0;
}
static int spmc_burst_owner(void)
{
u64 round;
int ret;
arena_subprog_init();
for (round = zero; round < TEST_SPMC_BURST_ROUNDS && can_loop; round++) {
ret = spmc_burst_owner_round(round);
if (ret)
goto err;
}
return spmc_validate_all_seen();
err:
test_abort = true;
return -EINVAL;
}
static int spmc_burst_stealer(void)
{
u64 round, val, active_epoch;
int ret;
arena_subprog_init();
for (round = zero; round < TEST_SPMC_BURST_ROUNDS && can_loop; round++) {
active_epoch = round * 2 + 1;
/*
* Wait till the owner prefills the queue then
* start stealing.
*/
ret = spmc_quiesce_on_owner(active_epoch);
if (ret)
return ret;
while (owner_epoch == active_epoch && can_loop) {
if (test_abort)
return -EINTR;
ret = spmc_steal(spmc, &val);
if (!ret) {
ret = spmc_update_stats(val, false);
if (ret)
return ret;
__sync_fetch_and_add(&round_steals, 1);
continue;
}
if (ret == -EAGAIN || ret == -ENOENT)
continue;
test_abort = true;
return ret;
}
__sync_fetch_and_add(&stealer_epoch, 1);
}
return 0;
}
DEFINE_PARALLEL_SPMC_TEST(burst, TEST_SPMC_BURST_TOTAL)
@@ -0,0 +1,15 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#pragma once
#ifdef __BPF__
/* Selftests use these tags for compatibility with test_progs. */
#define __test_tag(tag) __attribute__((btf_decl_tag("comment:" XSTR(__COUNTER__) ":" tag)))
#define __stderr(msg) __test_tag("test_expect_stderr=" msg)
#define __stderr_unpriv(msg) __test_tag("test_expect_stderr_unpriv=" msg)
#define XSTR(s) STR(s)
#define STR(s) #s
#endif
@@ -0,0 +1,968 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
#include <libarena/common.h>
#include <libarena/asan.h>
#include <libarena/rbtree.h>
typedef struct node_ctx __arena *node_ctx;
struct node_ctx {
struct rbnode rbnode;
node_ctx next;
};
static const u64 keys[] = { 51, 43, 37, 3, 301, 46, 383, 990, 776, 729, 871, 96, 189, 213,
376, 167, 131, 939, 626, 119, 374, 700, 772, 154, 883, 620, 641, 5,
428, 516, 105, 622, 988, 811, 931, 973, 246, 690, 934, 744, 210, 311,
32, 255, 960, 830, 523, 429, 541, 738, 705, 774, 715, 446, 98, 578,
777, 191, 279, 91, 767 };
static const u64 morekeys[] = { 173, 636, 1201, 8642, 5957, 3617, 4586, 8053, 6551, 7592, 1748, 1589, 8644, 9918, 6977,
4448, 5852, 4640, 9717, 2303, 7424, 7695, 2334, 8876, 8618, 5745, 7134, 2178, 5280, 2140, 1138,
5083, 8922, 1516, 2437, 2488, 4307, 4329, 5088, 8456, 5938, 1441, 1684, 5750, 721, 1107, 2089,
9737, 4687, 5016, 4849, 8193, 9603, 9147, 5992, 166, 6721, 812, 4144, 6237, 6509, 3466, 9255,
7767, 3960, 6759, 2968, 6046, 9784, 8395, 2619, 1711, 528, 6424, 9084, 3179, 1342, 5676, 9445,
5691, 6678, 8487, 1627, 998, 6178, 2229, 1987, 3319, 572, 169, 2161, 3018, 5439, 7287, 7265, 5995,
5003, 5857, 2836, 5634, 4735, 9261, 8287, 5359, 533, 1406, 9573, 4026, 714, 3956, 1722, 6395,
9648, 3887, 7185, 470, 4482, 4997, 841, 8913, 9946, 3999, 9357, 9847, 277, 8184, 8704, 6766, 3323,
5468, 8638, 7905, 8858, 6142, 3685, 3452, 4689, 8878, 8836, 158, 831, 7914, 3031, 8374, 4921,
4207, 3460, 5547, 3358, 1083, 4619, 7818, 2962, 4879, 4583, 2172, 8819, 9830, 1194, 2666, 9812,
5704, 8432, 5916, 6007, 6609, 4791, 1985, 3226, 2478, 9605, 5236, 8079, 3042, 1965, 3539, 9704,
4267, 6416, 760, 9968, 2983, 1190, 1964, 3211, 2870, 3106, 2794, 1542, 6916, 5986, 9096, 441,
5894, 8353, 7765, 3757, 5732, 88, 3091, 5637, 6042, 8447, 4073, 6923, 5491, 7010, 3663, 5029,
6162, 822, 4874, 7491, 5100, 3461, 6983, 2170, 1458, 1856, 648, 6272, 4887, 976, 2369, 5909, 4274,
3324, 6968, 2312, 2271, 8891, 6268, 6581, 1610, 8880, 6194, 6144, 9764, 6915, 829, 3774, 2265,
1752, 1314, 6377, 8760, 8004, 501, 4912, 9278, 1425, 9578, 7337, 307, 1885, 3151, 9617, 1647,
2458, 3702, 6091, 8902, 5663, 9378, 7640, 3336, 557, 1644, 6848, 1559, 8821, 266, 4330, 9790,
5920, 4222, 1143, 6248, 5792, 4847, 9726, 6303, 821, 6839, 6062, 7133, 3649, 9888, 2528, 1966,
5456, 4914, 3615, 1543, 3206, 3353, 6097, 2800, 1424, 9094, 7920, 7243, 1394, 5464, 1707, 576,
6524, 4261, 4187, 7889, 5336, 3377, 2921, 7244, 2766, 6584, 5514, 1387, 2957, 2258, 1077, 9979,
1128, 876, 4056, 4668, 4532, 1982, 7093, 4184, 5460, 7588, 4704, 6717, 61, 3959, 1826, 2294, 18,
8170, 9394, 8796, 7288, 7285, 7143, 148, 6676, 6603, 1051, 8225, 4169, 3230, 7697, 6971, 3454,
7501, 9514, 394, 2339, 4993, 5606, 6060, 1297, 8273, 3012, 157, 8181, 6765, 7207, 1005, 8833, 1914,
7456, 1846, 8375, 2741, 2074, 1712, 5286 };
SEC("syscall")
__weak int test_rbtree_find_nonexistent(void)
{
u64 key = 0xdeadbeef;
u64 value = 0;
int ret;
struct rbtree __arena *rbtree;
rbtree = rb_create(RB_ALLOC, RB_DEFAULT);
if (!rbtree)
return 1;
/* Should return -EINVAL */
ret = rb_find(rbtree, key, &value);
if (!ret)
return 2;
return rb_destroy(rbtree);
}
SEC("syscall")
__weak int test_rbtree_insert_existing(void)
{
u64 key = 525252;
u64 value = 24;
int ret;
struct rbtree __arena *rbtree;
rbtree = rb_create(RB_ALLOC, RB_DEFAULT);
if (!rbtree)
return 1;
ret = rb_insert(rbtree, key, value);
if (ret)
return 2;
/* Should return -EALREADY. */
ret = rb_insert(rbtree, key, value);
if (ret != -EALREADY) {
return 3;
}
return rb_destroy(rbtree);
}
SEC("syscall")
__weak int test_rbtree_update_existing(void)
{
u64 key = 33333;
u64 value;
int ret;
struct rbtree __arena *rbtree;
rbtree = rb_create(RB_ALLOC, RB_UPDATE);
if (!rbtree)
return 1;
value = 52;
ret = rb_insert(rbtree, key, value);
if (ret)
return 2;
ret = rb_find(rbtree, key, &value);
if (ret)
return 3;
if (value != 52)
return 4;
value = 65;
/* Should succeed. */
ret = rb_insert(rbtree, key, value);
if (ret)
return 5;
/* Should be updated. */
ret = rb_find(rbtree, key, &value);
if (ret)
return 6;
if (value != 65)
return 7;
return rb_destroy(rbtree);
}
SEC("syscall")
__weak int test_rbtree_insert_one(void)
{
u64 key = 202020;
u64 value = 0xbadcafe;
int ret;
struct rbtree __arena *rbtree;
rbtree = rb_create(RB_ALLOC, RB_UPDATE);
if (!rbtree)
return 1;
ret = rb_insert(rbtree, key, value);
if (ret)
return 2;
ret = rb_find(rbtree, key, &value);
if (ret)
return 3;
if (value != 0xbadcafe)
return 4;
return rb_destroy(rbtree);
}
SEC("syscall")
__weak int test_rbtree_insert_ten(void)
{
u64 key, value;
int ret, i;
struct rbtree __arena *rbtree;
rbtree = rb_create(RB_ALLOC, RB_UPDATE);
if (!rbtree)
return 1;
for (i = 0; i < 10 && can_loop; i++) {
key = keys[i];
ret = rb_insert(rbtree, key, 2 * key);
if (ret)
return 2 + 3 * i;
/* Read it back. */
ret = rb_find(rbtree, key, &value);
if (ret)
return 2 + 3 * i + 1;
if (value != 2 * key)
return 2 + 3 * i + 2;
}
/* Go find all inserted pairs. */
for (i = 0; i < 10 && can_loop; i++) {
key = keys[i];
ret = rb_find(rbtree, key, &value);
if (ret)
return 35 + 2 * i;
if (value != 2 * key)
return 35 + 2 * i + 1;
}
return rb_destroy(rbtree);
}
SEC("syscall")
__weak int test_rbtree_duplicate(void)
{
u64 key = 0x121212;
u64 value;
int ret, i;
struct rbtree __arena *rbtree;
rbtree = rb_create(RB_ALLOC, RB_DUPLICATE);
if (!rbtree)
return 1;
for (i = 0; i < 10 && can_loop; i++) {
ret = rb_insert(rbtree, key, 2 * key);
if (ret)
return 2 + 3 * i;
/* Read it back. */
ret = rb_find(rbtree, key, &value);
if (ret)
return 2 + 3 * i + 1;
if (value != 2 * key)
return 2 + 3 * i + 2;
}
/* Go find all inserted copies and remove them. */
for (i = 0; i < 10 && can_loop; i++) {
ret = rb_find(rbtree, key, &value);
if (ret) {
rb_print(rbtree);
return 35 + 3 * i;
}
if (value != 2 * key)
return 35 + 3 * i + 1;
ret = rb_remove(rbtree, key);
if (ret)
return 35 + 3 * i + 2;
}
return rb_destroy(rbtree);
}
static inline int
clean_up_noalloc_tree(struct rbtree __arena *rbtree)
{
node_ctx nodec;
int ret;
if (rbtree->alloc != RB_NOALLOC)
return -EINVAL;
/* Can't destroy an RB_NOALLOC tree that still has nodes. */
if (rb_destroy(rbtree) != -EBUSY)
return -EINVAL;
while (rbtree->root && can_loop) {
nodec = (node_ctx)arena_container_of(rbtree->root, struct node_ctx, rbnode);
ret = rb_remove_node(rbtree, &nodec->rbnode);
if (ret)
return ret;
arena_free(nodec);
}
return 0;
}
int insert_many(enum rbtree_alloc alloc, enum rbtree_insert_mode insert)
{
const size_t numkeys = sizeof(keys) / sizeof(keys[0]);
node_ctx nodec;
u64 key, value;
int ret;
int i;
struct rbtree __arena *rbtree;
rbtree = rb_create(alloc, insert);
if (!rbtree)
return 1;
for (i = 0; i < numkeys && can_loop; i++) {
key = keys[i];
if (rbtree->alloc != RB_ALLOC) {
nodec = arena_malloc(sizeof(*nodec));
if (!nodec) {
arena_stderr("out of memory\n");
return -ENOMEM;
}
nodec->rbnode.key = key;
nodec->rbnode.value = 2 * key;
ret = rb_insert_node(rbtree, &nodec->rbnode);
} else {
ret = rb_insert(rbtree, key, 2 * key);
}
if (ret)
return 2 + 3 * i;
/* Read it back. */
ret = rb_find(rbtree, key, &value);
if (ret)
return 2 + 3 * i + 1;
if (value != 2 * key)
return 2 + 3 * i + 2;
}
/* Go find all inserted pairs. */
for (i = 0; i < numkeys && can_loop; i++) {
key = keys[i];
ret = rb_find(rbtree, key, &value);
if (ret)
return 302 + 2 * i;
if (value != 2 * key)
return 302 + 2 * i + 1;
}
/* RB_ALLOC trees are destroyed while still having elements. */
if (rbtree->alloc == RB_ALLOC)
return rb_destroy(rbtree);
/* Otherwise manually clean up the tree. */
if (clean_up_noalloc_tree(rbtree))
return 5;
return rb_destroy(rbtree);
}
SEC("syscall")
__weak int test_rbtree_remove_one(void)
{
u64 key = 20, value = 5, newvalue;
int ret;
struct rbtree __arena *rbtree;
rbtree = rb_create(RB_ALLOC, RB_DEFAULT);
if (!rbtree)
return 1;
ret = rb_find(rbtree, key, &newvalue);
if (!ret)
return 2;
ret = rb_insert(rbtree, key, value);
if (ret)
return 3;
ret = rb_find(rbtree, key, &newvalue);
if (ret || value != newvalue)
return 4;
ret = rb_remove(rbtree, key);
if (ret)
return 5;
ret = rb_find(rbtree, key, &newvalue);
if (!ret)
return 6;
return rb_destroy(rbtree);
}
static __always_inline int remove_many_verify_all_present(struct rbtree __arena *rbtree)
{
const size_t numkeys = sizeof(morekeys) / sizeof(morekeys[0]);
u64 value;
int ret;
int i;
for (i = 0; i < numkeys && can_loop; i++) {
u64 key = morekeys[i];
ret = rb_find(rbtree, key, &value);
if (ret)
return -1;
if (value != 2 * key)
return -1;
}
return 0;
}
static __always_inline int remove_many_verify_remaining(struct rbtree __arena *rbtree)
{
const size_t numkeys = sizeof(morekeys) / sizeof(morekeys[0]);
u64 value;
int ret;
int i;
for (i = 0; i < numkeys && can_loop; i += 2) {
u64 key = morekeys[i];
ret = rb_find(rbtree, key, &value);
if (!ret)
return -1;
if (i + 1 >= numkeys)
break;
key = morekeys[i + 1];
ret = rb_find(rbtree, key, &value);
if (ret)
return -1;
if (value != 2 * key)
return -1;
}
for (i = 1; i < numkeys && can_loop; i += 2) {
u64 key = morekeys[i];
ret = rb_find(rbtree, key, &value);
if (ret)
return -1;
if (value != 2 * key)
return -1;
}
return 0;
}
static __noinline int remove_many_alloc(struct rbtree __arena *rbtree)
{
const size_t numkeys = sizeof(morekeys) / sizeof(morekeys[0]);
u64 value;
int ret;
int i;
for (i = 0; i < numkeys && can_loop; i++) {
u64 key = morekeys[i];
ret = rb_insert(rbtree, key, 2 * key);
if (ret)
return -1;
if (rb_integrity_check(rbtree)) {
arena_stderr("iteration %d\n", i);
return -EINVAL;
}
ret = rb_find(rbtree, key, &value);
if (ret)
return -1;
if (value != 2 * key)
return -1;
}
ret = remove_many_verify_all_present(rbtree);
if (ret)
return ret;
for (i = 0; i < numkeys && can_loop; i += 2) {
u64 key = morekeys[i];
ret = rb_remove(rbtree, key);
if (ret) {
arena_stderr("Failed to remove %ld\n", key);
return -1;
}
ret = rb_find(rbtree, key, &value);
if (!ret)
return -1;
}
return remove_many_verify_remaining(rbtree);
}
static __noinline int remove_many_noalloc(struct rbtree __arena *rbtree)
{
const size_t numkeys = sizeof(morekeys) / sizeof(morekeys[0]);
node_ctx first = NULL, last = NULL;
u64 value;
int ret;
int i;
for (i = 0; i < numkeys && can_loop; i++) {
u64 key = morekeys[i];
node_ctx nodec = arena_malloc(sizeof(*nodec));
if (!nodec) {
arena_stderr("out of memory\n");
return -ENOMEM;
}
nodec->rbnode.key = key;
nodec->rbnode.value = 2 * key;
nodec->next = NULL;
if (!first)
first = nodec;
if (last)
last->next = nodec;
last = nodec;
ret = rb_insert_node(rbtree, &nodec->rbnode);
if (ret)
return -1;
if (rb_integrity_check(rbtree)) {
arena_stderr("iteration %d\n", i);
return -EINVAL;
}
ret = rb_find(rbtree, key, &value);
if (ret)
return -1;
if (value != 2 * key)
return -1;
}
ret = remove_many_verify_all_present(rbtree);
if (ret)
return ret;
for (i = 0; i < numkeys && can_loop; i += 2) {
u64 key = morekeys[i];
node_ctx nodec = first;
if (!nodec || key != nodec->rbnode.key)
return -1;
first = nodec->next ? nodec->next->next : NULL;
ret = rb_remove_node(rbtree, &nodec->rbnode);
if (ret) {
arena_stderr("Failed to remove %ld\n", key);
return -1;
}
ret = rb_find(rbtree, key, &value);
if (!ret)
return -1;
}
return remove_many_verify_remaining(rbtree);
}
static inline int remove_many(enum rbtree_alloc alloc,
enum rbtree_insert_mode insert)
{
int ret;
struct rbtree __arena *rbtree;
rbtree = rb_create(alloc, insert);
if (!rbtree)
return -ENOMEM;
ret = (alloc == RB_ALLOC) ? remove_many_alloc(rbtree)
: remove_many_noalloc(rbtree);
if (ret)
return ret;
if (alloc == RB_ALLOC)
return rb_destroy(rbtree);
ret = clean_up_noalloc_tree(rbtree);
if (ret)
return ret;
return rb_destroy(rbtree);
}
SEC("syscall")
__weak int test_rbtree_insert_many_update(void)
{
return insert_many(RB_ALLOC, RB_UPDATE);
}
SEC("syscall")
__weak int test_rbtree_insert_many_noalloc(void)
{
return insert_many(RB_NOALLOC, RB_DUPLICATE);
}
SEC("syscall")
__weak int test_rbtree_remove_many_update(void)
{
return remove_many(RB_ALLOC, RB_UPDATE);
}
SEC("syscall")
__weak int test_rbtree_remove_many_noalloc(void)
{
return remove_many(RB_NOALLOC, RB_DUPLICATE);
}
SEC("syscall")
__weak int test_rbtree_add_remove_circular(void)
{
const size_t iters = 60;
const size_t prefill = 10;
const size_t numkeys = 50;
const size_t prefix = 400000;
u64 value, rmval;
int errval = 1;
u64 key;
int ret;
int i;
struct rbtree __arena *rbtree;
rbtree = rb_create(RB_ALLOC, RB_UPDATE);
if (!rbtree)
return 1;
for (i = 0; i < prefill && can_loop; i++) {
ret = rb_insert(rbtree, prefix + (i % numkeys), i);
if (ret)
return errval;
errval += 1;
}
errval = 2 * 1000 * 1000;
for (i = 0; i < prefill && can_loop; i++) {
/* Read it back. */
ret = rb_find(rbtree, prefix + (i % numkeys), &value);
if (ret)
return errval;
if (value != i)
return errval;
}
errval = 3 * 1000 * 1000;
for (i = prefill; i < iters && can_loop; i++) {
key = prefix + (i % numkeys);
ret = rb_find(rbtree, key, &value);
if (!ret) {
arena_stderr("Key %d already present\n", key);
return errval;
}
errval += 1;
ret = rb_insert(rbtree, key, i);
if (ret) {
arena_stderr("ITERATION %d\n", i);
rb_print(rbtree);
return errval;
}
rmval = i - prefill;
errval += 1;
ret = rb_find(rbtree, prefix + (rmval % numkeys), &value);
if (ret)
return errval;
errval += 1;
if (value != rmval)
return errval;
errval += 1;
ret = rb_remove(rbtree, prefix + (rmval % numkeys));
if (ret) {
arena_stderr("ITERATION %d\n", i);
return errval;
}
errval += 1;
}
for (i = 0; i < numkeys && can_loop; i++) {
rb_remove(rbtree, prefix + i);
}
return rb_destroy(rbtree);
}
SEC("syscall")
__weak int test_rbtree_add_remove_circular_reverse(void)
{
const size_t iters = 110;
const size_t prefill = 10;
const size_t numkeys = 50;
const size_t prefix = 500000;
u64 value, rmval;
int errval = 1;
u64 key;
int ret;
int i;
struct rbtree __arena *rbtree;
rbtree = rb_create(RB_ALLOC, RB_UPDATE);
if (!rbtree)
return 1;
for (i = 0; i < prefill && can_loop; i++) {
ret = rb_insert(rbtree, prefix - (i % numkeys), i);
if (ret)
return errval;
errval += 1;
}
errval = 2 * 1000 * 1000;
for (i = 0; i < prefill && can_loop; i++) {
/* Read it back. */
ret = rb_find(rbtree, prefix - (i % numkeys), &value);
if (ret)
return errval;
if (value != i)
return errval;
}
errval = 3 * 1000 * 1000;
for (i = prefill; i < iters && can_loop; i++) {
key = prefix - (i % numkeys);
ret = rb_find(rbtree, key, &value);
if (!ret) {
arena_stderr("Key %d already present\n", key);
return errval;
}
errval += 1;
ret = rb_insert(rbtree, key, i);
if (ret) {
arena_stderr("error %d on insert\n", ret);
rb_print(rbtree);
return errval;
}
rmval = i - prefill;
errval += 1;
ret = rb_find(rbtree, prefix - (rmval % numkeys), &value);
if (ret)
return errval;
errval += 1;
if (value != rmval)
return errval;
errval += 1;
ret = rb_remove(rbtree, prefix - (rmval % numkeys));
if (ret)
return errval;
errval += 1;
}
errval = 4 * 1000 * 1000;
for (i = 0; i < prefill && can_loop; i++) {
ret = rb_remove(rbtree, prefix - i);
if (ret) {
arena_stderr("Did not remove %d, error %d\n", prefix - i, ret);
return errval + i;
}
}
return rb_destroy(rbtree);
}
SEC("syscall")
__weak int test_rbtree_least_pop(void)
{
const size_t keys = 10;
u64 key, value;
int errval = 1;
int ret, i;
struct rbtree __arena *rbtree;
rbtree = rb_create(RB_ALLOC, RB_DEFAULT);
if (!rbtree)
return errval;
errval += 1;
for (i = 0; i < keys / 2 && can_loop; i++) {
ret = rb_insert(rbtree, i, i);
if (ret)
return errval;
errval += 1;
ret = rb_insert(rbtree, keys - 1 - i, keys - 1 - i);
if (ret)
return errval;
errval += 1;
ret = rb_least(rbtree, &key, &value);
if (ret)
return errval;
errval += 1;
if (key != 0 || value != 0)
return errval;
errval += 1;
}
errval = 1000;
for (i = 0; i < keys && can_loop; i++) {
ret = rb_least(rbtree, &key, &value);
if (ret) {
arena_stderr("rb_least failed with %d\n", ret);
return errval;
}
errval += 1;
if (key != i || value != i) {
arena_stderr("Got KV %ld/%ld expected %d\n", key, value, i);
return errval;
}
errval += 1;
ret = rb_pop(rbtree, &key, &value);
if (ret) {
arena_stderr("Error %d during pop on iter %d\n", ret, i);
return errval;
}
errval += 1;
if (key != i || value != i)
return errval;
}
return rb_destroy(rbtree);
}
/* Reject rb_pop() for RB_NOALLOC trees. */
SEC("syscall")
__weak int test_rbtree_noalloc_pop(void)
{
const u64 expect_value = 1;
const u64 expect_key = 0;
struct rbtree __arena *rbtree;
struct rbnode __arena *node;
u64 value = 0;
int ret;
rbtree = rb_create(RB_NOALLOC, RB_DEFAULT);
if (!rbtree)
return 1;
node = rb_node_alloc(expect_key, expect_value);
if (!node) {
rb_destroy(rbtree);
return 2;
}
ret = rb_insert_node(rbtree, node);
if (ret) {
rb_node_free(node);
rb_destroy(rbtree);
return 3;
}
ret = rb_pop(rbtree, NULL, &value);
if (ret != -EINVAL)
return 4;
ret = rb_find(rbtree, expect_key, &value);
if (ret)
return 5;
if (value != expect_value)
return 6;
ret = rb_remove_node(rbtree, node);
if (ret)
return 7;
rb_node_free(node);
return rb_destroy(rbtree);
}
SEC("syscall")
__weak int test_rbtree_alloc_check(void)
{
struct rbtree __arena *alloc, *noalloc;
struct rbnode __arena *node;
int ret;
alloc = rb_create(RB_ALLOC, RB_DEFAULT);
if (!alloc)
return 1;
noalloc = rb_create(RB_NOALLOC, RB_DEFAULT);
if (!noalloc)
return 2;
node = rb_node_alloc(0, 0);
if (!node)
return 3;
/*
* RB_ALLOC trees can use rb_insert, RB_NOALLOC trees can
* use rb_insert_node. RB_ALLOC and RB_NOALLOC trees cannot
* use each other's APIs.
*
* NOTE: This begs the question, why not different types? We
* want to partially share the API and that would require us
* to duplicate it.
*/
if (rb_insert(alloc, 0, 0))
return 4;
if (!rb_insert_node(alloc, node))
return 5;
if (!rb_remove_node(alloc, node))
return 6;
if (rb_remove(alloc, 0))
return 7;
if (rb_insert_node(noalloc, node))
return 8;
if (!rb_insert(noalloc, 0, 0))
return 9;
if (!rb_remove(noalloc, 0))
return 10;
if (rb_remove_node(noalloc, node))
return 11;
rb_node_free(node);
ret = rb_destroy(alloc);
if (ret)
return ret;
return rb_destroy(noalloc);
}
@@ -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;
}
@@ -0,0 +1,553 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <vmlinux.h>
#include <libarena/common.h>
#include <libarena/asan.h>
enum {
/*
* Is the access checked by check_region_inline
* a read or a write?
*/
ASAN_READ = 0x0U,
ASAN_WRITE = 0x1U,
};
/*
* Address sanitizer (ASAN) for arena-based BPF programs, inspired
* by KASAN.
*
* The API
* -------
*
* The implementation includes two kinds of components: Implementation
* of ASAN hooks injected by LLVM into the program, and API calls that
* allocators use to mark memory as valid or invalid. The full list is:
*
* LLVM stubs:
*
* void __asan_{load, store}<size>(intptr_t addr)
* Checks whether an access is valid. All variations covered
* by check_region_inline().
*
* void __asan_{store, load}((intptr_t addr, ssize_t size)
*
* void __asan_report_{load, store}<size>(intptr_t addr)
* Report an access violation for the program. Used when LLVM
* uses direct code generation for shadow map checks.
*
* void *__asan_memcpy(void *d, const void *s, size_t n)
* void *__asan_memmove(void *d, const void *s, size_t n)
* void *__asan_memset(void *p, int c, size_t n)
* Hooks for ASAN instrumentation of the LLVM mem* builtins.
* Currently unimplemented just like the builtins themselves.
*
* API methods:
*
* asan_init()
* Initialize the ASAN map for the arena.
*
* asan_poison()
* Mark a region of memory as poisoned. Accessing poisoned memory
* causes asan_report() to fire. Invoked during free().
*
* asan_unpoison()
* Mark a region as unpoisoned after alloc().
*
* asan_shadow_set()
* Check a byte's validity directly.
*
* The Algorithm In Brief
* ----------------------
* Each group of 8 bytes is mapped to a "granule" in the shadow map. This
* granule is the size of the byte and describes which bytes are valid.
* Possible values are:
*
* 0: All bytes are valid. Makes checks in the middle of an allocated region
* (most of them) fast.
* (0, 7]: How many consecutive bytes are valid, starting from the lowest one.
* The tradeoff is that we can't poison individual bytes in the middle of a
* valid region.
* [0x80, 0xff]: Special poison values, can be used to denote specific error
* modes (e.g., recently freed vs uninitialized memory).
*
* The mapping between a memory location and its shadow is:
* shadow_addr = shadow_base + (addr >> 3). We retain the 8:1 data:shadow
* ratio of existing ASAN implementations as a compromise between tracking
* granularity and space usage/scan overhead.
*/
#ifdef BPF_ARENA_ASAN
#pragma clang attribute push(__attribute__((no_sanitize("address"))), \
apply_to = function)
#define SHADOW_ALL_ZEROES ((u64)-1)
/*
* Canary variable for ASAN violations. Set to the offending address.
*/
volatile u64 asan_violated = 0;
/*
* Shadow map occupancy map.
*/
volatile u64 __asan_shadow_memory_dynamic_address;
volatile u32 asan_reported = false;
volatile bool asan_inited = false;
/*
* Set during program load.
*/
volatile bool asan_report_once = false;
/*
* BPF does not currently support the memset/memcpy/memcmp intrinsics.
* For large sequential copies, or assignments of large data structures,
* the frontend will generate an intrinsic that causes the BPF backend
* to exit due to a missing implementation. Provide a simple implementation
* just for memset to use it for poisoning/unpoisoning the map.
*/
__weak int asan_memset(s8 __arena *dst, s8 val, size_t size)
{
size_t i;
for (i = zero; i < size && can_loop; i++)
dst[i] = val;
return 0;
}
/* Validate a 1-byte access, always within a single byte. */
static __always_inline bool memory_is_poisoned_1(s8 __arena *addr)
{
s8 shadow_value = *(s8 __arena *)mem_to_shadow(addr);
/* Byte is 0, access is valid. */
if (likely(!shadow_value))
return false;
/*
* Byte is non-zero. Access is valid if granule offset in [0, shadow_value),
* so the memory is poisoned if shadow_value is negative or smaller than
* the granule's value.
*/
return ASAN_GRANULE(addr) >= shadow_value;
}
/* Validate a 2- 4-, 8-byte access, shadow spans up to 2 bytes. */
static __always_inline bool memory_is_poisoned_2_4_8(s8 __arena *addr, u64 size)
{
u64 end = (u64)addr + size - 1;
/*
* Region fully within a single byte (addition didn't
* overflow above ASAN_GRANULE).
*/
if (likely(ASAN_GRANULE(end) >= size - 1))
return memory_is_poisoned_1((s8 __arena *)end);
/*
* Otherwise first byte must be fully unpoisoned, and second byte
* must be unpoisoned up to the end of the accessed region.
*/
return *(s8 __arena *)mem_to_shadow(addr) || memory_is_poisoned_1((s8 __arena *)end);
}
__weak bool asan_shadow_set(void __arena *addr)
{
return memory_is_poisoned_1(addr);
}
static __always_inline u64 first_nonzero_byte(u64 addr, size_t size)
{
while (size && can_loop) {
if (unlikely(*(s8 __arena *)addr))
return addr;
addr += 1;
size -= 1;
}
return SHADOW_ALL_ZEROES;
}
static __always_inline bool memory_is_poisoned_n(s8 __arena *addr, u64 size)
{
u64 ret;
u64 start;
u64 end;
/* Size of [start, end] is end - start + 1. */
start = (u64)mem_to_shadow(addr);
end = (u64)mem_to_shadow(addr + size - 1);
ret = first_nonzero_byte(start, (end - start) + 1);
if (likely(ret == SHADOW_ALL_ZEROES))
return false;
return unlikely(ret != end || ASAN_GRANULE(addr + size - 1) >= *(s8 __arena *)end);
}
__weak int asan_report(s8 __arena *addr, size_t sz, u32 flags)
{
u32 reported = __sync_val_compare_and_swap(&asan_reported, false, true);
/* Only report the first ASAN violation. */
if (reported && asan_report_once)
return 0;
asan_violated = (u64)addr;
arena_stderr("Memory violation for address %p (0x%lx) for %s of size %ld\n",
addr, (u64)addr,
(flags & ASAN_WRITE) ? "write" : "read",
sz);
bpf_stream_print_stack(BPF_STDERR);
return 0;
}
static __always_inline bool check_asan_args(s8 __arena *addr, size_t size,
bool *result)
{
bool valid = true;
/* Size 0 accesses are valid even if the address is invalid. */
if (unlikely(size == 0))
goto confirmed_valid;
/*
* Wraparound is possible for values close to the the edge of the
* 4GiB boundary of the arena (last valid address is 1UL << 32 - 1).
*
*
* The wraparound detection below works for small sizes. check_asan_args is
* always called from the builtin ASAN checks, so 1 <= size <= 64. Even
* for storeN/loadN that we do not expect to encounter the intrinsics will
* not have a large enough size that:
*
* - addr + size > MAX_U32
* - (u32)(addr + size) > (u32) addr
*
* which would defeat wraparound detection.
*/
if (unlikely((u32)(u64)(addr + size) < (u32)(u64)addr))
goto confirmed_invalid;
return false;
confirmed_invalid:
valid = false;
/* FALLTHROUGH */
confirmed_valid:
*result = valid;
return true;
}
static __always_inline bool check_region_inline(intptr_t ptr, size_t size,
u32 flags)
{
s8 __arena *addr = (s8 __arena *)(u64)ptr;
bool is_poisoned, is_valid;
if (check_asan_args(addr, size, &is_valid)) {
if (!is_valid)
asan_report(addr, size, flags);
return is_valid;
}
switch (size) {
case 1:
is_poisoned = memory_is_poisoned_1(addr);
break;
case 2:
case 4:
case 8:
is_poisoned = memory_is_poisoned_2_4_8(addr, size);
break;
default:
is_poisoned = memory_is_poisoned_n(addr, size);
}
if (is_poisoned) {
asan_report(addr, size, flags);
return false;
}
return true;
}
/*
* __alias is not supported for BPF so define *__noabort() variants as wrappers.
*/
#define DEFINE_ASAN_LOAD_STORE(size) \
__hidden void __asan_store##size(intptr_t addr) \
{ \
check_region_inline(addr, size, ASAN_WRITE); \
} \
__hidden void __asan_store##size##_noabort(intptr_t addr) \
{ \
check_region_inline(addr, size, ASAN_WRITE); \
} \
__hidden void __asan_load##size(intptr_t addr) \
{ \
check_region_inline(addr, size, ASAN_READ); \
} \
__hidden void __asan_load##size##_noabort(intptr_t addr) \
{ \
check_region_inline(addr, size, ASAN_READ); \
} \
__hidden void __asan_report_store##size(intptr_t addr) \
{ \
asan_report((s8 __arena *)addr, size, ASAN_WRITE); \
} \
__hidden void __asan_report_store##size##_noabort(intptr_t addr) \
{ \
asan_report((s8 __arena *)addr, size, ASAN_WRITE); \
} \
__hidden void __asan_report_load##size(intptr_t addr) \
{ \
asan_report((s8 __arena *)addr, size, ASAN_READ); \
} \
__hidden void __asan_report_load##size##_noabort(intptr_t addr) \
{ \
asan_report((s8 __arena *)addr, size, ASAN_READ); \
}
DEFINE_ASAN_LOAD_STORE(1);
DEFINE_ASAN_LOAD_STORE(2);
DEFINE_ASAN_LOAD_STORE(4);
DEFINE_ASAN_LOAD_STORE(8);
void __asan_storeN(intptr_t addr, ssize_t size)
{
check_region_inline(addr, size, ASAN_WRITE);
}
void __asan_storeN_noabort(intptr_t addr, ssize_t size)
{
check_region_inline(addr, size, ASAN_WRITE);
}
void __asan_loadN(intptr_t addr, ssize_t size)
{
check_region_inline(addr, size, ASAN_READ);
}
void __asan_loadN_noabort(intptr_t addr, ssize_t size)
{
check_region_inline(addr, size, ASAN_READ);
}
/*
* We currently do not sanitize globals.
*/
void __asan_register_globals(intptr_t globals, size_t n)
{
}
void __asan_unregister_globals(intptr_t globals, size_t n)
{
}
/*
* We do not currently have memcpy/memmove/memset intrinsics
* in LLVM. Do not implement sanitization.
*/
void *__asan_memcpy(void *d, const void *s, size_t n)
{
arena_stderr("ASAN: Unexpected %s call", __func__);
return NULL;
}
void *__asan_memmove(void *d, const void *s, size_t n)
{
arena_stderr("ASAN: Unexpected %s call", __func__);
return NULL;
}
void *__asan_memset(void *p, int c, size_t n)
{
arena_stderr("ASAN: Unexpected %s call", __func__);
return NULL;
}
/*
* Poisoning code, used when we add more freed memory to the allocator by:
* a) pulling memory from the arena segment using bpf_arena_alloc_pages()
* b) freeing memory from application code
*/
__hidden __noasan int asan_poison(void __arena *addr, s8 val, size_t size)
{
s8 __arena *shadow;
size_t len;
/*
* Poisoning from a non-granule address makes no sense: We can only allocate
* memory to the application that has a granule-aligned starting address,
* and bpf_arena_alloc_pages returns page-aligned memory. A non-aligned
* addr then implies we're freeing a different address than the one we
* allocated.
*/
if (unlikely((u64)addr & ASAN_GRANULE_MASK))
return -EINVAL;
/*
* We cannot free an unaligned region because it'd be possible that we
* cannot describe the resulting poisoning state of the granule in
* the ASAN encoding.
*
* Every granule represents a region of memory that looks like the
* following (P for poisoned bytes, C for clear):
*
* <Clear> <Poisoned>
* [ C C C ... P P ]
*
* The value of the granule's shadow map is the number of clear bytes in
* it. We cannot represent granules with the following state:
*
* [ P P ... C C ... P P ]
*
* That would be possible if we could free unaligned regions, so prevent that.
*/
if (unlikely(size & ASAN_GRANULE_MASK))
return -EINVAL;
shadow = mem_to_shadow(addr);
len = size >> ASAN_SHADOW_SHIFT;
asan_memset(shadow, val, len);
return 0;
}
/*
* Unpoisoning code for marking memory as valid during allocation calls.
*
* Very similar to asan_poison, except we need to round up instead of
* down, then partially poison the last granule if necessary.
*
* Partial poisoning is useful for keeping the padding poisoned. Allocations
* are granule-aligned, so we we're reserving granule-aligned sizes for the
* allocation. However, we want to still treat accesses to the padding as
* invalid. Partial poisoning takes care of that. Freeing and poisoning the
* memory is still done in granule-aligned sizes and repoisons the already
* poisoned padding.
*/
__hidden __noasan int asan_unpoison(void __arena *addr, size_t size)
{
size_t partial = size & ASAN_GRANULE_MASK;
s8 __arena *shadow;
size_t len;
/*
* We cannot allocate in the middle of the granule. The ASAN shadow
* map encoding only describes regions of memory where every granule
* follows this format (P for poisoned, C for clear):
*
* <Clear> <Poisoned>
* [ C C C ... P P ]
*
* This is so we can use a single number in [0, ASAN_SHADOW_SCALE)
* to represent the poison state of the granule.
*/
if (unlikely((u64)addr & ASAN_GRANULE_MASK))
return -EINVAL;
shadow = mem_to_shadow(addr);
len = size >> ASAN_SHADOW_SHIFT;
asan_memset(shadow, 0, len);
/*
* If we are allocating a non-granule aligned region, we need to adjust
* the last byte of the shadow map to list how many bytes in the granule
* are unpoisoned. If the region is aligned, then the memset call above
* was enough.
*/
if (partial)
shadow[len] = partial;
return 0;
}
/*
* Initialize ASAN state when necessary. Triggered from userspace before
* allocator startup.
*/
SEC("syscall")
__weak __noasan int asan_init(struct asan_init_args *args)
{
u64 globals_pages = args->arena_globals_pages;
u64 all_pages = args->arena_all_pages;
u64 shadow_map, shadow_pgoff;
u64 shadow_pages;
if (asan_inited)
return 0;
/*
* Round up the shadow map size to the nearest page.
*/
shadow_pages = all_pages >> ASAN_SHADOW_SHIFT;
if ((all_pages & ((1 << ASAN_SHADOW_SHIFT) - 1)))
shadow_pages += 1;
if (all_pages > (1ULL << 32) / __PAGE_SIZE) {
arena_stderr("error: arena size %lx too large", all_pages);
return -EINVAL;
}
if (globals_pages > all_pages) {
arena_stderr("error: globals %lx do not fit in arena %lx",
globals_pages, all_pages);
return -EINVAL;
}
if (globals_pages + shadow_pages >= all_pages) {
arena_stderr("error: globals %lx do not leave room for shadow map %lx "
"(arena pages %lx)",
globals_pages, shadow_pages, all_pages);
return -EINVAL;
}
shadow_pgoff = all_pages - shadow_pages - globals_pages;
__asan_shadow_memory_dynamic_address = shadow_pgoff * __PAGE_SIZE;
/*
* Allocate the last (1/ASAN_SHADOW_SCALE)th of an arena's pages for the map
* We find the offset and size from the arena map.
*
* The allocated map pages are zeroed out, meaning all memory is marked as valid
* even if it's not allocated already. This is expected: Since the actual memory
* pages are not allocated, accesses to it will trigger page faults and will be
* reported through BPF streams. Any pages allocated through bpf_arena_alloc_pages
* should be poisoned by the allocator right after the call succeeds.
*/
shadow_map = (u64)bpf_arena_alloc_pages(
&arena, (void __arena *)__asan_shadow_memory_dynamic_address,
shadow_pages, NUMA_NO_NODE, 0);
if (!shadow_map) {
arena_stderr("Could not allocate shadow map\n");
__asan_shadow_memory_dynamic_address = 0;
return -ENOMEM;
}
asan_inited = true;
return 0;
}
#pragma clang attribute pop
#endif /* BPF_ARENA_ASAN */
__weak char _license[] SEC("license") = "GPL";
@@ -0,0 +1,903 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <libarena/common.h>
#include <libarena/asan.h>
#include <libarena/buddy.h>
/*
* Buddy allocator arena-based implementation.
*
* Memory is organized into chunks. These chunks
* cannot be coalesced or split. Allocating
* chunks allocates their memory eagerly.
*
* Internally, each chunk is organized into blocks.
* Blocks _can_ be coalesced/split, but only inside
* the chunk. Each block can be allocated or
* unallocated. If allocated, the entire block holds
* user data. If unallocated, the block is mostly
* invalid memory, with the exception of a header
* used for freelist tracking.
*
* The header is placed at an offset inside the block
* to prevent off-by-one errors from the previous block
* from trivially overwriting the header. Such an error
* is also not catchable by ASAN, since the header remains
* valid memory even after the block is freed. It is still
* theoretically possible for the header to be corrupted
* without being caught by ASAN, but harder.
*
* Since the allocator needs to track order information for
* both allocated and free blocks, and allocated blocks cannot
* store a header, the allocator also stores per-chunk order
* information in a reserved region at the beginning of the
* chunk. The header includes a bitmap with the order of blocks
* and their allocation state. It also includes the freelist
* heads for the allocation itself.
*/
enum {
BUDDY_POISONED = (s8)0xef,
/* Number of pages to be allocated per chunk. */
BUDDY_CHUNK_PAGES = BUDDY_CHUNK_BYTES / __PAGE_SIZE
};
static inline int buddy_lock(struct buddy __arena *buddy)
{
return arena_spin_lock(&buddy->lock);
}
static inline void buddy_unlock(struct buddy __arena *buddy)
{
arena_spin_unlock(&buddy->lock);
}
/*
* Reserve part of the arena address space for the allocator. We use
* this to get aligned addresses for the chunks, since the arena
* page alloc kfuncs do not support aligning to a boundary (in this
* case 1 MiB, see buddy.h on how this is derived).
*/
static int buddy_reserve_arena_vaddr(struct buddy __arena *buddy)
{
buddy->vaddr = 0;
return bpf_arena_reserve_pages(&arena,
(void __arena *)BUDDY_VADDR_OFFSET,
BUDDY_VADDR_SIZE / __PAGE_SIZE);
}
/*
* Free up any unused address space. Used only during teardown.
*/
static void buddy_unreserve_arena_vaddr(struct buddy __arena *buddy)
{
bpf_arena_free_pages(
&arena, (void __arena *)(BUDDY_VADDR_OFFSET + buddy->vaddr),
(BUDDY_VADDR_SIZE - buddy->vaddr) / __PAGE_SIZE);
buddy->vaddr = 0;
}
/*
* Carve out part of the reserved address space and hand it over
* to the buddy allocator.
*
* We are assuming the buddy allocator is the only allocator in the
* system, so there is no race between this function reserving a
* page range and some other allocator actually making the BPF call
* to really create and reserve it.
*
* However, bump allocation must still be atomic because this function
* is called without the buddy lock from multiple threads concurrently.
*/
__weak int buddy_alloc_arena_vaddr(struct buddy __arena *buddy, u64 *vaddrp)
{
u64 vaddr, old, new;
if (!buddy || !vaddrp)
return -EINVAL;
do {
vaddr = buddy->vaddr;
new = vaddr + BUDDY_CHUNK_BYTES;
if (new > BUDDY_VADDR_SIZE)
return -EINVAL;
old = __sync_val_compare_and_swap(&buddy->vaddr, vaddr, new);
} while (old != vaddr && can_loop);
if (old != vaddr)
return -EINVAL;
*vaddrp = BUDDY_VADDR_OFFSET + vaddr;
return 0;
}
static u64 arena_next_pow2(__u64 n)
{
n--;
n |= n >> 1;
n |= n >> 2;
n |= n >> 4;
n |= n >> 8;
n |= n >> 16;
n |= n >> 32;
n++;
return n;
}
__weak
int idx_set_allocated(struct buddy_chunk __arena *chunk, u64 idx, bool allocated)
{
bool already_allocated;
if (unlikely(idx >= BUDDY_CHUNK_ITEMS)) {
arena_stderr("setting state of invalid idx (%ld, max %d)\n", idx,
BUDDY_CHUNK_ITEMS);
return -EINVAL;
}
already_allocated = chunk->allocated[idx / 8] & (1 << (idx % 8));
if (unlikely(already_allocated == allocated)) {
arena_stderr("Double %s of idx %ld for chunk %p",
allocated ? "alloc" : "free",
idx, chunk);
return -EINVAL;
}
if (allocated)
chunk->allocated[idx / 8] |= 1 << (idx % 8);
else
chunk->allocated[idx / 8] &= ~(1 << (idx % 8));
return 0;
}
static int idx_is_allocated(struct buddy_chunk __arena *chunk, u64 idx, bool *allocated)
{
if (unlikely(idx >= BUDDY_CHUNK_ITEMS)) {
arena_stderr("getting state of invalid idx (%llu, max %d)\n", idx,
BUDDY_CHUNK_ITEMS);
return -EINVAL;
}
*allocated = chunk->allocated[idx / 8] & (1 << (idx % 8));
return 0;
}
__weak
int idx_set_order(struct buddy_chunk __arena *chunk, u64 idx, u8 order)
{
u8 prev_order;
if (unlikely(order >= BUDDY_CHUNK_NUM_ORDERS)) {
arena_stderr("setting invalid order %u\n", order);
return -EINVAL;
}
if (unlikely(idx >= BUDDY_CHUNK_ITEMS)) {
arena_stderr("setting order of invalid idx (%d, max %d)\n", idx,
BUDDY_CHUNK_ITEMS);
return -EINVAL;
}
/*
* We store two order instances per byte, one per nibble.
* Retain the existing nibble.
*/
prev_order = chunk->orders[idx / 2];
if (idx & 0x1) {
order &= 0xf;
order |= (prev_order & 0xf0);
} else {
order <<= 4;
order |= (prev_order & 0xf);
}
chunk->orders[idx / 2] = order;
return 0;
}
static u8 idx_get_order(struct buddy_chunk __arena *chunk, u64 idx)
{
u8 result;
_Static_assert(BUDDY_CHUNK_NUM_ORDERS <= 16,
"order must fit in 4 bits");
if (unlikely(idx >= BUDDY_CHUNK_ITEMS)) {
arena_stderr("getting order of invalid idx %u\n", idx);
return BUDDY_CHUNK_NUM_ORDERS;
}
result = chunk->orders[idx / 2];
return (idx & 0x1) ? (result & 0xf) : (result >> 4);
}
static void __arena *idx_to_addr(struct buddy_chunk __arena *chunk, size_t idx)
{
u64 address;
if (unlikely(idx >= BUDDY_CHUNK_ITEMS)) {
arena_stderr("translating invalid idx %u\n", idx);
return NULL;
}
/*
* The data blocks start in the chunk after the metadata block.
* We find the actual address by indexing into the region at an
* BUDDY_MIN_ALLOC_BYTES granularity, the minimum allowed.
* The index number already accounts for the fact that the first
* blocks in the chunk are occupied by the metadata, so we do
* not need to offset it.
*/
address = (u64)chunk + (idx * BUDDY_MIN_ALLOC_BYTES);
return (void __arena *)address;
}
static struct buddy_header __arena *idx_to_header(struct buddy_chunk __arena *chunk, size_t idx)
{
bool allocated;
u64 address;
if (unlikely(idx_is_allocated(chunk, idx, &allocated))) {
arena_stderr("accessing invalid idx 0x%lx\n", idx);
return NULL;
}
if (unlikely(allocated)) {
arena_stderr("accessing allocated idx 0x%lx as header\n", idx);
return NULL;
}
address = (u64)idx_to_addr(chunk, idx);
if (!address)
return NULL;
/*
* Offset the header within the block. This avoids accidental overwrites
* to the header because of off-by-one errors when using adjacent blocks.
*
* The offset has been chosen as a compromise between ASAN effectiveness
* and allocator granularity:
* 1) ASAN dictates valid data runs are 8-byte aligned.
* 2) We want to keep a low minimum allocation size (currently 16).
*
* As a result, we have only two possible positions for the header: Bytes
* 0 and 8. Keeping the header in byte 0 means off-by-ones from the previous
* block touch the header, and, since the header must be accessible, ASAN
* will not trigger. Keeping the header on byte 8 means off-by-one errors from
* the previous block are caught by ASAN. Negative offsets are rarer, so
* while accesses into the block from the next block are possible, they are
* less probable.
*/
return (struct buddy_header __arena *)(address + BUDDY_HEADER_OFF);
}
static void header_add_freelist(struct buddy_chunk __arena *chunk, struct buddy_header __arena *header,
u64 idx, u8 order)
{
struct buddy_header __arena *tmp_header;
idx_set_order(chunk, idx, order);
header->next_index = chunk->freelists[order];
header->prev_index = BUDDY_CHUNK_ITEMS;
if (header->next_index != BUDDY_CHUNK_ITEMS) {
tmp_header = idx_to_header(chunk, header->next_index);
tmp_header->prev_index = idx;
}
chunk->freelists[order] = idx;
}
static void header_remove_freelist(struct buddy_chunk __arena *chunk,
struct buddy_header __arena *header, u8 order)
{
struct buddy_header __arena *tmp_header;
if (header->prev_index != BUDDY_CHUNK_ITEMS) {
tmp_header = idx_to_header(chunk, header->prev_index);
tmp_header->next_index = header->next_index;
}
if (header->next_index != BUDDY_CHUNK_ITEMS) {
tmp_header = idx_to_header(chunk, header->next_index);
tmp_header->prev_index = header->prev_index;
}
/* Pop off the list head if necessary. */
if (idx_to_header(chunk, chunk->freelists[order]) == header)
chunk->freelists[order] = header->next_index;
header->prev_index = BUDDY_CHUNK_ITEMS;
header->next_index = BUDDY_CHUNK_ITEMS;
}
static u64 size_to_order(size_t size)
{
u64 order;
/*
* Legal sizes are [1, 4GiB] (the biggest possible arena).
* Of course, sizes close to GiB are practically impossible
* to fulfill and allocation will fail, but that's taken care
* of by the caller.
*/
if (unlikely(size == 0 || size > (1UL << 32))) {
arena_stderr("illegal size request %lu\n", size);
return 64;
}
/*
* To find the order of the allocation we find the first power of two
* >= the requested size, take the log2, then adjust it for the minimum
* allocation size by removing the minimum shift from it. Requests
* smaller than the minimum allocation size are rounded up.
*/
order = arena_fls(arena_next_pow2(size)) - 1;
if (order < BUDDY_MIN_ALLOC_SHIFT)
return 0;
return order - BUDDY_MIN_ALLOC_SHIFT;
}
__weak
int add_leftovers_to_freelist(struct buddy_chunk __arena *chunk, u32 cur_idx,
u64 min_order, u64 max_order)
{
struct buddy_header __arena *header;
u64 ord;
u32 idx;
for (ord = min_order; ord < max_order && can_loop; ord++) {
/* Mark the buddy as free and add it to the freelists. */
idx = cur_idx + (1 << ord);
header = idx_to_header(chunk, idx);
if (unlikely(!header)) {
arena_stderr("idx %u has no header", idx);
return -EINVAL;
}
asan_unpoison(header, sizeof(*header));
header_add_freelist(chunk, header, idx, ord);
}
return 0;
}
static struct buddy_chunk __arena *buddy_chunk_get(struct buddy __arena *buddy)
{
u64 order, ord, min_order, max_order;
struct buddy_chunk __arena *chunk;
size_t left;
int power2;
u64 vaddr;
u32 idx;
int ret;
/*
* Step 1: Allocate a properly aligned chunk, and
* prep it for insertion into the buddy allocator.
* We don't need the allocator lock until step 2.
*/
ret = buddy_alloc_arena_vaddr(buddy, &vaddr);
if (ret)
return NULL;
/* Addresses must be aligned to the chunk boundary. */
if (vaddr % BUDDY_CHUNK_BYTES)
return NULL;
/* Unreserve the address space. */
bpf_arena_free_pages(&arena, (void __arena *)vaddr,
BUDDY_CHUNK_PAGES);
chunk = bpf_arena_alloc_pages(&arena, (void __arena *)vaddr,
BUDDY_CHUNK_PAGES, NUMA_NO_NODE, 0);
if (!chunk) {
arena_stderr("[ALLOC FAILED]");
return NULL;
}
if (buddy_lock(buddy)) {
/*
* We cannot reclaim the vaddr space, but that is ok - this
* operation should always succeed. The error path is to catch
* accidental deadlocks that will cause -ENOMEMs to the program as
* the allocator fails to refill itself, in which case vaddr usage
* is the least of our worries.
*/
bpf_arena_free_pages(&arena, (void __arena *)vaddr, BUDDY_CHUNK_PAGES);
return NULL;
}
asan_poison(chunk, BUDDY_POISONED, BUDDY_CHUNK_PAGES * __PAGE_SIZE);
/* Unpoison the chunk itself. */
asan_unpoison(chunk, sizeof(*chunk));
/* Mark all freelists as empty. */
for (ord = zero; ord < BUDDY_CHUNK_NUM_ORDERS && can_loop; ord++)
chunk->freelists[ord] = BUDDY_CHUNK_ITEMS;
/*
* Initialize the chunk by carving out a page range to hold the metadata
* struct above, then dumping the rest of the pages into the allocator.
*/
_Static_assert(BUDDY_CHUNK_PAGES * __PAGE_SIZE >=
BUDDY_MIN_ALLOC_BYTES *
BUDDY_CHUNK_ITEMS,
"chunk must fit within the allocation");
/*
* Step 2: Reserve a chunk for the chunk metadata, then breaks
* the rest of the full allocation into the different buckets.
* We allocating the memory by grabbing blocks of progressively
* smaller sizes from the allocator, which are guaranteed to be
* continuous.
*
* This operation also populates the allocator.
*
* Algorithm:
*
* - max_order: The last order allocation we made
* - left: How many bytes are left to allocate
* - cur_index: Current index into the top-level block we are
* allocating from.
*
* Step 3:
* - Find the largest power-of-2 allocation still smaller than left (infimum)
* - Reserve a chunk of that size, along with its buddy
* - For every order from [infimum + 1, last order), carve out a block
* and put it into the allocator.
*
* Example: Chunk size 0b1010000 (80 bytes)
*
* Step 1:
*
* idx infimum 1 << max_order
* 0 64 128 1 << 20
* |________|_________|______________________|
*
* Blocks set aside:
* [0, 64) - Completely allocated
* [64, 128) - Will be further split in the next iteration
*
* Blocks added to the allocator:
* [128, 256)
* [256, 512)
* ...
* [1 << 18, 1 << 19)
* [1 << 19, 1 << 20)
*
* Step 2:
*
* idx infimum idx + 1 << max_order
* 64 80 96 64 + 1 << 6 = 128
* |________|_________|______________________|
*
* Blocks set aside:
* [64, 80) - Completely allocated
*
* Blocks added to the allocator:
* [80, 96) - left == 0 so the buddy is unused and marked as freed
* [96, 128)
*/
max_order = BUDDY_CHUNK_NUM_ORDERS;
left = sizeof(*chunk);
idx = 0;
while (left && can_loop) {
power2 = arena_fls(left) - 1;
/*
* Note: The condition below only triggers to catch serious bugs
* early. There is no sane way to undo any block insertions from
* the allocated chunk, so just leak any leftover allocations,
* emit a diagnostic, unlock and exit.
*
*/
if (unlikely(power2 >= BUDDY_CHUNK_NUM_ORDERS)) {
arena_stderr(
"buddy chunk metadata require allocation of order %d\n",
power2);
arena_stderr(
"chunk has size of 0x%lx bytes (left %lx bytes)\n",
sizeof(*chunk), left);
buddy_unlock(buddy);
return NULL;
}
/* Round up allocations that are too small. */
left -= (power2 >= BUDDY_MIN_ALLOC_SHIFT) ? 1 << power2 : left;
order = (power2 >= BUDDY_MIN_ALLOC_SHIFT) ? power2 - BUDDY_MIN_ALLOC_SHIFT : 0;
if (idx_set_allocated(chunk, idx, true)) {
buddy_unlock(buddy);
return NULL;
}
/*
* Starting an order above the one we allocated, populate
* the allocator with free blocks. If this is the last
* allocation (left == 0), also mark the buddy as free.
*
* See comment above about error handling: The error path
* is only there as a way to mitigate deeply buggy allocator
* states by emitting a diagnostic in add_leftovers_to_freelist()
* and leaking any memory not added in the freelists.
*/
min_order = left ? order + 1 : order;
if (add_leftovers_to_freelist(chunk, idx, min_order, max_order)) {
buddy_unlock(buddy);
return NULL;
}
/* Adjust the index. */
idx += 1 << order;
max_order = order;
}
buddy_unlock(buddy);
return chunk;
}
__weak int buddy_init(struct buddy __arena *buddy)
{
struct buddy_chunk __arena *chunk;
int ret;
if (!asan_ready())
return -EINVAL;
/* Reserve enough address space to ensure allocations are aligned. */
ret = buddy_reserve_arena_vaddr(buddy);
if (ret)
return ret;
_Static_assert(BUDDY_CHUNK_PAGES > 0,
"chunk must use one or more pages");
chunk = buddy_chunk_get(buddy);
if (buddy_lock(buddy)) {
bpf_arena_free_pages(&arena, chunk, BUDDY_CHUNK_PAGES);
return -EINVAL;
}
/* Chunk is already properly unpoisoned if allocated. */
if (chunk)
chunk->next = buddy->first_chunk;
/* Put the chunk at the beginning of the list. */
buddy->first_chunk = chunk;
buddy_unlock(buddy);
return chunk ? 0 : -ENOMEM;
}
/*
* Destroy the allocator. This does not check whether there are any allocations
* currently in use, so any pages being accessed will start taking arena faults.
* We do not take a lock because we are freeing arena pages, and nobody should
* be using the allocator at that point in the execution.
*/
__weak int buddy_destroy(struct buddy __arena *buddy)
{
struct buddy_chunk __arena *chunk, *next;
if (!buddy)
return -EINVAL;
/*
* Traverse all buddy chunks and free them back to the arena
* with the same granularity they were allocated with.
*/
for (chunk = buddy->first_chunk; chunk && can_loop; chunk = next) {
next = chunk->next;
/* Wholesale poison the entire block. */
asan_poison(chunk, BUDDY_POISONED,
BUDDY_CHUNK_PAGES * __PAGE_SIZE);
bpf_arena_free_pages(&arena, chunk, BUDDY_CHUNK_PAGES);
}
/* Free up any part of the address space that did not get used. */
buddy_unreserve_arena_vaddr(buddy);
/* Clear all fields. */
buddy->first_chunk = NULL;
return 0;
}
__weak u64 buddy_chunk_alloc(struct buddy_chunk __arena *chunk, int order_req)
{
struct buddy_header __arena *header, *tmp_header, *next_header;
u32 idx, tmpidx, retidx;
u64 address;
u64 order = 0;
u64 i;
for (order = order_req; order < BUDDY_CHUNK_NUM_ORDERS && can_loop; order++) {
if (chunk->freelists[order] != BUDDY_CHUNK_ITEMS)
break;
}
if (order >= BUDDY_CHUNK_NUM_ORDERS)
return (u64)NULL;
retidx = chunk->freelists[order];
header = idx_to_header(chunk, retidx);
if (unlikely(!header))
return (u64) NULL;
chunk->freelists[order] = header->next_index;
if (header->next_index != BUDDY_CHUNK_ITEMS) {
next_header = idx_to_header(chunk, header->next_index);
next_header->prev_index = BUDDY_CHUNK_ITEMS;
}
header->prev_index = BUDDY_CHUNK_ITEMS;
header->next_index = BUDDY_CHUNK_ITEMS;
if (idx_set_order(chunk, retidx, order_req))
return (u64)NULL;
if (idx_set_allocated(chunk, retidx, true))
return (u64)NULL;
/*
* Do not unpoison the address yet, will be done by the caller
* because the caller has the exact allocation size requested.
*/
address = (u64)idx_to_addr(chunk, retidx);
if (!address)
return (u64)NULL;
/* If we allocated from a larger-order chunk, split the buddies. */
for (i = order_req; i < order && can_loop; i++) {
/*
* Flip the bit for the current order (the bit is guaranteed
* to be 0, so just add 1 << i).
*/
idx = retidx + (1 << i);
/* Add the buddy of the allocation to the free list. */
header = idx_to_header(chunk, idx);
/* Unpoison the buddy header */
asan_unpoison(header, sizeof(*header));
if (idx_set_order(chunk, idx, i))
return (u64)NULL;
/* Push the header to the beginning of the freelists list. */
tmpidx = chunk->freelists[i];
header->prev_index = BUDDY_CHUNK_ITEMS;
header->next_index = tmpidx;
if (tmpidx != BUDDY_CHUNK_ITEMS) {
tmp_header = idx_to_header(chunk, tmpidx);
tmp_header->prev_index = idx;
}
chunk->freelists[i] = idx;
}
return address;
}
/* Scan the existing chunks for available memory. */
static u64 buddy_alloc_from_existing_chunks(struct buddy __arena *buddy, int order)
{
struct buddy_chunk __arena *chunk;
u64 address;
for (chunk = buddy->first_chunk; chunk != NULL && can_loop;
chunk = chunk->next) {
address = buddy_chunk_alloc(chunk, order);
if (address)
return address;
}
return (u64)NULL;
}
/*
* Try an allocation from a newly allocated chunk. Also
* incorporate the chunk into the linked list.
*/
static u64 buddy_alloc_from_new_chunk(struct buddy __arena *buddy, struct buddy_chunk __arena *chunk, int order)
{
u64 address;
if (buddy_lock(buddy))
return (u64)NULL;
/*
* Add the chunk into the allocator and try
* to allocate specifically from that chunk.
*/
chunk->next = buddy->first_chunk;
buddy->first_chunk = chunk;
address = buddy_chunk_alloc(buddy->first_chunk, order);
buddy_unlock(buddy);
return (u64)address;
}
__weak
void __arena *buddy_alloc(struct buddy __arena *buddy, size_t size)
{
void __arena *address = NULL;
struct buddy_chunk __arena *chunk;
int order;
if (!buddy)
return NULL;
order = size_to_order(size);
if (order >= BUDDY_CHUNK_NUM_ORDERS || order < 0) {
arena_stderr("invalid order %d (sz %lu)\n", order, size);
return NULL;
}
if (buddy_lock(buddy))
return NULL;
address = (u8 __arena *)buddy_alloc_from_existing_chunks(buddy, order);
buddy_unlock(buddy);
if (address)
goto done;
/* Get a new chunk. */
chunk = buddy_chunk_get(buddy);
if (chunk)
address = (u8 __arena *)buddy_alloc_from_new_chunk(buddy, chunk, order);
done:
/* If we failed to allocate memory, return NULL. */
if (!address)
return NULL;
/*
* Unpoison exactly the amount of bytes requested. If the
* data is smaller than the header, we must poison any
* unused bytes that were part of the header.
*/
if (size < BUDDY_HEADER_OFF + sizeof(struct buddy_header __arena))
asan_poison(address + BUDDY_HEADER_OFF, BUDDY_POISONED,
sizeof(struct buddy_header __arena));
asan_unpoison(address, size);
return address;
}
static __always_inline int buddy_free_unlocked(struct buddy __arena *buddy, u64 addr)
{
struct buddy_header __arena *header, *buddy_header;
u64 idx, buddy_idx, tmp_idx;
struct buddy_chunk __arena *chunk;
bool allocated;
u8 order;
int ret;
if (!buddy)
return -EINVAL;
if (addr & (BUDDY_MIN_ALLOC_BYTES - 1)) {
arena_stderr("Freeing unaligned address %llx\n", addr);
return -EINVAL;
}
/* Get (chunk, idx) out of the address. */
chunk = (void __arena *)(addr & ~BUDDY_CHUNK_OFFSET_MASK);
idx = (addr & BUDDY_CHUNK_OFFSET_MASK) / BUDDY_MIN_ALLOC_BYTES;
/* Mark the block as unallocated so we can access the header. */
ret = idx_set_allocated(chunk, idx, false);
if (ret)
return ret;
order = idx_get_order(chunk, idx);
header = idx_to_header(chunk, idx);
/* The header is in the block itself, keep it unpoisoned. */
asan_poison((u8 __arena *)addr, BUDDY_POISONED,
BUDDY_MIN_ALLOC_BYTES << order);
asan_unpoison(header, sizeof(*header));
/*
* Coalescing loop. Merge with free buddies of equal order.
* For every coalescing step, keep the left buddy and
* drop the right buddy's header.
*/
for (; order < BUDDY_CHUNK_NUM_ORDERS && can_loop; order++) {
buddy_idx = idx ^ (1 << order);
/* Check if the buddy is actually free. */
idx_is_allocated(chunk, buddy_idx, &allocated);
if (allocated)
break;
/*
* If buddy is not the same order as the chunk
* being freed, then we're done coalescing.
*/
if (idx_get_order(chunk, buddy_idx) != order)
break;
buddy_header = idx_to_header(chunk, buddy_idx);
header_remove_freelist(chunk, buddy_header, order);
/* Keep the left header out of the two buddies, drop the other one. */
if (buddy_idx < idx) {
tmp_idx = idx;
idx = buddy_idx;
buddy_idx = tmp_idx;
}
/* Remove the buddy from the freelists so that we can merge it. */
idx_set_order(chunk, buddy_idx, order);
buddy_header = idx_to_header(chunk, buddy_idx);
asan_poison(buddy_header, BUDDY_POISONED,
sizeof(*buddy_header));
}
/* Header properly freed but not in any freelists yet .*/
idx_set_order(chunk, idx, order);
header = idx_to_header(chunk, idx);
header_add_freelist(chunk, header, idx, order);
return 0;
}
__weak int buddy_free(struct buddy __arena *buddy, void __arena *addr)
{
int ret;
if (!buddy)
return -EINVAL;
/* Freeing NULL is a valid no-op. */
if (!addr)
return 0;
ret = buddy_lock(buddy);
if (ret)
return ret;
ret = buddy_free_unlocked(buddy, (u64)addr);
buddy_unlock(buddy);
return ret;
}
__weak char _license[] SEC("license") = "GPL";
@@ -0,0 +1,52 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <libarena/common.h>
#include <libarena/asan.h>
#include <libarena/buddy.h>
const volatile u32 zero = 0;
struct buddy __arena buddy;
int arena_fls(__u64 word)
{
if (!word)
return 0;
return 64 - __builtin_clzll(word);
}
SEC("syscall")
__weak int arena_get_info(struct arena_get_info_args *args)
{
args->arena_base = arena_base(&arena);
return 0;
}
SEC("syscall")
__weak int arena_alloc_reserve(struct arena_alloc_reserve_args *args)
{
return bpf_arena_reserve_pages(&arena, NULL, args->nr_pages);
}
SEC("syscall")
__weak int arena_buddy_reset(void)
{
buddy_destroy(&buddy);
return buddy_init(&buddy);
}
__weak void __arena *arena_malloc(size_t size)
{
return buddy_alloc(&buddy, size);
}
__weak void arena_free(void __arena *ptr)
{
buddy_free(&buddy, ptr);
}
char _license[] SEC("license") = "GPL";
File diff suppressed because it is too large Load Diff
@@ -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;
}