xref: /netbsd-src/sys/kern/subr_thmap.c (revision 53b02e147d4ed531c0d2a5ca9b3e8026ba3e99b5)
1 /*	$NetBSD: subr_thmap.c,v 1.7 2020/08/31 20:22:57 riastradh Exp $	*/
2 
3 /*-
4  * Copyright (c) 2018 Mindaugas Rasiukevicius <rmind at noxt eu>
5  * All rights reserved.
6  *
7  * Redistribution and use in source and binary forms, with or without
8  * modification, are permitted provided that the following conditions
9  * are met:
10  * 1. Redistributions of source code must retain the above copyright
11  *    notice, this list of conditions and the following disclaimer.
12  * 2. Redistributions in binary form must reproduce the above copyright
13  *    notice, this list of conditions and the following disclaimer in the
14  *    documentation and/or other materials provided with the distribution.
15  *
16  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
17  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
19  * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
20  * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
21  * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
22  * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
23  * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
24  * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
25  * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
26  * SUCH DAMAGE.
27  *
28  * Upstream: https://github.com/rmind/thmap/
29  */
30 
31 /*
32  * Concurrent trie-hash map.
33  *
34  * The data structure is conceptually a radix trie on hashed keys.
35  * Keys are hashed using a 32-bit function.  The root level is a special
36  * case: it is managed using the compare-and-swap (CAS) atomic operation
37  * and has a fanout of 64.  The subsequent levels are constructed using
38  * intermediate nodes with a fanout of 16 (using 4 bits).  As more levels
39  * are created, more blocks of the 32-bit hash value might be generated
40  * by incrementing the seed parameter of the hash function.
41  *
42  * Concurrency
43  *
44  * - READERS: Descending is simply walking through the slot values of
45  *   the intermediate nodes.  It is lock-free as there is no intermediate
46  *   state: the slot is either empty or has a pointer to the child node.
47  *   The main assumptions here are the following:
48  *
49  *   i) modifications must preserve consistency with the respect to the
50  *   readers i.e. the readers can only see the valid node values;
51  *
52  *   ii) any invalid view must "fail" the reads, e.g. by making them
53  *   re-try from the root; this is a case for deletions and is achieved
54  *   using the NODE_DELETED flag.
55  *
56  *   iii) the node destruction must be synchronized with the readers,
57  *   e.g. by using the Epoch-based reclamation or other techniques.
58  *
59  * - WRITERS AND LOCKING: Each intermediate node has a spin-lock (which
60  *   is implemented using the NODE_LOCKED bit) -- it provides mutual
61  *   exclusion amongst concurrent writers.  The lock order for the nodes
62  *   is "bottom-up" i.e. they are locked as we ascend the trie.  A key
63  *   constraint here is that parent pointer never changes.
64  *
65  * - DELETES: In addition to writer's locking, the deletion keeps the
66  *   intermediate nodes in a valid state and sets the NODE_DELETED flag,
67  *   to indicate that the readers must re-start the walk from the root.
68  *   As the levels are collapsed, NODE_DELETED gets propagated up-tree.
69  *   The leaf nodes just stay as-is until they are reclaimed.
70  *
71  * - ROOT LEVEL: The root level is a special case, as it is implemented
72  *   as an array (rather than intermediate node).  The root-level slot can
73  *   only be set using CAS and it can only be set to a valid intermediate
74  *   node.  The root-level slot can only be cleared when the node it points
75  *   at becomes empty, is locked and marked as NODE_DELETED (this causes
76  *   the insert/delete operations to re-try until the slot is set to NULL).
77  *
78  * References:
79  *
80  *	W. Litwin, 1981, Trie Hashing.
81  *	Proceedings of the 1981 ACM SIGMOD, p. 19-29
82  *	https://dl.acm.org/citation.cfm?id=582322
83  *
84  *	P. L. Lehman and S. B. Yao.
85  *	Efficient locking for concurrent operations on B-trees.
86  *	ACM TODS, 6(4):650-670, 1981
87  *	https://www.csd.uoc.gr/~hy460/pdf/p650-lehman.pdf
88  */
89 
90 #ifdef _KERNEL
91 #include <sys/cdefs.h>
92 #include <sys/param.h>
93 #include <sys/types.h>
94 #include <sys/thmap.h>
95 #include <sys/kmem.h>
96 #include <sys/lock.h>
97 #include <sys/atomic.h>
98 #include <sys/hash.h>
99 #include <sys/cprng.h>
100 #define THMAP_RCSID(a) __KERNEL_RCSID(0, a)
101 #else
102 #include <stdio.h>
103 #include <stdlib.h>
104 #include <stdbool.h>
105 #include <stddef.h>
106 #include <inttypes.h>
107 #include <string.h>
108 #include <limits.h>
109 #define THMAP_RCSID(a) __RCSID(a)
110 
111 #include "thmap.h"
112 #include "utils.h"
113 #endif
114 
115 THMAP_RCSID("$NetBSD: subr_thmap.c,v 1.7 2020/08/31 20:22:57 riastradh Exp $");
116 
117 #include <crypto/blake2/blake2s.h>
118 
119 /*
120  * NetBSD kernel wrappers
121  */
122 #ifdef _KERNEL
123 #define	ASSERT KASSERT
124 #define	atomic_thread_fence(x) membar_sync()
125 #define	atomic_compare_exchange_weak_explicit_32(p, e, n, m1, m2) \
126     (atomic_cas_32((p), *(e), (n)) == *(e))
127 #define	atomic_compare_exchange_weak_explicit_ptr(p, e, n, m1, m2) \
128     (atomic_cas_ptr((p), *(void **)(e), (void *)(n)) == *(void **)(e))
129 #define	atomic_exchange_explicit(o, n, m1) atomic_swap_ptr((o), (n))
130 #define	murmurhash3 murmurhash2
131 #endif
132 
133 /*
134  * The root level fanout is 64 (indexed by the last 6 bits of the hash
135  * value XORed with the length).  Each subsequent level, represented by
136  * intermediate nodes, has a fanout of 16 (using 4 bits).
137  *
138  * The hash function produces 32-bit values.
139  */
140 
141 #define	HASHVAL_SEEDLEN	(16)
142 #define	HASHVAL_BITS	(32)
143 #define	HASHVAL_MOD	(HASHVAL_BITS - 1)
144 #define	HASHVAL_SHIFT	(5)
145 
146 #define	ROOT_BITS	(6)
147 #define	ROOT_SIZE	(1 << ROOT_BITS)
148 #define	ROOT_MASK	(ROOT_SIZE - 1)
149 #define	ROOT_MSBITS	(HASHVAL_BITS - ROOT_BITS)
150 
151 #define	LEVEL_BITS	(4)
152 #define	LEVEL_SIZE	(1 << LEVEL_BITS)
153 #define	LEVEL_MASK	(LEVEL_SIZE - 1)
154 
155 /*
156  * Instead of raw pointers, we use offsets from the base address.
157  * This accommodates the use of this data structure in shared memory,
158  * where mappings can be in different address spaces.
159  *
160  * The pointers must be aligned, since pointer tagging is used to
161  * differentiate the intermediate nodes from leaves.  We reserve the
162  * least significant bit.
163  */
164 typedef uintptr_t thmap_ptr_t;
165 typedef uintptr_t atomic_thmap_ptr_t;			// C11 _Atomic
166 
167 #define	THMAP_NULL		((thmap_ptr_t)0)
168 
169 #define	THMAP_LEAF_BIT		(0x1)
170 
171 #define	THMAP_ALIGNED_P(p)	(((uintptr_t)(p) & 3) == 0)
172 #define	THMAP_ALIGN(p)		((uintptr_t)(p) & ~(uintptr_t)3)
173 #define	THMAP_INODE_P(p)	(((uintptr_t)(p) & THMAP_LEAF_BIT) == 0)
174 
175 #define	THMAP_GETPTR(th, p)	((void *)((th)->baseptr + (uintptr_t)(p)))
176 #define	THMAP_GETOFF(th, p)	((thmap_ptr_t)((uintptr_t)(p) - (th)->baseptr))
177 #define	THMAP_NODE(th, p)	THMAP_GETPTR(th, THMAP_ALIGN(p))
178 
179 /*
180  * State field.
181  */
182 
183 #define	NODE_LOCKED		(1U << 31)		// lock (writers)
184 #define	NODE_DELETED		(1U << 30)		// node deleted
185 #define	NODE_COUNT(s)		((s) & 0x3fffffff)	// slot count mask
186 
187 /*
188  * There are two types of nodes:
189  * - Intermediate nodes -- arrays pointing to another level or a leaf;
190  * - Leaves, which store a key-value pair.
191  */
192 
193 typedef struct {
194 	uint32_t		state;			// C11 _Atomic
195 	thmap_ptr_t		parent;
196 	atomic_thmap_ptr_t	slots[LEVEL_SIZE];
197 } thmap_inode_t;
198 
199 #define	THMAP_INODE_LEN	sizeof(thmap_inode_t)
200 
201 typedef struct {
202 	thmap_ptr_t	key;
203 	size_t		len;
204 	void *		val;
205 } thmap_leaf_t;
206 
207 typedef struct {
208 	const uint8_t *	seed;		// secret seed
209 	unsigned	rslot;		// root-level slot index
210 	unsigned	level;		// current level in the tree
211 	unsigned	hashidx;	// current hash index (block of bits)
212 	uint32_t	hashval;	// current hash value
213 } thmap_query_t;
214 
215 typedef struct {
216 	uintptr_t	addr;
217 	size_t		len;
218 	void *		next;
219 } thmap_gc_t;
220 
221 #define	THMAP_ROOT_LEN	(sizeof(thmap_ptr_t) * ROOT_SIZE)
222 
223 struct thmap {
224 	uintptr_t		baseptr;
225 	atomic_thmap_ptr_t *	root;
226 	unsigned		flags;
227 	const thmap_ops_t *	ops;
228 	thmap_gc_t *		gc_list;		// C11 _Atomic
229 	uint8_t			seed[HASHVAL_SEEDLEN];
230 };
231 
232 static void	stage_mem_gc(thmap_t *, uintptr_t, size_t);
233 
234 /*
235  * A few low-level helper routines.
236  */
237 
238 static uintptr_t
239 alloc_wrapper(size_t len)
240 {
241 	return (uintptr_t)kmem_intr_alloc(len, KM_NOSLEEP);
242 }
243 
244 static void
245 free_wrapper(uintptr_t addr, size_t len)
246 {
247 	kmem_intr_free((void *)addr, len);
248 }
249 
250 static const thmap_ops_t thmap_default_ops = {
251 	.alloc = alloc_wrapper,
252 	.free = free_wrapper
253 };
254 
255 /*
256  * NODE LOCKING.
257  */
258 
259 #ifdef DIAGNOSTIC
260 static inline bool
261 node_locked_p(thmap_inode_t *node)
262 {
263 	return (atomic_load_relaxed(&node->state) & NODE_LOCKED) != 0;
264 }
265 #endif
266 
267 static void
268 lock_node(thmap_inode_t *node)
269 {
270 	unsigned bcount = SPINLOCK_BACKOFF_MIN;
271 	uint32_t s;
272 again:
273 	s = atomic_load_relaxed(&node->state);
274 	if (s & NODE_LOCKED) {
275 		SPINLOCK_BACKOFF(bcount);
276 		goto again;
277 	}
278 	/* Acquire from prior release in unlock_node.() */
279 	if (!atomic_compare_exchange_weak_explicit_32(&node->state,
280 	    &s, s | NODE_LOCKED, memory_order_acquire, memory_order_relaxed)) {
281 		bcount = SPINLOCK_BACKOFF_MIN;
282 		goto again;
283 	}
284 }
285 
286 static void
287 unlock_node(thmap_inode_t *node)
288 {
289 	uint32_t s = atomic_load_relaxed(&node->state) & ~NODE_LOCKED;
290 
291 	ASSERT(node_locked_p(node));
292 	/* Release to subsequent acquire in lock_node(). */
293 	atomic_store_release(&node->state, s);
294 }
295 
296 /*
297  * HASH VALUE AND KEY OPERATIONS.
298  */
299 
300 static inline uint32_t
301 hash(const uint8_t seed[static HASHVAL_SEEDLEN], const void *key, size_t len,
302     uint32_t level)
303 {
304 	struct blake2s B;
305 	uint32_t h;
306 
307 	if (level == 0)
308 		return murmurhash3(key, len, 0);
309 
310 	/*
311 	 * Byte order is not significant here because this is
312 	 * intentionally secret and independent for each thmap.
313 	 *
314 	 * XXX We get 32 bytes of output at a time; we could march
315 	 * through them sequentially rather than throwing away 28 bytes
316 	 * and recomputing BLAKE2 each time.  But the number of
317 	 * iterations ought to be geometric in the collision
318 	 * probability at each level which should be very small anyway.
319 	 */
320 	blake2s_init(&B, sizeof h, seed, HASHVAL_SEEDLEN);
321 	blake2s_update(&B, &level, sizeof level);
322 	blake2s_update(&B, key, len);
323 	blake2s_final(&B, &h);
324 
325 	return h;
326 }
327 
328 static inline void
329 hashval_init(thmap_query_t *query, const uint8_t seed[static HASHVAL_SEEDLEN],
330     const void * restrict key, size_t len)
331 {
332 	const uint32_t hashval = hash(seed, key, len, 0);
333 
334 	query->seed = seed;
335 	query->rslot = ((hashval >> ROOT_MSBITS) ^ len) & ROOT_MASK;
336 	query->level = 0;
337 	query->hashval = hashval;
338 	query->hashidx = 0;
339 }
340 
341 /*
342  * hashval_getslot: given the key, compute the hash (if not already cached)
343  * and return the offset for the current level.
344  */
345 static unsigned
346 hashval_getslot(thmap_query_t *query, const void * restrict key, size_t len)
347 {
348 	const unsigned offset = query->level * LEVEL_BITS;
349 	const unsigned shift = offset & HASHVAL_MOD;
350 	const unsigned i = offset >> HASHVAL_SHIFT;
351 
352 	if (query->hashidx != i) {
353 		/* Generate a hash value for a required range. */
354 		query->hashval = hash(query->seed, key, len, i);
355 		query->hashidx = i;
356 	}
357 	return (query->hashval >> shift) & LEVEL_MASK;
358 }
359 
360 static unsigned
361 hashval_getleafslot(const thmap_t *thmap,
362     const thmap_leaf_t *leaf, unsigned level)
363 {
364 	const void *key = THMAP_GETPTR(thmap, leaf->key);
365 	const unsigned offset = level * LEVEL_BITS;
366 	const unsigned shift = offset & HASHVAL_MOD;
367 	const unsigned i = offset >> HASHVAL_SHIFT;
368 
369 	return (hash(thmap->seed, key, leaf->len, i) >> shift) & LEVEL_MASK;
370 }
371 
372 static inline unsigned
373 hashval_getl0slot(const thmap_t *thmap, const thmap_query_t *query,
374     const thmap_leaf_t *leaf)
375 {
376 	if (__predict_true(query->hashidx == 0)) {
377 		return query->hashval & LEVEL_MASK;
378 	}
379 	return hashval_getleafslot(thmap, leaf, 0);
380 }
381 
382 static bool
383 key_cmp_p(const thmap_t *thmap, const thmap_leaf_t *leaf,
384     const void * restrict key, size_t len)
385 {
386 	const void *leafkey = THMAP_GETPTR(thmap, leaf->key);
387 	return len == leaf->len && memcmp(key, leafkey, len) == 0;
388 }
389 
390 /*
391  * INTER-NODE OPERATIONS.
392  */
393 
394 static thmap_inode_t *
395 node_create(thmap_t *thmap, thmap_inode_t *parent)
396 {
397 	thmap_inode_t *node;
398 	uintptr_t p;
399 
400 	p = thmap->ops->alloc(THMAP_INODE_LEN);
401 	if (!p) {
402 		return NULL;
403 	}
404 	node = THMAP_GETPTR(thmap, p);
405 	ASSERT(THMAP_ALIGNED_P(node));
406 
407 	memset(node, 0, THMAP_INODE_LEN);
408 	if (parent) {
409 		/* Not yet published, no need for ordering. */
410 		atomic_store_relaxed(&node->state, NODE_LOCKED);
411 		node->parent = THMAP_GETOFF(thmap, parent);
412 	}
413 	return node;
414 }
415 
416 static void
417 node_insert(thmap_inode_t *node, unsigned slot, thmap_ptr_t child)
418 {
419 	ASSERT(node_locked_p(node) || node->parent == THMAP_NULL);
420 	ASSERT((atomic_load_relaxed(&node->state) & NODE_DELETED) == 0);
421 	ASSERT(atomic_load_relaxed(&node->slots[slot]) == THMAP_NULL);
422 
423 	ASSERT(NODE_COUNT(atomic_load_relaxed(&node->state)) < LEVEL_SIZE);
424 
425 	/*
426 	 * If node is public already, caller is responsible for issuing
427 	 * release fence; if node is not public, no ordering is needed.
428 	 * Hence relaxed ordering.
429 	 */
430 	atomic_store_relaxed(&node->slots[slot], child);
431 	atomic_store_relaxed(&node->state,
432 	    atomic_load_relaxed(&node->state) + 1);
433 }
434 
435 static void
436 node_remove(thmap_inode_t *node, unsigned slot)
437 {
438 	ASSERT(node_locked_p(node));
439 	ASSERT((atomic_load_relaxed(&node->state) & NODE_DELETED) == 0);
440 	ASSERT(atomic_load_relaxed(&node->slots[slot]) != THMAP_NULL);
441 
442 	ASSERT(NODE_COUNT(atomic_load_relaxed(&node->state)) > 0);
443 	ASSERT(NODE_COUNT(atomic_load_relaxed(&node->state)) <= LEVEL_SIZE);
444 
445 	/* Element will be GC-ed later; no need for ordering here. */
446 	atomic_store_relaxed(&node->slots[slot], THMAP_NULL);
447 	atomic_store_relaxed(&node->state,
448 	    atomic_load_relaxed(&node->state) - 1);
449 }
450 
451 /*
452  * LEAF OPERATIONS.
453  */
454 
455 static thmap_leaf_t *
456 leaf_create(const thmap_t *thmap, const void *key, size_t len, void *val)
457 {
458 	thmap_leaf_t *leaf;
459 	uintptr_t leaf_off, key_off;
460 
461 	leaf_off = thmap->ops->alloc(sizeof(thmap_leaf_t));
462 	if (!leaf_off) {
463 		return NULL;
464 	}
465 	leaf = THMAP_GETPTR(thmap, leaf_off);
466 	ASSERT(THMAP_ALIGNED_P(leaf));
467 
468 	if ((thmap->flags & THMAP_NOCOPY) == 0) {
469 		/*
470 		 * Copy the key.
471 		 */
472 		key_off = thmap->ops->alloc(len);
473 		if (!key_off) {
474 			thmap->ops->free(leaf_off, sizeof(thmap_leaf_t));
475 			return NULL;
476 		}
477 		memcpy(THMAP_GETPTR(thmap, key_off), key, len);
478 		leaf->key = key_off;
479 	} else {
480 		/* Otherwise, we use a reference. */
481 		leaf->key = (uintptr_t)key;
482 	}
483 	leaf->len = len;
484 	leaf->val = val;
485 	return leaf;
486 }
487 
488 static void
489 leaf_free(const thmap_t *thmap, thmap_leaf_t *leaf)
490 {
491 	if ((thmap->flags & THMAP_NOCOPY) == 0) {
492 		thmap->ops->free(leaf->key, leaf->len);
493 	}
494 	thmap->ops->free(THMAP_GETOFF(thmap, leaf), sizeof(thmap_leaf_t));
495 }
496 
497 static thmap_leaf_t *
498 get_leaf(const thmap_t *thmap, thmap_inode_t *parent, unsigned slot)
499 {
500 	thmap_ptr_t node;
501 
502 	/* Consume from prior release in thmap_put(). */
503 	node = atomic_load_consume(&parent->slots[slot]);
504 	if (THMAP_INODE_P(node)) {
505 		return NULL;
506 	}
507 	return THMAP_NODE(thmap, node);
508 }
509 
510 /*
511  * ROOT OPERATIONS.
512  */
513 
514 /*
515  * root_try_put: Try to set a root pointer at query->rslot.
516  *
517  * => Implies release operation on success.
518  * => Implies no ordering on failure.
519  */
520 static inline bool
521 root_try_put(thmap_t *thmap, const thmap_query_t *query, thmap_leaf_t *leaf)
522 {
523 	thmap_ptr_t expected;
524 	const unsigned i = query->rslot;
525 	thmap_inode_t *node;
526 	thmap_ptr_t nptr;
527 	unsigned slot;
528 
529 	/*
530 	 * Must pre-check first.  No ordering required because we will
531 	 * check again before taking any actions, and start over if
532 	 * this changes from null.
533 	 */
534 	if (atomic_load_relaxed(&thmap->root[i])) {
535 		return false;
536 	}
537 
538 	/*
539 	 * Create an intermediate node.  Since there is no parent set,
540 	 * it will be created unlocked and the CAS operation will
541 	 * release it to readers.
542 	 */
543 	node = node_create(thmap, NULL);
544 	slot = hashval_getl0slot(thmap, query, leaf);
545 	node_insert(node, slot, THMAP_GETOFF(thmap, leaf) | THMAP_LEAF_BIT);
546 	nptr = THMAP_GETOFF(thmap, node);
547 again:
548 	if (atomic_load_relaxed(&thmap->root[i])) {
549 		thmap->ops->free(nptr, THMAP_INODE_LEN);
550 		return false;
551 	}
552 	/* Release to subsequent consume in find_edge_node(). */
553 	expected = THMAP_NULL;
554 	if (!atomic_compare_exchange_weak_explicit_ptr(&thmap->root[i], &expected,
555 	    nptr, memory_order_release, memory_order_relaxed)) {
556 		goto again;
557 	}
558 	return true;
559 }
560 
561 /*
562  * find_edge_node: given the hash, traverse the tree to find the edge node.
563  *
564  * => Returns an aligned (clean) pointer to the parent node.
565  * => Returns the slot number and sets current level.
566  */
567 static thmap_inode_t *
568 find_edge_node(const thmap_t *thmap, thmap_query_t *query,
569     const void * restrict key, size_t len, unsigned *slot)
570 {
571 	thmap_ptr_t root_slot;
572 	thmap_inode_t *parent;
573 	thmap_ptr_t node;
574 	unsigned off;
575 
576 	ASSERT(query->level == 0);
577 
578 	/* Consume from prior release in root_try_put(). */
579 	root_slot = atomic_load_consume(&thmap->root[query->rslot]);
580 	parent = THMAP_NODE(thmap, root_slot);
581 	if (!parent) {
582 		return NULL;
583 	}
584 descend:
585 	off = hashval_getslot(query, key, len);
586 	/* Consume from prior release in thmap_put(). */
587 	node = atomic_load_consume(&parent->slots[off]);
588 
589 	/* Descend the tree until we find a leaf or empty slot. */
590 	if (node && THMAP_INODE_P(node)) {
591 		parent = THMAP_NODE(thmap, node);
592 		query->level++;
593 		goto descend;
594 	}
595 	/*
596 	 * NODE_DELETED does not become stale until GC runs, which
597 	 * cannot happen while we are in the middle of an operation,
598 	 * hence relaxed ordering.
599 	 */
600 	if (atomic_load_relaxed(&parent->state) & NODE_DELETED) {
601 		return NULL;
602 	}
603 	*slot = off;
604 	return parent;
605 }
606 
607 /*
608  * find_edge_node_locked: traverse the tree, like find_edge_node(),
609  * but attempt to lock the edge node.
610  *
611  * => Returns NULL if the deleted node is found.  This indicates that
612  *    the caller must re-try from the root, as the root slot might have
613  *    changed too.
614  */
615 static thmap_inode_t *
616 find_edge_node_locked(const thmap_t *thmap, thmap_query_t *query,
617     const void * restrict key, size_t len, unsigned *slot)
618 {
619 	thmap_inode_t *node;
620 	thmap_ptr_t target;
621 retry:
622 	/*
623 	 * Find the edge node and lock it!  Re-check the state since
624 	 * the tree might change by the time we acquire the lock.
625 	 */
626 	node = find_edge_node(thmap, query, key, len, slot);
627 	if (!node) {
628 		/* The root slot is empty -- let the caller decide. */
629 		query->level = 0;
630 		return NULL;
631 	}
632 	lock_node(node);
633 	if (__predict_false(atomic_load_relaxed(&node->state) & NODE_DELETED)) {
634 		/*
635 		 * The node has been deleted.  The tree might have a new
636 		 * shape now, therefore we must re-start from the root.
637 		 */
638 		unlock_node(node);
639 		query->level = 0;
640 		return NULL;
641 	}
642 	target = atomic_load_relaxed(&node->slots[*slot]);
643 	if (__predict_false(target && THMAP_INODE_P(target))) {
644 		/*
645 		 * The target slot has been changed and it is now an
646 		 * intermediate node.  Re-start from the top internode.
647 		 */
648 		unlock_node(node);
649 		query->level = 0;
650 		goto retry;
651 	}
652 	return node;
653 }
654 
655 /*
656  * thmap_get: lookup a value given the key.
657  */
658 void *
659 thmap_get(thmap_t *thmap, const void *key, size_t len)
660 {
661 	thmap_query_t query;
662 	thmap_inode_t *parent;
663 	thmap_leaf_t *leaf;
664 	unsigned slot;
665 
666 	hashval_init(&query, thmap->seed, key, len);
667 	parent = find_edge_node(thmap, &query, key, len, &slot);
668 	if (!parent) {
669 		return NULL;
670 	}
671 	leaf = get_leaf(thmap, parent, slot);
672 	if (!leaf) {
673 		return NULL;
674 	}
675 	if (!key_cmp_p(thmap, leaf, key, len)) {
676 		return NULL;
677 	}
678 	return leaf->val;
679 }
680 
681 /*
682  * thmap_put: insert a value given the key.
683  *
684  * => If the key is already present, return the associated value.
685  * => Otherwise, on successful insert, return the given value.
686  */
687 void *
688 thmap_put(thmap_t *thmap, const void *key, size_t len, void *val)
689 {
690 	thmap_query_t query;
691 	thmap_leaf_t *leaf, *other;
692 	thmap_inode_t *parent, *child;
693 	unsigned slot, other_slot;
694 	thmap_ptr_t target;
695 
696 	/*
697 	 * First, pre-allocate and initialize the leaf node.
698 	 */
699 	leaf = leaf_create(thmap, key, len, val);
700 	if (__predict_false(!leaf)) {
701 		return NULL;
702 	}
703 	hashval_init(&query, thmap->seed, key, len);
704 retry:
705 	/*
706 	 * Try to insert into the root first, if its slot is empty.
707 	 */
708 	if (root_try_put(thmap, &query, leaf)) {
709 		/* Success: the leaf was inserted; no locking involved. */
710 		return val;
711 	}
712 
713 	/*
714 	 * Release node via store in node_insert (*) to subsequent
715 	 * consume in get_leaf() or find_edge_node().
716 	 */
717 	atomic_thread_fence(memory_order_release);
718 
719 	/*
720 	 * Find the edge node and the target slot.
721 	 */
722 	parent = find_edge_node_locked(thmap, &query, key, len, &slot);
723 	if (!parent) {
724 		goto retry;
725 	}
726 	target = atomic_load_relaxed(&parent->slots[slot]); // tagged offset
727 	if (THMAP_INODE_P(target)) {
728 		/*
729 		 * Empty slot: simply insert the new leaf.  The release
730 		 * fence is already issued for us.
731 		 */
732 		target = THMAP_GETOFF(thmap, leaf) | THMAP_LEAF_BIT;
733 		node_insert(parent, slot, target); /* (*) */
734 		goto out;
735 	}
736 
737 	/*
738 	 * Collision or duplicate.
739 	 */
740 	other = THMAP_NODE(thmap, target);
741 	if (key_cmp_p(thmap, other, key, len)) {
742 		/*
743 		 * Duplicate.  Free the pre-allocated leaf and
744 		 * return the present value.
745 		 */
746 		leaf_free(thmap, leaf);
747 		val = other->val;
748 		goto out;
749 	}
750 descend:
751 	/*
752 	 * Collision -- expand the tree.  Create an intermediate node
753 	 * which will be locked (NODE_LOCKED) for us.  At this point,
754 	 * we advance to the next level.
755 	 */
756 	child = node_create(thmap, parent);
757 	if (__predict_false(!child)) {
758 		leaf_free(thmap, leaf);
759 		val = NULL;
760 		goto out;
761 	}
762 	query.level++;
763 
764 	/*
765 	 * Insert the other (colliding) leaf first.  The new child is
766 	 * not yet published, so memory order is relaxed.
767 	 */
768 	other_slot = hashval_getleafslot(thmap, other, query.level);
769 	target = THMAP_GETOFF(thmap, other) | THMAP_LEAF_BIT;
770 	node_insert(child, other_slot, target);
771 
772 	/*
773 	 * Insert the intermediate node into the parent node.
774 	 * It becomes the new parent for the our new leaf.
775 	 *
776 	 * Ensure that stores to the child (and leaf) reach global
777 	 * visibility before it gets inserted to the parent, as
778 	 * consumed by get_leaf() or find_edge_node().
779 	 */
780 	atomic_store_release(&parent->slots[slot], THMAP_GETOFF(thmap, child));
781 
782 	unlock_node(parent);
783 	ASSERT(node_locked_p(child));
784 	parent = child;
785 
786 	/*
787 	 * Get the new slot and check for another collision
788 	 * at the next level.
789 	 */
790 	slot = hashval_getslot(&query, key, len);
791 	if (slot == other_slot) {
792 		/* Another collision -- descend and expand again. */
793 		goto descend;
794 	}
795 
796 	/*
797 	 * Insert our new leaf once we expanded enough.  The release
798 	 * fence is already issued for us.
799 	 */
800 	target = THMAP_GETOFF(thmap, leaf) | THMAP_LEAF_BIT;
801 	node_insert(parent, slot, target); /* (*) */
802 out:
803 	unlock_node(parent);
804 	return val;
805 }
806 
807 /*
808  * thmap_del: remove the entry given the key.
809  */
810 void *
811 thmap_del(thmap_t *thmap, const void *key, size_t len)
812 {
813 	thmap_query_t query;
814 	thmap_leaf_t *leaf;
815 	thmap_inode_t *parent;
816 	unsigned slot;
817 	void *val;
818 
819 	hashval_init(&query, thmap->seed, key, len);
820 	parent = find_edge_node_locked(thmap, &query, key, len, &slot);
821 	if (!parent) {
822 		/* Root slot empty: not found. */
823 		return NULL;
824 	}
825 	leaf = get_leaf(thmap, parent, slot);
826 	if (!leaf || !key_cmp_p(thmap, leaf, key, len)) {
827 		/* Not found. */
828 		unlock_node(parent);
829 		return NULL;
830 	}
831 
832 	/* Remove the leaf. */
833 	ASSERT(THMAP_NODE(thmap, atomic_load_relaxed(&parent->slots[slot]))
834 	    == leaf);
835 	node_remove(parent, slot);
836 
837 	/*
838 	 * Collapse the levels if removing the last item.
839 	 */
840 	while (query.level &&
841 	    NODE_COUNT(atomic_load_relaxed(&parent->state)) == 0) {
842 		thmap_inode_t *node = parent;
843 
844 		ASSERT(atomic_load_relaxed(&node->state) == NODE_LOCKED);
845 
846 		/*
847 		 * Ascend one level up.
848 		 * => Mark our current parent as deleted.
849 		 * => Lock the parent one level up.
850 		 */
851 		query.level--;
852 		slot = hashval_getslot(&query, key, len);
853 		parent = THMAP_NODE(thmap, node->parent);
854 		ASSERT(parent != NULL);
855 
856 		lock_node(parent);
857 		ASSERT((atomic_load_relaxed(&parent->state) & NODE_DELETED)
858 		    == 0);
859 
860 		/*
861 		 * Lock is exclusive, so nobody else can be writing at
862 		 * the same time, and no need for atomic R/M/W, but
863 		 * readers may read without the lock and so need atomic
864 		 * load/store.  No ordering here needed because the
865 		 * entry itself stays valid until GC.
866 		 */
867 		atomic_store_relaxed(&node->state,
868 		    atomic_load_relaxed(&node->state) | NODE_DELETED);
869 		unlock_node(node); // memory_order_release
870 
871 		ASSERT(THMAP_NODE(thmap,
872 		    atomic_load_relaxed(&parent->slots[slot])) == node);
873 		node_remove(parent, slot);
874 
875 		/* Stage the removed node for G/C. */
876 		stage_mem_gc(thmap, THMAP_GETOFF(thmap, node), THMAP_INODE_LEN);
877 	}
878 
879 	/*
880 	 * If the top node is empty, then we need to remove it from the
881 	 * root level.  Mark the node as deleted and clear the slot.
882 	 *
883 	 * Note: acquiring the lock on the top node effectively prevents
884 	 * the root slot from changing.
885 	 */
886 	if (NODE_COUNT(atomic_load_relaxed(&parent->state)) == 0) {
887 		const unsigned rslot = query.rslot;
888 		const thmap_ptr_t nptr =
889 		    atomic_load_relaxed(&thmap->root[rslot]);
890 
891 		ASSERT(query.level == 0);
892 		ASSERT(parent->parent == THMAP_NULL);
893 		ASSERT(THMAP_GETOFF(thmap, parent) == nptr);
894 
895 		/* Mark as deleted and remove from the root-level slot. */
896 		atomic_store_relaxed(&parent->state,
897 		    atomic_load_relaxed(&parent->state) | NODE_DELETED);
898 		atomic_store_relaxed(&thmap->root[rslot], THMAP_NULL);
899 
900 		stage_mem_gc(thmap, nptr, THMAP_INODE_LEN);
901 	}
902 	unlock_node(parent);
903 
904 	/*
905 	 * Save the value and stage the leaf for G/C.
906 	 */
907 	val = leaf->val;
908 	if ((thmap->flags & THMAP_NOCOPY) == 0) {
909 		stage_mem_gc(thmap, leaf->key, leaf->len);
910 	}
911 	stage_mem_gc(thmap, THMAP_GETOFF(thmap, leaf), sizeof(thmap_leaf_t));
912 	return val;
913 }
914 
915 /*
916  * G/C routines.
917  */
918 
919 static void
920 stage_mem_gc(thmap_t *thmap, uintptr_t addr, size_t len)
921 {
922 	thmap_gc_t *head, *gc;
923 
924 	gc = kmem_intr_alloc(sizeof(thmap_gc_t), KM_NOSLEEP);
925 	gc->addr = addr;
926 	gc->len = len;
927 retry:
928 	head = atomic_load_relaxed(&thmap->gc_list);
929 	gc->next = head; // not yet published
930 
931 	/* Release to subsequent acquire in thmap_stage_gc(). */
932 	if (!atomic_compare_exchange_weak_explicit_ptr(&thmap->gc_list, &head, gc,
933 	    memory_order_release, memory_order_relaxed)) {
934 		goto retry;
935 	}
936 }
937 
938 void *
939 thmap_stage_gc(thmap_t *thmap)
940 {
941 	/* Acquire from prior release in stage_mem_gc(). */
942 	return atomic_exchange_explicit(&thmap->gc_list, NULL,
943 	    memory_order_acquire);
944 }
945 
946 void
947 thmap_gc(thmap_t *thmap, void *ref)
948 {
949 	thmap_gc_t *gc = ref;
950 
951 	while (gc) {
952 		thmap_gc_t *next = gc->next;
953 		thmap->ops->free(gc->addr, gc->len);
954 		kmem_intr_free(gc, sizeof(thmap_gc_t));
955 		gc = next;
956 	}
957 }
958 
959 /*
960  * thmap_create: construct a new trie-hash map object.
961  */
962 thmap_t *
963 thmap_create(uintptr_t baseptr, const thmap_ops_t *ops, unsigned flags)
964 {
965 	thmap_t *thmap;
966 	uintptr_t root;
967 
968 	/*
969 	 * Setup the map object.
970 	 */
971 	if (!THMAP_ALIGNED_P(baseptr)) {
972 		return NULL;
973 	}
974 	thmap = kmem_zalloc(sizeof(thmap_t), KM_SLEEP);
975 	if (!thmap) {
976 		return NULL;
977 	}
978 	thmap->baseptr = baseptr;
979 	thmap->ops = ops ? ops : &thmap_default_ops;
980 	thmap->flags = flags;
981 
982 	if ((thmap->flags & THMAP_SETROOT) == 0) {
983 		/* Allocate the root level. */
984 		root = thmap->ops->alloc(THMAP_ROOT_LEN);
985 		thmap->root = THMAP_GETPTR(thmap, root);
986 		if (!thmap->root) {
987 			kmem_free(thmap, sizeof(thmap_t));
988 			return NULL;
989 		}
990 		memset(thmap->root, 0, THMAP_ROOT_LEN);
991 		atomic_thread_fence(memory_order_release); /* XXX */
992 	}
993 
994 	cprng_strong(kern_cprng, thmap->seed, sizeof thmap->seed, 0);
995 
996 	return thmap;
997 }
998 
999 int
1000 thmap_setroot(thmap_t *thmap, uintptr_t root_off)
1001 {
1002 	if (thmap->root) {
1003 		return -1;
1004 	}
1005 	thmap->root = THMAP_GETPTR(thmap, root_off);
1006 	atomic_thread_fence(memory_order_release); /* XXX */
1007 	return 0;
1008 }
1009 
1010 uintptr_t
1011 thmap_getroot(const thmap_t *thmap)
1012 {
1013 	return THMAP_GETOFF(thmap, thmap->root);
1014 }
1015 
1016 void
1017 thmap_destroy(thmap_t *thmap)
1018 {
1019 	uintptr_t root = THMAP_GETOFF(thmap, thmap->root);
1020 	void *ref;
1021 
1022 	ref = thmap_stage_gc(thmap);
1023 	thmap_gc(thmap, ref);
1024 
1025 	if ((thmap->flags & THMAP_SETROOT) == 0) {
1026 		thmap->ops->free(root, THMAP_ROOT_LEN);
1027 	}
1028 	kmem_free(thmap, sizeof(thmap_t));
1029 }
1030