kernel/bpf/range_tree.c

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