1/* SPDX-License-Identifier: GPL-2.0 */
2/*
3 * Latched RB-trees
4 *
5 * Copyright (C) 2015 Intel Corp., Peter Zijlstra <peterz@infradead.org>
6 *
7 * Since RB-trees have non-atomic modifications they're not immediately suited
8 * for RCU/lockless queries. Even though we made RB-tree lookups non-fatal for
9 * lockless lookups; we cannot guarantee they return a correct result.
10 *
11 * The simplest solution is a seqlock + RB-tree, this will allow lockless
12 * lookups; but has the constraint (inherent to the seqlock) that read sides
13 * cannot nest in write sides.
14 *
15 * If we need to allow unconditional lookups (say as required for NMI context
16 * usage) we need a more complex setup; this data structure provides this by
17 * employing the latch technique -- see @write_seqcount_latch_begin -- to
18 * implement a latched RB-tree which does allow for unconditional lookups by
19 * virtue of always having (at least) one stable copy of the tree.
20 *
21 * However, while we have the guarantee that there is at all times one stable
22 * copy, this does not guarantee an iteration will not observe modifications.
23 * What might have been a stable copy at the start of the iteration, need not
24 * remain so for the duration of the iteration.
25 *
26 * Therefore, this does require a lockless RB-tree iteration to be non-fatal;
27 * see the comment in lib/rbtree.c. Note however that we only require the first
28 * condition -- not seeing partial stores -- because the latch thing isolates
29 * us from loops. If we were to interrupt a modification the lookup would be
30 * pointed at the stable tree and complete while the modification was halted.
31 */
32
33#ifndef RB_TREE_LATCH_H
34#define RB_TREE_LATCH_H
35
36#include <linux/rbtree.h>
37#include <linux/seqlock.h>
38#include <linux/rcupdate.h>
39
40struct latch_tree_node {
41 struct rb_node node[2];
42};
43
44struct latch_tree_root {
45 seqcount_latch_t seq;
46 struct rb_root tree[2];
47};
48
49/**
50 * latch_tree_ops - operators to define the tree order
51 * @less: used for insertion; provides the (partial) order between two elements.
52 * @comp: used for lookups; provides the order between the search key and an element.
53 *
54 * The operators are related like:
55 *
56 * comp(a->key,b) < 0 := less(a,b)
57 * comp(a->key,b) > 0 := less(b,a)
58 * comp(a->key,b) == 0 := !less(a,b) && !less(b,a)
59 *
60 * If these operators define a partial order on the elements we make no
61 * guarantee on which of the elements matching the key is found. See
62 * latch_tree_find().
63 */
64struct latch_tree_ops {
65 bool (*less)(struct latch_tree_node *a, struct latch_tree_node *b);
66 int (*comp)(void *key, struct latch_tree_node *b);
67};
68
69static __always_inline struct latch_tree_node *
70__lt_from_rb(struct rb_node *node, int idx)
71{
72 return container_of(node, struct latch_tree_node, node[idx]);
73}
74
75static __always_inline void
76__lt_insert(struct latch_tree_node *ltn, struct latch_tree_root *ltr, int idx,
77 bool (*less)(struct latch_tree_node *a, struct latch_tree_node *b))
78{
79 struct rb_root *root = &ltr->tree[idx];
80 struct rb_node **link = &root->rb_node;
81 struct rb_node *node = &ltn->node[idx];
82 struct rb_node *parent = NULL;
83 struct latch_tree_node *ltp;
84
85 while (*link) {
86 parent = *link;
87 ltp = __lt_from_rb(node: parent, idx);
88
89 if (less(ltn, ltp))
90 link = &parent->rb_left;
91 else
92 link = &parent->rb_right;
93 }
94
95 rb_link_node_rcu(node, parent, rb_link: link);
96 rb_insert_color(node, root);
97}
98
99static __always_inline void
100__lt_erase(struct latch_tree_node *ltn, struct latch_tree_root *ltr, int idx)
101{
102 rb_erase(&ltn->node[idx], &ltr->tree[idx]);
103}
104
105static __always_inline struct latch_tree_node *
106__lt_find(void *key, struct latch_tree_root *ltr, int idx,
107 int (*comp)(void *key, struct latch_tree_node *node))
108{
109 struct rb_node *node = rcu_dereference_raw(ltr->tree[idx].rb_node);
110 struct latch_tree_node *ltn;
111 int c;
112
113 while (node) {
114 ltn = __lt_from_rb(node, idx);
115 c = comp(key, ltn);
116
117 if (c < 0)
118 node = rcu_dereference_raw(node->rb_left);
119 else if (c > 0)
120 node = rcu_dereference_raw(node->rb_right);
121 else
122 return ltn;
123 }
124
125 return NULL;
126}
127
128/**
129 * latch_tree_insert() - insert @node into the trees @root
130 * @node: nodes to insert
131 * @root: trees to insert @node into
132 * @ops: operators defining the node order
133 *
134 * It inserts @node into @root in an ordered fashion such that we can always
135 * observe one complete tree. See the comment for write_seqcount_latch_begin().
136 *
137 * The inserts use rcu_assign_pointer() to publish the element such that the
138 * tree structure is stored before we can observe the new @node.
139 *
140 * All modifications (latch_tree_insert, latch_tree_remove) are assumed to be
141 * serialized.
142 */
143static __always_inline void
144latch_tree_insert(struct latch_tree_node *node,
145 struct latch_tree_root *root,
146 const struct latch_tree_ops *ops)
147{
148 write_seqcount_latch_begin(s: &root->seq);
149 __lt_insert(ltn: node, ltr: root, idx: 0, less: ops->less);
150 write_seqcount_latch(s: &root->seq);
151 __lt_insert(ltn: node, ltr: root, idx: 1, less: ops->less);
152 write_seqcount_latch_end(s: &root->seq);
153}
154
155/**
156 * latch_tree_erase() - removes @node from the trees @root
157 * @node: nodes to remote
158 * @root: trees to remove @node from
159 * @ops: operators defining the node order
160 *
161 * Removes @node from the trees @root in an ordered fashion such that we can
162 * always observe one complete tree. See the comment for
163 * write_seqcount_latch_begin().
164 *
165 * It is assumed that @node will observe one RCU quiescent state before being
166 * reused of freed.
167 *
168 * All modifications (latch_tree_insert, latch_tree_remove) are assumed to be
169 * serialized.
170 */
171static __always_inline void
172latch_tree_erase(struct latch_tree_node *node,
173 struct latch_tree_root *root,
174 const struct latch_tree_ops *ops)
175{
176 write_seqcount_latch_begin(s: &root->seq);
177 __lt_erase(ltn: node, ltr: root, idx: 0);
178 write_seqcount_latch(s: &root->seq);
179 __lt_erase(ltn: node, ltr: root, idx: 1);
180 write_seqcount_latch_end(s: &root->seq);
181}
182
183/**
184 * latch_tree_find() - find the node matching @key in the trees @root
185 * @key: search key
186 * @root: trees to search for @key
187 * @ops: operators defining the node order
188 *
189 * Does a lockless lookup in the trees @root for the node matching @key.
190 *
191 * It is assumed that this is called while holding the appropriate RCU read
192 * side lock.
193 *
194 * If the operators define a partial order on the elements (there are multiple
195 * elements which have the same key value) it is undefined which of these
196 * elements will be found. Nor is it possible to iterate the tree to find
197 * further elements with the same key value.
198 *
199 * Returns: a pointer to the node matching @key or NULL.
200 */
201static __always_inline struct latch_tree_node *
202latch_tree_find(void *key, struct latch_tree_root *root,
203 const struct latch_tree_ops *ops)
204{
205 struct latch_tree_node *node;
206 unsigned int seq;
207
208 do {
209 seq = read_seqcount_latch(s: &root->seq);
210 node = __lt_find(key, ltr: root, idx: seq & 1, comp: ops->comp);
211 } while (read_seqcount_latch_retry(s: &root->seq, start: seq));
212
213 return node;
214}
215
216#endif /* RB_TREE_LATCH_H */
217