1 // SPDX-License-Identifier: GPL-2.0 2 /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */ 3 4 #include <vmlinux.h> 5 #include <bpf/bpf_tracing.h> 6 #include <bpf/bpf_helpers.h> 7 #include <bpf/bpf_core_read.h> 8 #include "bpf_experimental.h" 9 10 struct node_data { 11 long key; 12 long data; 13 struct bpf_rb_node node; 14 }; 15 16 long less_callback_ran = -1; 17 long removed_key = -1; 18 long first_data[2] = {-1, -1}; 19 20 #define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8))) 21 private(A) struct bpf_spin_lock glock; 22 private(A) struct bpf_rb_root groot __contains(node_data, node); 23 24 static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b) 25 { 26 struct node_data *node_a; 27 struct node_data *node_b; 28 29 node_a = container_of(a, struct node_data, node); 30 node_b = container_of(b, struct node_data, node); 31 less_callback_ran = 1; 32 33 return node_a->key < node_b->key; 34 } 35 36 static long __add_three(struct bpf_rb_root *root, struct bpf_spin_lock *lock) 37 { 38 struct node_data *n, *m; 39 40 n = bpf_obj_new(typeof(*n)); 41 if (!n) 42 return 1; 43 n->key = 5; 44 45 m = bpf_obj_new(typeof(*m)); 46 if (!m) { 47 bpf_obj_drop(n); 48 return 2; 49 } 50 m->key = 1; 51 52 bpf_spin_lock(&glock); 53 bpf_rbtree_add(&groot, &n->node, less); 54 bpf_rbtree_add(&groot, &m->node, less); 55 bpf_spin_unlock(&glock); 56 57 n = bpf_obj_new(typeof(*n)); 58 if (!n) 59 return 3; 60 n->key = 3; 61 62 bpf_spin_lock(&glock); 63 bpf_rbtree_add(&groot, &n->node, less); 64 bpf_spin_unlock(&glock); 65 return 0; 66 } 67 68 SEC("tc") 69 long rbtree_add_nodes(void *ctx) 70 { 71 return __add_three(&groot, &glock); 72 } 73 74 SEC("tc") 75 long rbtree_add_and_remove(void *ctx) 76 { 77 struct bpf_rb_node *res = NULL; 78 struct node_data *n, *m; 79 80 n = bpf_obj_new(typeof(*n)); 81 if (!n) 82 goto err_out; 83 n->key = 5; 84 85 m = bpf_obj_new(typeof(*m)); 86 if (!m) 87 goto err_out; 88 m->key = 3; 89 90 bpf_spin_lock(&glock); 91 bpf_rbtree_add(&groot, &n->node, less); 92 bpf_rbtree_add(&groot, &m->node, less); 93 res = bpf_rbtree_remove(&groot, &n->node); 94 bpf_spin_unlock(&glock); 95 96 n = container_of(res, struct node_data, node); 97 removed_key = n->key; 98 99 bpf_obj_drop(n); 100 101 return 0; 102 err_out: 103 if (n) 104 bpf_obj_drop(n); 105 if (m) 106 bpf_obj_drop(m); 107 return 1; 108 } 109 110 SEC("tc") 111 long rbtree_first_and_remove(void *ctx) 112 { 113 struct bpf_rb_node *res = NULL; 114 struct node_data *n, *m, *o; 115 116 n = bpf_obj_new(typeof(*n)); 117 if (!n) 118 return 1; 119 n->key = 3; 120 n->data = 4; 121 122 m = bpf_obj_new(typeof(*m)); 123 if (!m) 124 goto err_out; 125 m->key = 5; 126 m->data = 6; 127 128 o = bpf_obj_new(typeof(*o)); 129 if (!o) 130 goto err_out; 131 o->key = 1; 132 o->data = 2; 133 134 bpf_spin_lock(&glock); 135 bpf_rbtree_add(&groot, &n->node, less); 136 bpf_rbtree_add(&groot, &m->node, less); 137 bpf_rbtree_add(&groot, &o->node, less); 138 139 res = bpf_rbtree_first(&groot); 140 if (!res) { 141 bpf_spin_unlock(&glock); 142 return 2; 143 } 144 145 o = container_of(res, struct node_data, node); 146 first_data[0] = o->data; 147 148 res = bpf_rbtree_remove(&groot, &o->node); 149 bpf_spin_unlock(&glock); 150 151 o = container_of(res, struct node_data, node); 152 removed_key = o->key; 153 154 bpf_obj_drop(o); 155 156 bpf_spin_lock(&glock); 157 res = bpf_rbtree_first(&groot); 158 if (!res) { 159 bpf_spin_unlock(&glock); 160 return 3; 161 } 162 163 o = container_of(res, struct node_data, node); 164 first_data[1] = o->data; 165 bpf_spin_unlock(&glock); 166 167 return 0; 168 err_out: 169 if (n) 170 bpf_obj_drop(n); 171 if (m) 172 bpf_obj_drop(m); 173 return 1; 174 } 175 176 char _license[] SEC("license") = "GPL"; 177