1 // SPDX-License-Identifier: GPL-2.0-only 2 /* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */ 3 #include <linux/interval_tree_generic.h> 4 #include <linux/slab.h> 5 #include <linux/bpf.h> 6 #include "range_tree.h" 7 8 /* 9 * struct range_tree is a data structure used to allocate contiguous memory 10 * ranges in bpf arena. It's a large bitmap. The contiguous sequence of bits is 11 * represented by struct range_node or 'rn' for short. 12 * rn->rn_rbnode links it into an interval tree while 13 * rn->rb_range_size links it into a second rbtree sorted by size of the range. 14 * __find_range() performs binary search and best fit algorithm to find the 15 * range less or equal requested size. 16 * range_tree_clear/set() clears or sets a range of bits in this bitmap. The 17 * adjacent ranges are merged or split at the same time. 18 * 19 * The split/merge logic is based/borrowed from XFS's xbitmap32 added 20 * in commit 6772fcc8890a ("xfs: convert xbitmap to interval tree"). 21 * 22 * The implementation relies on external lock to protect rbtree-s. 23 * The alloc/free of range_node-s is done via kmalloc_nolock(). 24 * 25 * bpf arena is using range_tree to represent unallocated slots. 26 * At init time: 27 * range_tree_set(rt, 0, max); 28 * Then: 29 * start = range_tree_find(rt, len); 30 * if (start >= 0) 31 * range_tree_clear(rt, start, len); 32 * to find free range and mark slots as allocated and later: 33 * range_tree_set(rt, start, len); 34 * to mark as unallocated after use. 35 */ 36 struct range_node { 37 struct rb_node rn_rbnode; 38 struct rb_node rb_range_size; 39 u32 rn_start; 40 u32 rn_last; /* inclusive */ 41 u32 __rn_subtree_last; 42 }; 43 44 static struct range_node *rb_to_range_node(struct rb_node *rb) 45 { 46 return rb_entry(rb, struct range_node, rb_range_size); 47 } 48 49 static u32 rn_size(struct range_node *rn) 50 { 51 return rn->rn_last - rn->rn_start + 1; 52 } 53 54 /* Find range that fits best to requested size */ 55 static inline struct range_node *__find_range(struct range_tree *rt, u32 len) 56 { 57 struct rb_node *rb = rt->range_size_root.rb_root.rb_node; 58 struct range_node *best = NULL; 59 60 while (rb) { 61 struct range_node *rn = rb_to_range_node(rb); 62 63 if (len <= rn_size(rn)) { 64 best = rn; 65 rb = rb->rb_right; 66 } else { 67 rb = rb->rb_left; 68 } 69 } 70 71 return best; 72 } 73 74 s64 range_tree_find(struct range_tree *rt, u32 len) 75 { 76 struct range_node *rn; 77 78 rn = __find_range(rt, len); 79 if (!rn) 80 return -ENOENT; 81 return rn->rn_start; 82 } 83 84 /* Insert the range into rbtree sorted by the range size */ 85 static inline void __range_size_insert(struct range_node *rn, 86 struct rb_root_cached *root) 87 { 88 struct rb_node **link = &root->rb_root.rb_node, *rb = NULL; 89 u64 size = rn_size(rn); 90 bool leftmost = true; 91 92 while (*link) { 93 rb = *link; 94 if (size > rn_size(rb_to_range_node(rb))) { 95 link = &rb->rb_left; 96 } else { 97 link = &rb->rb_right; 98 leftmost = false; 99 } 100 } 101 102 rb_link_node(&rn->rb_range_size, rb, link); 103 rb_insert_color_cached(&rn->rb_range_size, root, leftmost); 104 } 105 106 #define START(node) ((node)->rn_start) 107 #define LAST(node) ((node)->rn_last) 108 109 INTERVAL_TREE_DEFINE(struct range_node, rn_rbnode, u32, 110 __rn_subtree_last, START, LAST, 111 static inline __maybe_unused, 112 __range_it) 113 114 static inline __maybe_unused void 115 range_it_insert(struct range_node *rn, struct range_tree *rt) 116 { 117 __range_size_insert(rn, &rt->range_size_root); 118 __range_it_insert(rn, &rt->it_root); 119 } 120 121 static inline __maybe_unused void 122 range_it_remove(struct range_node *rn, struct range_tree *rt) 123 { 124 rb_erase_cached(&rn->rb_range_size, &rt->range_size_root); 125 RB_CLEAR_NODE(&rn->rb_range_size); 126 __range_it_remove(rn, &rt->it_root); 127 } 128 129 static inline __maybe_unused struct range_node * 130 range_it_iter_first(struct range_tree *rt, u32 start, u32 last) 131 { 132 return __range_it_iter_first(&rt->it_root, start, last); 133 } 134 135 /* Clear the range in this range tree */ 136 int range_tree_clear(struct range_tree *rt, u32 start, u32 len) 137 { 138 u32 last = start + len - 1; 139 struct range_node *new_rn; 140 struct range_node *rn; 141 142 while ((rn = range_it_iter_first(rt, start, last))) { 143 if (rn->rn_start < start && rn->rn_last > last) { 144 u32 old_last = rn->rn_last; 145 146 /* Overlaps with the entire clearing range */ 147 range_it_remove(rn, rt); 148 rn->rn_last = start - 1; 149 range_it_insert(rn, rt); 150 151 /* Add a range */ 152 new_rn = kmalloc_nolock(sizeof(struct range_node), __GFP_ACCOUNT, 153 NUMA_NO_NODE); 154 if (!new_rn) 155 return -ENOMEM; 156 new_rn->rn_start = last + 1; 157 new_rn->rn_last = old_last; 158 range_it_insert(new_rn, rt); 159 } else if (rn->rn_start < start) { 160 /* Overlaps with the left side of the clearing range */ 161 range_it_remove(rn, rt); 162 rn->rn_last = start - 1; 163 range_it_insert(rn, rt); 164 } else if (rn->rn_last > last) { 165 /* Overlaps with the right side of the clearing range */ 166 range_it_remove(rn, rt); 167 rn->rn_start = last + 1; 168 range_it_insert(rn, rt); 169 break; 170 } else { 171 /* in the middle of the clearing range */ 172 range_it_remove(rn, rt); 173 kfree_nolock(rn); 174 } 175 } 176 return 0; 177 } 178 179 /* Is the whole range set ? */ 180 int is_range_tree_set(struct range_tree *rt, u32 start, u32 len) 181 { 182 u32 last = start + len - 1; 183 struct range_node *left; 184 185 /* Is this whole range set ? */ 186 left = range_it_iter_first(rt, start, last); 187 if (left && left->rn_start <= start && left->rn_last >= last) 188 return 0; 189 return -ESRCH; 190 } 191 192 /* Set the range in this range tree */ 193 int range_tree_set(struct range_tree *rt, u32 start, u32 len) 194 { 195 u32 last = start + len - 1; 196 struct range_node *right; 197 struct range_node *left; 198 int err; 199 200 /* Is this whole range already set ? */ 201 left = range_it_iter_first(rt, start, last); 202 if (left && left->rn_start <= start && left->rn_last >= last) 203 return 0; 204 205 /* Clear out everything in the range we want to set. */ 206 err = range_tree_clear(rt, start, len); 207 if (err) 208 return err; 209 210 /* Do we have a left-adjacent range ? */ 211 left = range_it_iter_first(rt, start - 1, start - 1); 212 if (left && left->rn_last + 1 != start) 213 return -EFAULT; 214 215 /* Do we have a right-adjacent range ? */ 216 right = range_it_iter_first(rt, last + 1, last + 1); 217 if (right && right->rn_start != last + 1) 218 return -EFAULT; 219 220 if (left && right) { 221 /* Combine left and right adjacent ranges */ 222 range_it_remove(left, rt); 223 range_it_remove(right, rt); 224 left->rn_last = right->rn_last; 225 range_it_insert(left, rt); 226 kfree_nolock(right); 227 } else if (left) { 228 /* Combine with the left range */ 229 range_it_remove(left, rt); 230 left->rn_last = last; 231 range_it_insert(left, rt); 232 } else if (right) { 233 /* Combine with the right range */ 234 range_it_remove(right, rt); 235 right->rn_start = start; 236 range_it_insert(right, rt); 237 } else { 238 left = kmalloc_nolock(sizeof(struct range_node), __GFP_ACCOUNT, NUMA_NO_NODE); 239 if (!left) 240 return -ENOMEM; 241 left->rn_start = start; 242 left->rn_last = last; 243 range_it_insert(left, rt); 244 } 245 return 0; 246 } 247 248 void range_tree_destroy(struct range_tree *rt) 249 { 250 struct range_node *rn; 251 252 while ((rn = range_it_iter_first(rt, 0, -1U))) { 253 range_it_remove(rn, rt); 254 kfree_nolock(rn); 255 } 256 } 257 258 void range_tree_init(struct range_tree *rt) 259 { 260 rt->it_root = RB_ROOT_CACHED; 261 rt->range_size_root = RB_ROOT_CACHED; 262 } 263