xref: /linux/kernel/sched/ext/arena.c (revision 7db28abbea0f7dc1ec4fdfdc149db5fbd9e4c994)
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 
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 
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  */
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  */
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  */
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 
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