include/linux/rbtree_latch.h

0001 /* SPDX-License-Identifier: GPL-2.0 */
0002 /*
0003  * Latched RB-trees
0004  *
0005  * Copyright (C) 2015 Intel Corp., Peter Zijlstra <peterz@infradead.org>
0006  *
0007  * Since RB-trees have non-atomic modifications they're not immediately suited
0008  * for RCU/lockless queries. Even though we made RB-tree lookups non-fatal for
0009  * lockless lookups; we cannot guarantee they return a correct result.
0010  *
0011  * The simplest solution is a seqlock + RB-tree, this will allow lockless
0012  * lookups; but has the constraint (inherent to the seqlock) that read sides
0013  * cannot nest in write sides.
0014  *
0015  * If we need to allow unconditional lookups (say as required for NMI context
0016  * usage) we need a more complex setup; this data structure provides this by
0017  * employing the latch technique -- see @raw_write_seqcount_latch -- to
0018  * implement a latched RB-tree which does allow for unconditional lookups by
0019  * virtue of always having (at least) one stable copy of the tree.
0020  *
0021  * However, while we have the guarantee that there is at all times one stable
0022  * copy, this does not guarantee an iteration will not observe modifications.
0023  * What might have been a stable copy at the start of the iteration, need not
0024  * remain so for the duration of the iteration.
0025  *
0026  * Therefore, this does require a lockless RB-tree iteration to be non-fatal;
0027  * see the comment in lib/rbtree.c. Note however that we only require the first
0028  * condition -- not seeing partial stores -- because the latch thing isolates
0029  * us from loops. If we were to interrupt a modification the lookup would be
0030  * pointed at the stable tree and complete while the modification was halted.
0031  */
0032
0033 #ifndef RB_TREE_LATCH_H
0034 #define RB_TREE_LATCH_H
0035
0036 #include <linux/rbtree.h>
0037 #include <linux/seqlock.h>
0038 #include <linux/rcupdate.h>
0039
0040 struct latch_tree_node {
0041     struct rb_node node[2];
0042 };
0043
0044 struct latch_tree_root {
0045     seqcount_latch_t    seq;
0046     struct rb_root      tree[2];
0047 };
0048
0049 /**
0050  * latch_tree_ops - operators to define the tree order
0051  * @less: used for insertion; provides the (partial) order between two elements.
0052  * @comp: used for lookups; provides the order between the search key and an element.
0053  *
0054  * The operators are related like:
0055  *
0056  *  comp(a->key,b) < 0  := less(a,b)
0057  *  comp(a->key,b) > 0  := less(b,a)
0058  *  comp(a->key,b) == 0 := !less(a,b) && !less(b,a)
0059  *
0060  * If these operators define a partial order on the elements we make no
0061  * guarantee on which of the elements matching the key is found. See
0062  * latch_tree_find().
0063  */
0064 struct latch_tree_ops {
0065     bool (*less)(struct latch_tree_node *a, struct latch_tree_node *b);
0066     int  (*comp)(void *key,                 struct latch_tree_node *b);
0067 };
0068
0069 static __always_inline struct latch_tree_node *
0070 __lt_from_rb(struct rb_node *node, int idx)
0071 {
0072     return container_of(node, struct latch_tree_node, node[idx]);
0073 }
0074
0075 static __always_inline void
0076 __lt_insert(struct latch_tree_node *ltn, struct latch_tree_root *ltr, int idx,
0077         bool (*less)(struct latch_tree_node *a, struct latch_tree_node *b))
0078 {
0079     struct rb_root *root = &ltr->tree[idx];
0080     struct rb_node **link = &root->rb_node;
0081     struct rb_node *node = &ltn->node[idx];
0082     struct rb_node *parent = NULL;
0083     struct latch_tree_node *ltp;
0084
0085     while (*link) {
0086         parent = *link;
0087         ltp = __lt_from_rb(parent, idx);
0088
0089         if (less(ltn, ltp))
0090             link = &parent->rb_left;
0091         else
0092             link = &parent->rb_right;
0093     }
0094
0095     rb_link_node_rcu(node, parent, link);
0096     rb_insert_color(node, root);
0097 }
0098
0099 static __always_inline void
0100 __lt_erase(struct latch_tree_node *ltn, struct latch_tree_root *ltr, int idx)
0101 {
0102     rb_erase(&ltn->node[idx], &ltr->tree[idx]);
0103 }
0104
0105 static __always_inline struct latch_tree_node *
0106 __lt_find(void *key, struct latch_tree_root *ltr, int idx,
0107       int (*comp)(void *key, struct latch_tree_node *node))
0108 {
0109     struct rb_node *node = rcu_dereference_raw(ltr->tree[idx].rb_node);
0110     struct latch_tree_node *ltn;
0111     int c;
0112
0113     while (node) {
0114         ltn = __lt_from_rb(node, idx);
0115         c = comp(key, ltn);
0116
0117         if (c < 0)
0118             node = rcu_dereference_raw(node->rb_left);
0119         else if (c > 0)
0120             node = rcu_dereference_raw(node->rb_right);
0121         else
0122             return ltn;
0123     }
0124
0125     return NULL;
0126 }
0127
0128 /**
0129  * latch_tree_insert() - insert @node into the trees @root
0130  * @node: nodes to insert
0131  * @root: trees to insert @node into
0132  * @ops: operators defining the node order
0133  *
0134  * It inserts @node into @root in an ordered fashion such that we can always
0135  * observe one complete tree. See the comment for raw_write_seqcount_latch().
0136  *
0137  * The inserts use rcu_assign_pointer() to publish the element such that the
0138  * tree structure is stored before we can observe the new @node.
0139  *
0140  * All modifications (latch_tree_insert, latch_tree_remove) are assumed to be
0141  * serialized.
0142  */
0143 static __always_inline void
0144 latch_tree_insert(struct latch_tree_node *node,
0145           struct latch_tree_root *root,
0146           const struct latch_tree_ops *ops)
0147 {
0148     raw_write_seqcount_latch(&root->seq);
0149     __lt_insert(node, root, 0, ops->less);
0150     raw_write_seqcount_latch(&root->seq);
0151     __lt_insert(node, root, 1, ops->less);
0152 }
0153
0154 /**
0155  * latch_tree_erase() - removes @node from the trees @root
0156  * @node: nodes to remote
0157  * @root: trees to remove @node from
0158  * @ops: operators defining the node order
0159  *
0160  * Removes @node from the trees @root in an ordered fashion such that we can
0161  * always observe one complete tree. See the comment for
0162  * raw_write_seqcount_latch().
0163  *
0164  * It is assumed that @node will observe one RCU quiescent state before being
0165  * reused of freed.
0166  *
0167  * All modifications (latch_tree_insert, latch_tree_remove) are assumed to be
0168  * serialized.
0169  */
0170 static __always_inline void
0171 latch_tree_erase(struct latch_tree_node *node,
0172          struct latch_tree_root *root,
0173          const struct latch_tree_ops *ops)
0174 {
0175     raw_write_seqcount_latch(&root->seq);
0176     __lt_erase(node, root, 0);
0177     raw_write_seqcount_latch(&root->seq);
0178     __lt_erase(node, root, 1);
0179 }
0180
0181 /**
0182  * latch_tree_find() - find the node matching @key in the trees @root
0183  * @key: search key
0184  * @root: trees to search for @key
0185  * @ops: operators defining the node order
0186  *
0187  * Does a lockless lookup in the trees @root for the node matching @key.
0188  *
0189  * It is assumed that this is called while holding the appropriate RCU read
0190  * side lock.
0191  *
0192  * If the operators define a partial order on the elements (there are multiple
0193  * elements which have the same key value) it is undefined which of these
0194  * elements will be found. Nor is it possible to iterate the tree to find
0195  * further elements with the same key value.
0196  *
0197  * Returns: a pointer to the node matching @key or NULL.
0198  */
0199 static __always_inline struct latch_tree_node *
0200 latch_tree_find(void *key, struct latch_tree_root *root,
0201         const struct latch_tree_ops *ops)
0202 {
0203     struct latch_tree_node *node;
0204     unsigned int seq;
0205
0206     do {
0207         seq = raw_read_seqcount_latch(&root->seq);
0208         node = __lt_find(key, root, seq & 1, ops->comp);
0209     } while (read_seqcount_latch_retry(&root->seq, seq));
0210
0211     return node;
0212 }
0213
0214 #endif /* RB_TREE_LATCH_H */