[PATCH 2/3] bpf: Add sleepable arena page allocation path

[Emil Tsalapatis <emil@etsalapatis.com>]

The bpf_arena_alloc_pages() function currently only allocates pages
inside a spinlock critical section with IRQs off. This forces the use
of alloc_pages_nolock() in the BPF allocator, even when the caller is
a sleepable BPF function. This in turn causes allocation failures even
in cases where falling into the allocator slow path and possibly
sleeping would eventually succeed. This can be triggered consistently
by heavy BPF arena users like scx.

Add a separate arena page allocation path just for sleepable callers.
The path preallocates the arena memory to be added to the tree before
taking the critical section.

Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 include/linux/bpf.h  |  6 ++++
 kernel/bpf/arena.c   | 76 +++++++++++++++++++++++++++++++++++++++++---
 kernel/bpf/syscall.c |  8 +----
 3 files changed, 79 insertions(+), 11 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index ffa5626411ac..d15ed7a3879b 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -710,6 +710,12 @@ void bpf_map_free_internal_structs(struct bpf_map *map, void *obj);
 int bpf_dynptr_from_file_sleepable(struct file *file, u32 flags,
 				   struct bpf_dynptr *ptr__uninit);
 
+static inline bool is_bpf_alloc_nonsleepable(void)
+{
+	return preempt_count() > 0 || irqs_disabled() ||
+		IS_ENABLED(CONFIG_PREEMPT_RT);
+}
+
 #if defined(CONFIG_MMU) && defined(CONFIG_64BIT)
 void *bpf_arena_alloc_pages_non_sleepable(void *p__map, void *addr__ign, u32 page_cnt, int node_id,
 					  u64 flags);
diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index da356989786a..126ddae763f9 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -683,8 +683,73 @@ static int arena_adjust_tree(struct bpf_arena *arena, long uaddr, long page_cnt,
 	return range_tree_clear(&arena->rt, *pgoff, page_cnt);
 }
 
-static long arena_alloc_pages_internal(struct bpf_arena *arena, long page_cnt,
-		long uaddr, long pgoff, int node_id, bool sleepable)
+static long arena_alloc_pages_sleepable(struct bpf_arena *arena, long page_cnt,
+		long uaddr, long pgoff, int node_id)
+{
+	u64 kern_vm_start = bpf_arena_get_kern_vm_start(arena);
+	struct apply_range_data data;
+	struct page **pages = NULL;
+	unsigned long flags;
+	u32 uaddr32;
+	int ret, i;
+
+	pages = kvcalloc(page_cnt, sizeof(struct page *), GFP_KERNEL_ACCOUNT);
+	if (!pages)
+		return 0;
+
+	ret = bpf_map_alloc_pages(&arena->map, node_id, page_cnt, pages);
+	if (ret) {
+		kvfree(pages);
+		return 0;
+	}
+
+	data.i = 0;
+	data.pages = pages;
+	data.arena = arena;
+
+	if (raw_res_spin_lock_irqsave(&arena->spinlock, flags))
+		goto out_free_pages;
+
+	ret = arena_adjust_tree(arena, uaddr, page_cnt, &pgoff);
+	if (ret) {
+		raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+		goto out_free_pages;
+	}
+
+	uaddr32 = (u32)(arena->user_vm_start + pgoff * PAGE_SIZE);
+
+	ret = apply_to_page_range(&init_mm, kern_vm_start + uaddr32,
+			page_cnt << PAGE_SHIFT, apply_range_set_cb, &data);
+
+	if (ret)
+		goto out_unmap;
+
+	flush_vmap_cache(kern_vm_start + uaddr32, page_cnt << PAGE_SHIFT);
+	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+
+	kvfree(pages);
+
+	return clear_lo32(arena->user_vm_start) + uaddr32;
+
+out_unmap:
+	/* data.i pages were mapped, undo only those who failed to map */
+	flush_vmap_cache(kern_vm_start + uaddr32, data.i << PAGE_SHIFT);
+	range_tree_set(&arena->rt, pgoff + data.i, page_cnt - data.i);
+
+	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+
+	arena_free_pages(arena, uaddr32, data.i, true);
+
+out_free_pages:
+	for (i = data.i; i < page_cnt; i++)
+		__free_page(pages[i]);
+
+	kvfree(pages);
+	return 0;
+}
+
+static long arena_alloc_pages_non_sleepable(struct bpf_arena *arena, long page_cnt,
+		long uaddr, long pgoff, int node_id)
 {
 	u64 kern_vm_start = bpf_arena_get_kern_vm_start(arena);
 	struct apply_range_data data;
@@ -762,7 +827,7 @@ static long arena_alloc_pages_internal(struct bpf_arena *arena, long page_cnt,
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
 	if (mapped) {
 		flush_vmap_cache(kern_vm_start + uaddr32, mapped << PAGE_SHIFT);
-		arena_free_pages(arena, uaddr32, mapped, sleepable);
+		arena_free_pages(arena, uaddr32, mapped, false);
 	}
 
 out_free_pages:
@@ -802,7 +867,10 @@ static long arena_alloc_pages(struct bpf_arena *arena, long uaddr, long page_cnt
 
 	bpf_map_memcg_enter(&arena->map, &old_memcg, &new_memcg);
 
-	addr = arena_alloc_pages_internal(arena, page_cnt, uaddr, pgoff, node_id, sleepable);
+	if (sleepable && !is_bpf_alloc_nonsleepable())
+		addr = arena_alloc_pages_sleepable(arena, page_cnt, uaddr, pgoff, node_id);
+	else
+		addr = arena_alloc_pages_non_sleepable(arena, page_cnt, uaddr, pgoff, node_id);
 
 	bpf_map_memcg_exit(old_memcg, new_memcg);
 
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index 6874ba1424af..700ed7a99579 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -596,15 +596,9 @@ static void bpf_map_release_memcg(struct bpf_map *map)
 }
 #endif
 
-static bool can_alloc_pages(void)
-{
-	return preempt_count() == 0 && !irqs_disabled() &&
-		!IS_ENABLED(CONFIG_PREEMPT_RT);
-}
-
 static struct page *__bpf_alloc_page(int nid)
 {
-	if (!can_alloc_pages())
+	if (is_bpf_alloc_nonsleepable())
 		return alloc_pages_nolock(__GFP_ACCOUNT, nid, 0);
 
 	return alloc_pages_node(nid,
-- 
2.54.0