xref: /linux/tools/testing/selftests/bpf/bpf_arena_alloc.h (revision 24168c5e6dfbdd5b414f048f47f75d64533296ca)
1 /* SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause) */
2 /* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */
3 #pragma once
4 #include "bpf_arena_common.h"
5 
6 #ifndef __round_mask
7 #define __round_mask(x, y) ((__typeof__(x))((y)-1))
8 #endif
9 #ifndef round_up
10 #define round_up(x, y) ((((x)-1) | __round_mask(x, y))+1)
11 #endif
12 
13 #ifdef __BPF__
14 #define NR_CPUS (sizeof(struct cpumask) * 8)
15 
16 static void __arena * __arena page_frag_cur_page[NR_CPUS];
17 static int __arena page_frag_cur_offset[NR_CPUS];
18 
19 /* Simple page_frag allocator */
20 static inline void __arena* bpf_alloc(unsigned int size)
21 {
22 	__u64 __arena *obj_cnt;
23 	__u32 cpu = bpf_get_smp_processor_id();
24 	void __arena *page = page_frag_cur_page[cpu];
25 	int __arena *cur_offset = &page_frag_cur_offset[cpu];
26 	int offset;
27 
28 	size = round_up(size, 8);
29 	if (size >= PAGE_SIZE - 8)
30 		return NULL;
31 	if (!page) {
32 refill:
33 		page = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
34 		if (!page)
35 			return NULL;
36 		cast_kern(page);
37 		page_frag_cur_page[cpu] = page;
38 		*cur_offset = PAGE_SIZE - 8;
39 		obj_cnt = page + PAGE_SIZE - 8;
40 		*obj_cnt = 0;
41 	} else {
42 		cast_kern(page);
43 		obj_cnt = page + PAGE_SIZE - 8;
44 	}
45 
46 	offset = *cur_offset - size;
47 	if (offset < 0)
48 		goto refill;
49 
50 	(*obj_cnt)++;
51 	*cur_offset = offset;
52 	return page + offset;
53 }
54 
55 static inline void bpf_free(void __arena *addr)
56 {
57 	__u64 __arena *obj_cnt;
58 
59 	addr = (void __arena *)(((long)addr) & ~(PAGE_SIZE - 1));
60 	obj_cnt = addr + PAGE_SIZE - 8;
61 	if (--(*obj_cnt) == 0)
62 		bpf_arena_free_pages(&arena, addr, 1);
63 }
64 #else
65 static inline void __arena* bpf_alloc(unsigned int size) { return NULL; }
66 static inline void bpf_free(void __arena *addr) {}
67 #endif
68