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
less(struct bpf_rb_node * a,const struct bpf_rb_node * b)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
__add_three(struct bpf_rb_root * root,struct bpf_spin_lock * lock)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")
rbtree_add_nodes(void * ctx)69 long rbtree_add_nodes(void *ctx)
70 {
71 return __add_three(&groot, &glock);
72 }
73
74 SEC("tc")
rbtree_add_and_remove(void * ctx)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")
rbtree_first_and_remove(void * ctx)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