1 /* SPDX-License-Identifier: GPL-2.0 */
2 /*
3 * BPF extensible scheduler class: Documentation/scheduler/sched-ext.rst
4 *
5 * scx_arena_pool: kernel-side sub-allocator over BPF-arena pages.
6 *
7 * Each chunk added to @sch->arena_pool comes from one
8 * bpf_arena_alloc_pages_sleepable() call and is registered at the
9 * kernel-side mapping address.
10 *
11 * Allocations grow the pool on demand. Underlying arena pages are released
12 * when the arena map itself is torn down.
13 *
14 * Copyright (c) 2026 Meta Platforms, Inc. and affiliates.
15 * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
16 */
17 #include <linux/genalloc.h>
18
19 #include "internal.h"
20 #include "arena.h"
21
22 enum scx_arena_consts {
23 SCX_ARENA_MIN_ORDER = 3, /* 8-byte minimum sub-allocation */
24 SCX_ARENA_GROW_PAGES = 4, /* per growth */
25 };
26
scx_arena_pool_init(struct scx_sched * sch)27 s32 scx_arena_pool_init(struct scx_sched *sch)
28 {
29 if (!sch->arena_map)
30 return 0;
31
32 sch->arena_pool = gen_pool_create(SCX_ARENA_MIN_ORDER, NUMA_NO_NODE);
33 if (!sch->arena_pool)
34 return -ENOMEM;
35 return 0;
36 }
37
scx_arena_clear_chunk(struct gen_pool * pool,struct gen_pool_chunk * chunk,void * data)38 static void scx_arena_clear_chunk(struct gen_pool *pool, struct gen_pool_chunk *chunk,
39 void *data)
40 {
41 int order = pool->min_alloc_order;
42 size_t chunk_sz = chunk->end_addr - chunk->start_addr + 1;
43 unsigned long end_bit = chunk_sz >> order;
44 unsigned long b, e;
45
46 for_each_set_bitrange(b, e, chunk->bits, end_bit)
47 gen_pool_free(pool, chunk->start_addr + (b << order),
48 (e - b) << order);
49 }
50
51 /*
52 * Tear down the pool. Outstanding gen_pool allocations are freed via
53 * scx_arena_clear_chunk() so gen_pool_destroy() doesn't BUG. The underlying
54 * arena pages are released when the arena map itself is torn down.
55 */
scx_arena_pool_destroy(struct scx_sched * sch)56 void scx_arena_pool_destroy(struct scx_sched *sch)
57 {
58 if (!sch->arena_pool)
59 return;
60 gen_pool_for_each_chunk(sch->arena_pool, scx_arena_clear_chunk, NULL);
61 gen_pool_destroy(sch->arena_pool);
62 sch->arena_pool = NULL;
63 }
64
65 /*
66 * Grow the pool by @page_cnt pages. bpf_arena_alloc_pages_sleepable() and
67 * gen_pool_add() (which calls vzalloc(GFP_KERNEL)) require a sleepable
68 * context.
69 */
scx_arena_grow(struct scx_sched * sch,u32 page_cnt)70 static int scx_arena_grow(struct scx_sched *sch, u32 page_cnt)
71 {
72 void *p;
73 int ret;
74
75 if (!sch->arena_map || !sch->arena_pool)
76 return -EINVAL;
77
78 p = bpf_arena_alloc_pages_sleepable(sch->arena_map, NULL,
79 page_cnt, NUMA_NO_NODE, 0);
80 if (!p)
81 return -ENOMEM;
82
83 ret = gen_pool_add(sch->arena_pool,
84 (unsigned long)scx_arena_to_kaddr(sch, p),
85 page_cnt * PAGE_SIZE, NUMA_NO_NODE);
86 if (ret) {
87 bpf_arena_free_pages_non_sleepable(sch->arena_map, p, page_cnt);
88 return ret;
89 }
90 return 0;
91 }
92
93 /*
94 * Allocate @size bytes from the arena pool. Returns kernel VA on success, NULL
95 * on failure. May grow the pool via scx_arena_grow() which sleeps. Caller must
96 * be in a GFP_KERNEL context.
97 */
scx_arena_alloc(struct scx_sched * sch,size_t size)98 void *scx_arena_alloc(struct scx_sched *sch, size_t size)
99 {
100 unsigned long kern_va;
101 u32 page_cnt;
102
103 might_sleep();
104
105 if (!sch->arena_pool)
106 return NULL;
107
108 while (true) {
109 kern_va = gen_pool_alloc(sch->arena_pool, size);
110 if (kern_va)
111 break;
112 page_cnt = max_t(u32, SCX_ARENA_GROW_PAGES,
113 (size + PAGE_SIZE - 1) >> PAGE_SHIFT);
114 if (scx_arena_grow(sch, page_cnt))
115 return NULL;
116 }
117
118 return (void *)kern_va;
119 }
120
scx_arena_free(struct scx_sched * sch,void * kern_va,size_t size)121 void scx_arena_free(struct scx_sched *sch, void *kern_va, size_t size)
122 {
123 if (sch->arena_pool && kern_va)
124 gen_pool_free(sch->arena_pool, (unsigned long)kern_va, size);
125 }
126