bpf: Use kmalloc_nolock() in range tree

The range tree uses bpf_mem_alloc() that is safe to be called from all
contexts and uses a pre-allocated pool of memory to serve these
allocations.

Replace bpf_mem_alloc() with kmalloc_nolock() as it can be called safely
from all contexts and is more scalable than bpf_mem_alloc().

Remove the migrate_disable/enable pairs as they were only needed for
bpf_mem_alloc() as it does per-cpu operations, kmalloc_nolock() doesn't
need this.

Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
Link: https://lore.kernel.org/r/20251106170608.4800-1-puranjay@kernel.org
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/kernel/bpf/range_tree.c b/kernel/bpf/range_tree.c
index 37b80a23ae1aed31c84f2676bd1214be6d7fbbff..99c63d982c5d67531d7962fc943fe48f3a6ab25e 100644 (file)
--- a/kernel/bpf/range_tree.c
+++ b/kernel/bpf/range_tree.c
@@ -2,7 +2,6 @@
 /* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */
 #include <linux/interval_tree_generic.h>
 #include <linux/slab.h>
-#include <linux/bpf_mem_alloc.h>
 #include <linux/bpf.h>
 #include "range_tree.h"
 
@@ -21,7 +20,7 @@
  * in commit 6772fcc8890a ("xfs: convert xbitmap to interval tree").
  *
  * The implementation relies on external lock to protect rbtree-s.
- * The alloc/free of range_node-s is done via bpf_mem_alloc.
+ * The alloc/free of range_node-s is done via kmalloc_nolock().
  *
  * bpf arena is using range_tree to represent unallocated slots.
  * At init time:
@@ -150,9 +149,7 @@ int range_tree_clear(struct range_tree *rt, u32 start, u32 len)
                        range_it_insert(rn, rt);
 
                        /* Add a range */
-                       migrate_disable();
-                       new_rn = bpf_mem_alloc(&bpf_global_ma, sizeof(struct range_node));
-                       migrate_enable();
+                       new_rn = kmalloc_nolock(sizeof(struct range_node), 0, NUMA_NO_NODE);
                        if (!new_rn)
                                return -ENOMEM;
                        new_rn->rn_start = last + 1;
@@ -172,9 +169,7 @@ int range_tree_clear(struct range_tree *rt, u32 start, u32 len)
                } else {
                        /* in the middle of the clearing range */
                        range_it_remove(rn, rt);
-                       migrate_disable();
-                       bpf_mem_free(&bpf_global_ma, rn);
-                       migrate_enable();
+                       kfree_nolock(rn);
                }
        }
        return 0;
@@ -227,9 +222,7 @@ int range_tree_set(struct range_tree *rt, u32 start, u32 len)
                range_it_remove(right, rt);
                left->rn_last = right->rn_last;
                range_it_insert(left, rt);
-               migrate_disable();
-               bpf_mem_free(&bpf_global_ma, right);
-               migrate_enable();
+               kfree_nolock(right);
        } else if (left) {
                /* Combine with the left range */
                range_it_remove(left, rt);
@@ -241,9 +234,7 @@ int range_tree_set(struct range_tree *rt, u32 start, u32 len)
                right->rn_start = start;
                range_it_insert(right, rt);
        } else {
-               migrate_disable();
-               left = bpf_mem_alloc(&bpf_global_ma, sizeof(struct range_node));
-               migrate_enable();
+               left = kmalloc_nolock(sizeof(struct range_node), 0, NUMA_NO_NODE);
                if (!left)
                        return -ENOMEM;
                left->rn_start = start;
@@ -259,7 +250,7 @@ void range_tree_destroy(struct range_tree *rt)
 
        while ((rn = range_it_iter_first(rt, 0, -1U))) {
                range_it_remove(rn, rt);
-               bpf_mem_free(&bpf_global_ma, rn);
+               kfree_nolock(rn);
        }
 }