Add an arena-based bump allocator to the arenalib. The allocator is aimed towards data allocated at program startup whose lifetime is that of the BPF program. An example is CPU-related information for sched-ext scheduler. Signed-off-by: Emil Tsalapatis (Meta) --- .../selftests/bpf/libarena/include/bump.h | 20 ++ .../selftests/bpf/libarena/src/bump.bpf.c | 212 ++++++++++++++++++ 2 files changed, 232 insertions(+) create mode 100644 tools/testing/selftests/bpf/libarena/include/bump.h create mode 100644 tools/testing/selftests/bpf/libarena/src/bump.bpf.c diff --git a/tools/testing/selftests/bpf/libarena/include/bump.h b/tools/testing/selftests/bpf/libarena/include/bump.h new file mode 100644 index 000000000000..e41a7045c56f --- /dev/null +++ b/tools/testing/selftests/bpf/libarena/include/bump.h @@ -0,0 +1,20 @@ +#pragma once + +#ifdef __BPF__ + +u64 bump_alloc_internal(size_t bytes, size_t alignment); +#define bump_alloc(bytes, alignment) ((void __arena *)bump_alloc_internal((bytes), (alignment))) +int bump_init(size_t max_alloc_pages); +int bump_destroy(void); +int bump_memlimit(u64 lim_memusage); + +#endif /* __BPF__ */ + +struct bump { + size_t max_contig_bytes; + void __arena *memory; + size_t off; + size_t lim_memusage; + size_t cur_memusage; +}; + diff --git a/tools/testing/selftests/bpf/libarena/src/bump.bpf.c b/tools/testing/selftests/bpf/libarena/src/bump.bpf.c new file mode 100644 index 000000000000..7c0c5d9052ea --- /dev/null +++ b/tools/testing/selftests/bpf/libarena/src/bump.bpf.c @@ -0,0 +1,212 @@ +/* + * SPDX-License-Identifier: GPL-2.0 + * Copyright (c) 2024-2025 Meta Platforms, Inc. and affiliates. + * Copyright (c) 2024-2025 Tejun Heo + * Copyright (c) 2024-2025 Emil Tsalapatis + */ + +/* + * Static allocation module used to allocate arena memory for + * whose lifetime is that of the BPF program. Data is rarely + * allocated, mostly at program init, and never freed. The + * memory returned by this code is typeless so it avoids us + * having to define an allocator for each type. + */ + +#include +#include +#include + +/* Maximum memory that can be allocated by the arena. */ +#define ARENA_MAX_MEMORY (1ULL << 20) + +private(STATIC_ALLOC_LOCK) struct bpf_spin_lock static_lock; + +private(STATIC_ALLOC) struct bump bump; + +const s8 STATIC_POISON_UNINIT = 0xff; + +struct bump_ll; +struct bump_ll { + struct bump_ll __arena *next; +}; +typedef struct bump_ll __arena bump_ll_t; + +__weak u64 bump_alloc_internal(size_t bytes, size_t alignment) +{ + void __arena *memory, *old; + bump_ll_t *oldll, *newll; + size_t alloc_bytes; + size_t alloc_pages; + void __arena *ptr; + size_t padding; + u64 addr; + + /* + * Allocated addresses must be aligned to the nearest granule, + * and since we're stack allocating this implies that allocations + * sizes are also aligned. + */ + alignment = round_up(alignment, 1 << ASAN_SHADOW_SHIFT); + + bpf_spin_lock(&static_lock); + + /* Round up the current offset. */ + addr = (__u64)bump.memory + bump.off; + + padding = round_up(addr, alignment) - addr; + alloc_bytes = bytes + padding; + + if (alloc_bytes > bump.max_contig_bytes) { + bpf_spin_unlock(&static_lock); + bpf_printk("invalid request %ld, max is %ld\n", alloc_bytes, + bump.max_contig_bytes); + return (u64)NULL; + } + + /* + * The code assumes that the maximum static allocation + * size is significantly larger than the typical allocation + * size, so it does not attempt to alleviate memory + * fragmentation. + */ + if (bump.off + alloc_bytes > bump.max_contig_bytes) { + if (bump.cur_memusage + bump.max_contig_bytes > + bump.lim_memusage) { + bpf_spin_unlock(&static_lock); + bpf_printk("allocator memory limit exceeded"); + return (u64)NULL; + } + + old = bump.memory; + + alloc_pages = bump.max_contig_bytes / __PAGE_SIZE; + + memory = bpf_arena_alloc_pages(&arena, NULL, alloc_pages, + NUMA_NO_NODE, 0); + if (!memory) { + bpf_spin_unlock(&static_lock); + bpf_printk("failed to allocate memory"); + return (u64)NULL; + } + + asan_poison(memory, STATIC_POISON_UNINIT, + bump.max_contig_bytes); + + /* Keep a list of allocated blocks to free on allocator destruction. */ + oldll = (bump_ll_t *)old; + newll = (bump_ll_t *)memory; + asan_unpoison(newll, sizeof(*newll)); + newll->next = oldll; + + /* + * Switch to new memory block, reset offset, + * and recalculate base address. + */ + bump.memory = memory; + bump.off = sizeof(*newll); + addr = (__u64)bump.memory + bump.off; + + /* + * We changed the base address. Recompute the padding. + */ + padding = round_up(addr, alignment) - addr; + alloc_bytes = bytes + padding; + + bump.cur_memusage += bump.max_contig_bytes; + } + + ptr = (void __arena *)(addr + padding); + asan_unpoison(ptr, bytes); + + bump.off += alloc_bytes; + + bpf_spin_unlock(&static_lock); + + return (u64)ptr; +} + +__weak int bump_destroy(void) +{ + size_t alloc_pages = bump.max_contig_bytes / __PAGE_SIZE; + bump_ll_t *ll, *llnext; + int i; + + for (ll = bump.memory; ll && can_loop; ll = llnext) { + llnext = ll->next; + asan_unpoison(ll, bump.max_contig_bytes); + bpf_arena_free_pages(&arena, ll, alloc_pages); + } + + for (i = 0; i < sizeof(bump) && can_loop; i++) { + ((u8 *)&bump)[i] = 0; + } + + return 0; +} + +__weak int bump_init(size_t alloc_pages) +{ + size_t max_bytes = alloc_pages * __PAGE_SIZE; + void __arena *memory; + bump_ll_t *ll; + int ret; + + memory = bpf_arena_alloc_pages(&arena, NULL, alloc_pages, NUMA_NO_NODE, + 0); + if (!memory) { + bpf_printk("Failed to allocate %d pages", alloc_pages); + return -ENOMEM; + } + + ret = asan_poison(memory, STATIC_POISON_UNINIT, max_bytes); + if (ret) + bpf_printk("Error %d: by poisoning", ret); + + ret = asan_unpoison(memory, sizeof(*ll)); + if (ret) + bpf_printk("Error %d: by poisoning", ret); + + ll = (bump_ll_t *)memory; + ll->next = NULL; + + /* We reserve sizeof(*ll) for the embedded linked list. */ + bump = (struct bump){ + .max_contig_bytes = max_bytes, + .off = sizeof(*ll), + .memory = memory, + .lim_memusage = ARENA_MAX_MEMORY, + .cur_memusage = max_bytes, + }; + + return 0; +} + +__weak int bump_memlimit(u64 lim_memusage) +{ + bpf_spin_lock(&static_lock); + + if (lim_memusage > ARENA_MAX_MEMORY) + goto error; + + /* We always allocate at a page granularity. */ + if (lim_memusage % __PAGE_SIZE) + goto error; + + /* Have we already overshot the limit? */ + if (lim_memusage < bump.cur_memusage) + goto error; + + bump.lim_memusage = lim_memusage; + + bpf_spin_unlock(&static_lock); + + return 0; + +error: + bpf_spin_unlock(&static_lock); + + return -EINVAL; +} + +__weak char _license[] SEC("license") = "GPL"; -- 2.47.3