rbtree_search.c

  1// SPDX-License-Identifier: GPL-2.0
  2/* Copyright (c) 2025 Meta Platforms, Inc. and affiliates. */
  3
  4#include <vmlinux.h>
  5#include <bpf/bpf_helpers.h>
  6#include "bpf_misc.h"
  7#include "bpf_experimental.h"
  8
  9struct node_data {
 10	struct bpf_refcount ref;
 11	struct bpf_rb_node r0;
 12	struct bpf_rb_node r1;
 13	int key0;
 14	int key1;
 15};
 16
 17#define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8)))
 18private(A) struct bpf_spin_lock glock0;
 19private(A) struct bpf_rb_root groot0 __contains(node_data, r0);
 20
 21private(B) struct bpf_spin_lock glock1;
 22private(B) struct bpf_rb_root groot1 __contains(node_data, r1);
 23
 24#define rb_entry(ptr, type, member) container_of(ptr, type, member)
 25#define NR_NODES 16
 26
 27int zero = 0;
 28
 29static bool less0(struct bpf_rb_node *a, const struct bpf_rb_node *b)
 30{
 31	struct node_data *node_a;
 32	struct node_data *node_b;
 33
 34	node_a = rb_entry(a, struct node_data, r0);
 35	node_b = rb_entry(b, struct node_data, r0);
 36
 37	return node_a->key0 < node_b->key0;
 38}
 39
 40static bool less1(struct bpf_rb_node *a, const struct bpf_rb_node *b)
 41{
 42	struct node_data *node_a;
 43	struct node_data *node_b;
 44
 45	node_a = rb_entry(a, struct node_data, r1);
 46	node_b = rb_entry(b, struct node_data, r1);
 47
 48	return node_a->key1 < node_b->key1;
 49}
 50
 51SEC("syscall")
 52__retval(0)
 53long rbtree_search(void *ctx)
 54{
 55	struct bpf_rb_node *rb_n, *rb_m, *gc_ns[NR_NODES];
 56	long lookup_key = NR_NODES / 2;
 57	struct node_data *n, *m;
 58	int i, nr_gc = 0;
 59
 60	for (i = zero; i < NR_NODES && can_loop; i++) {
 61		n = bpf_obj_new(typeof(*n));
 62		if (!n)
 63			return __LINE__;
 64
 65		m = bpf_refcount_acquire(n);
 66
 67		n->key0 = i;
 68		m->key1 = i;
 69
 70		bpf_spin_lock(&glock0);
 71		bpf_rbtree_add(&groot0, &n->r0, less0);
 72		bpf_spin_unlock(&glock0);
 73
 74		bpf_spin_lock(&glock1);
 75		bpf_rbtree_add(&groot1, &m->r1, less1);
 76		bpf_spin_unlock(&glock1);
 77	}
 78
 79	n = NULL;
 80	bpf_spin_lock(&glock0);
 81	rb_n = bpf_rbtree_root(&groot0);
 82	while (can_loop) {
 83		if (!rb_n) {
 84			bpf_spin_unlock(&glock0);
 85			return __LINE__;
 86		}
 87
 88		n = rb_entry(rb_n, struct node_data, r0);
 89		if (lookup_key == n->key0)
 90			break;
 91		if (nr_gc < NR_NODES)
 92			gc_ns[nr_gc++] = rb_n;
 93		if (lookup_key < n->key0)
 94			rb_n = bpf_rbtree_left(&groot0, rb_n);
 95		else
 96			rb_n = bpf_rbtree_right(&groot0, rb_n);
 97	}
 98
 99	if (!n || lookup_key != n->key0) {
100		bpf_spin_unlock(&glock0);
101		return __LINE__;
102	}
103
104	for (i = 0; i < nr_gc; i++) {
105		rb_n = gc_ns[i];
106		gc_ns[i] = bpf_rbtree_remove(&groot0, rb_n);
107	}
108
109	m = bpf_refcount_acquire(n);
110	bpf_spin_unlock(&glock0);
111
112	for (i = 0; i < nr_gc; i++) {
113		rb_n = gc_ns[i];
114		if (rb_n) {
115			n = rb_entry(rb_n, struct node_data, r0);
116			bpf_obj_drop(n);
117		}
118	}
119
120	if (!m)
121		return __LINE__;
122
123	bpf_spin_lock(&glock1);
124	rb_m = bpf_rbtree_remove(&groot1, &m->r1);
125	bpf_spin_unlock(&glock1);
126	bpf_obj_drop(m);
127	if (!rb_m)
128		return __LINE__;
129	bpf_obj_drop(rb_entry(rb_m, struct node_data, r1));
130
131	return 0;
132}
133
134#define TEST_ROOT(dolock)				\
135SEC("syscall")						\
136__failure __msg(MSG)					\
137long test_root_spinlock_##dolock(void *ctx)		\
138{							\
139	struct bpf_rb_node *rb_n;			\
140	__u64 jiffies = 0;				\
141							\
142	if (dolock)					\
143		bpf_spin_lock(&glock0);			\
144	rb_n = bpf_rbtree_root(&groot0);		\
145	if (rb_n)					\
146		jiffies = bpf_jiffies64();		\
147	if (dolock)					\
148		bpf_spin_unlock(&glock0);		\
149							\
150	return !!jiffies;				\
151}
152
153#define TEST_LR(op, dolock)				\
154SEC("syscall")						\
155__failure __msg(MSG)					\
156long test_##op##_spinlock_##dolock(void *ctx)		\
157{							\
158	struct bpf_rb_node *rb_n;			\
159	struct node_data *n;				\
160	__u64 jiffies = 0;				\
161							\
162	bpf_spin_lock(&glock0);				\
163	rb_n = bpf_rbtree_root(&groot0);		\
164	if (!rb_n) {					\
165		bpf_spin_unlock(&glock0);		\
166		return 1;				\
167	}						\
168	n = rb_entry(rb_n, struct node_data, r0);	\
169	n = bpf_refcount_acquire(n);			\
170	bpf_spin_unlock(&glock0);			\
171	if (!n)						\
172		return 1;				\
173							\
174	if (dolock)					\
175		bpf_spin_lock(&glock0);			\
176	rb_n = bpf_rbtree_##op(&groot0, &n->r0);	\
177	if (rb_n)					\
178		jiffies = bpf_jiffies64();		\
179	if (dolock)					\
180		bpf_spin_unlock(&glock0);		\
181							\
182	return !!jiffies;				\
183}
184
185/*
186 * Use a separate MSG macro instead of passing to TEST_XXX(..., MSG)
187 * to ensure the message itself is not in the bpf prog lineinfo
188 * which the verifier includes in its log.
189 * Otherwise, the test_loader will incorrectly match the prog lineinfo
190 * instead of the log generated by the verifier.
191 */
192#define MSG "call bpf_rbtree_root{{.+}}; R0{{(_w)?}}=rcu_ptr_or_null_node_data(id={{[0-9]+}},non_own_ref"
193TEST_ROOT(true)
194#undef MSG
195#define MSG "call bpf_rbtree_{{(left|right).+}}; R0{{(_w)?}}=rcu_ptr_or_null_node_data(id={{[0-9]+}},non_own_ref"
196TEST_LR(left,  true)
197TEST_LR(right, true)
198#undef MSG
199
200#define MSG "bpf_spin_lock at off=0 must be held for bpf_rb_root"
201TEST_ROOT(false)
202TEST_LR(left, false)
203TEST_LR(right, false)
204#undef MSG
205
206char _license[] SEC("license") = "GPL";