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 */
bpf_alloc(unsigned int size)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
bpf_free(void __arena * addr)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
bpf_alloc(unsigned int size)65 static inline void __arena* bpf_alloc(unsigned int size) { return NULL; }
bpf_free(void __arena * addr)66 static inline void bpf_free(void __arena *addr) {}
67 #endif
68