From b8254de5f092f9b51c0a2537858813c59adc560f Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Thu, 24 Oct 2024 17:29:51 -0700
Subject: [PATCH v3 2/4] raidxtree.h: support shared iteration.

This commit supports a shared iteration operation on a radix tree with
multiple processes. The radix tree must be in shared mode to start a
shared itereation. Parallel workers can attach the shared iteration
using the iterator handle given by the leader process. Same as the
normal interation, it's guarnteed that the shared iteration returns
key-values in an ascending order.

Author:
Reviewed-by:
Discussion: https://postgr.es/m/
---
 src/include/lib/radixtree.h | 221 +++++++++++++++++++++++++++++++-----
 1 file changed, 190 insertions(+), 31 deletions(-)

diff --git a/src/include/lib/radixtree.h b/src/include/lib/radixtree.h
index 88bf695e3f..b93553200d 100644
--- a/src/include/lib/radixtree.h
+++ b/src/include/lib/radixtree.h
@@ -177,6 +177,9 @@
 #define RT_ATTACH RT_MAKE_NAME(attach)
 #define RT_DETACH RT_MAKE_NAME(detach)
 #define RT_GET_HANDLE RT_MAKE_NAME(get_handle)
+#define RT_BEGIN_ITERATE_SHARED RT_MAKE_NAME(begin_iterate_shared)
+#define RT_ATTACH_ITERATE_SHARED RT_MAKE_NAME(attach_iterate_shared)
+#define RT_GET_ITER_HANDLE RT_MAKE_NAME(get_iter_handle)
 #define RT_LOCK_EXCLUSIVE RT_MAKE_NAME(lock_exclusive)
 #define RT_LOCK_SHARE RT_MAKE_NAME(lock_share)
 #define RT_UNLOCK RT_MAKE_NAME(unlock)
@@ -236,15 +239,19 @@
 #define RT_SHRINK_NODE_16 RT_MAKE_NAME(shrink_child_16)
 #define RT_SHRINK_NODE_48 RT_MAKE_NAME(shrink_child_48)
 #define RT_SHRINK_NODE_256 RT_MAKE_NAME(shrink_child_256)
+#define RT_INITIALIZE_ITER RT_MAKE_NAME(initialize_iter)
 #define RT_NODE_ITERATE_NEXT RT_MAKE_NAME(node_iterate_next)
 #define RT_VERIFY_NODE RT_MAKE_NAME(verify_node)
 
 /* type declarations */
 #define RT_RADIX_TREE RT_MAKE_NAME(radix_tree)
 #define RT_RADIX_TREE_CONTROL RT_MAKE_NAME(radix_tree_control)
+#define RT_ITER_CONTROL RT_MAKE_NAME(iter_control)
 #define RT_ITER RT_MAKE_NAME(iter)
 #ifdef RT_SHMEM
 #define RT_HANDLE RT_MAKE_NAME(handle)
+#define RT_ITER_CONTROL_SHARED RT_MAKE_NAME(iter_control_shared)
+#define RT_ITER_HANDLE RT_MAKE_NAME(iter_handle)
 #endif
 #define RT_NODE RT_MAKE_NAME(node)
 #define RT_CHILD_PTR RT_MAKE_NAME(child_ptr)
@@ -270,6 +277,7 @@ typedef struct RT_ITER RT_ITER;
 
 #ifdef RT_SHMEM
 typedef dsa_pointer RT_HANDLE;
+typedef dsa_pointer RT_ITER_HANDLE;
 #endif
 
 #ifdef RT_SHMEM
@@ -687,6 +695,7 @@ typedef struct RT_RADIX_TREE_CONTROL
 	RT_HANDLE	handle;
 	uint32		magic;
 	LWLock		lock;
+	int			tranche_id;
 #endif
 
 	RT_PTR_ALLOC root;
@@ -740,11 +749,9 @@ typedef struct RT_NODE_ITER
 	int			idx;
 }			RT_NODE_ITER;
 
-/* state for iterating over the whole radix tree */
-struct RT_ITER
+/* Contain the iteration state data */
+typedef struct RT_ITER_CONTROL
 {
-	RT_RADIX_TREE *tree;
-
 	/*
 	 * A stack to track iteration for each level. Level 0 is the lowest (or
 	 * leaf) level
@@ -755,8 +762,36 @@ struct RT_ITER
 
 	/* The key constructed during iteration */
 	uint64		key;
-};
+}			RT_ITER_CONTROL;
+
+#ifdef RT_SHMEM
+/* Contain the shared iteration state data */
+typedef struct RT_ITER_CONTROL_SHARED
+{
+	/* Actual shared iteration state data */
+	RT_ITER_CONTROL common;
+
+	/* protect the control data */
+	LWLock		lock;
+
+	RT_ITER_HANDLE handle;
+	pg_atomic_uint32 refcnt;
+}			RT_ITER_CONTROL_SHARED;
+#endif
+
+/* state for iterating over the whole radix tree */
+struct RT_ITER
+{
+	RT_RADIX_TREE *tree;
 
+	/* pointing to either local memory or DSA */
+	RT_ITER_CONTROL *ctl;
+
+#ifdef RT_SHMEM
+	/* True if the iterator is for shared iteration */
+	bool		shared;
+#endif
+};
 
 /* verification (available only in assert-enabled builds) */
 static void RT_VERIFY_NODE(RT_NODE * node);
@@ -1848,6 +1883,7 @@ RT_CREATE(MemoryContext ctx)
 	tree->ctl = (RT_RADIX_TREE_CONTROL *) dsa_get_address(dsa, dp);
 	tree->ctl->handle = dp;
 	tree->ctl->magic = RT_RADIX_TREE_MAGIC;
+	tree->ctl->tranche_id = tranche_id;
 	LWLockInitialize(&tree->ctl->lock, tranche_id);
 #else
 	tree->ctl = (RT_RADIX_TREE_CONTROL *) palloc0(sizeof(RT_RADIX_TREE_CONTROL));
@@ -1900,6 +1936,9 @@ RT_ATTACH(dsa_area *dsa, RT_HANDLE handle)
 	dsa_pointer control;
 
 	tree = (RT_RADIX_TREE *) palloc0(sizeof(RT_RADIX_TREE));
+	tree->iter_context = AllocSetContextCreate(CurrentMemoryContext,
+											   RT_STR(RT_PREFIX) "_radix_tree iter context",
+											   ALLOCSET_SMALL_SIZES);
 
 	/* Find the control object in shared memory */
 	control = handle;
@@ -2072,35 +2111,86 @@ RT_FREE(RT_RADIX_TREE * tree)
 
 /***************** ITERATION *****************/
 
+/* Common routine to initialize the given iterator */
+static void
+RT_INITIALIZE_ITER(RT_RADIX_TREE * tree, RT_ITER * iter)
+{
+	RT_CHILD_PTR root;
+
+	iter->tree = tree;
+
+	Assert(RT_PTR_ALLOC_IS_VALID(tree->ctl->root));
+	root.alloc = iter->tree->ctl->root;
+	RT_PTR_SET_LOCAL(tree, &root);
+
+	iter->ctl->top_level = iter->tree->ctl->start_shift / RT_SPAN;
+
+	/* Set the root to start */
+	iter->ctl->cur_level = iter->ctl->top_level;
+	iter->ctl->node_iters[iter->ctl->cur_level].node = root;
+	iter->ctl->node_iters[iter->ctl->cur_level].idx = 0;
+}
+
 /*
  * Create and return the iterator for the given radix tree.
  *
- * Taking a lock in shared mode during the iteration is the caller's
- * responsibility.
+ * Taking a lock on a radix tree in shared mode during the iteration is the
+ * caller's responsibility.
  */
 RT_SCOPE	RT_ITER *
 RT_BEGIN_ITERATE(RT_RADIX_TREE * tree)
 {
 	RT_ITER    *iter;
-	RT_CHILD_PTR root;
 
 	iter = (RT_ITER *) MemoryContextAllocZero(tree->iter_context,
 											  sizeof(RT_ITER));
-	iter->tree = tree;
+	iter->ctl = (RT_ITER_CONTROL *) MemoryContextAllocZero(tree->iter_context,
+														   sizeof(RT_ITER_CONTROL));
 
-	Assert(RT_PTR_ALLOC_IS_VALID(tree->ctl->root));
-	root.alloc = iter->tree->ctl->root;
-	RT_PTR_SET_LOCAL(tree, &root);
+	RT_INITIALIZE_ITER(tree, iter);
 
-	iter->top_level = iter->tree->ctl->start_shift / RT_SPAN;
+#ifdef RT_SHMEM
+	/* we will non-shared iteration on a shared radix tree */
+	iter->shared = false;
+#endif
 
-	/* Set the root to start */
-	iter->cur_level = iter->top_level;
-	iter->node_iters[iter->cur_level].node = root;
-	iter->node_iters[iter->cur_level].idx = 0;
+	return iter;
+}
+
+#ifdef RT_SHMEM
+/*
+ * Create and return the shared iterator for the given shard radix tree.
+ *
+ * Taking a lock on a radix tree in shared mode during the shared iteration to
+ * prevent concurrent writes is the caller's responsibility.
+ */
+RT_SCOPE	RT_ITER *
+RT_BEGIN_ITERATE_SHARED(RT_RADIX_TREE * tree)
+{
+	RT_ITER    *iter;
+	RT_ITER_CONTROL_SHARED *ctl_shared;
+	dsa_pointer dp;
+
+	/* The radix tree must be in shared mode */
+	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
+
+	dp = dsa_allocate(tree->dsa, sizeof(RT_ITER_CONTROL_SHARED));
+	ctl_shared = (RT_ITER_CONTROL_SHARED *) dsa_get_address(tree->dsa, dp);
+	ctl_shared->handle = dp;
+	LWLockInitialize(&ctl_shared->lock, tree->ctl->tranche_id);
+	pg_atomic_init_u32(&ctl_shared->refcnt, 1);
+
+	iter = (RT_ITER *) MemoryContextAllocZero(tree->iter_context,
+											  sizeof(RT_ITER));
+
+	iter->ctl = (RT_ITER_CONTROL *) ctl_shared;
+	iter->shared = true;
+
+	RT_INITIALIZE_ITER(tree, iter);
 
 	return iter;
 }
+#endif
 
 /*
  * Scan the inner node and return the next child pointer if one exists, otherwise
@@ -2114,12 +2204,18 @@ RT_NODE_ITERATE_NEXT(RT_ITER * iter, int level)
 	RT_CHILD_PTR node;
 	RT_PTR_ALLOC *slot = NULL;
 
+	node_iter = &(iter->ctl->node_iters[level]);
+	node = node_iter->node;
+
 #ifdef RT_SHMEM
-	Assert(iter->tree->ctl->magic == RT_RADIX_TREE_MAGIC);
-#endif
 
-	node_iter = &(iter->node_iters[level]);
-	node = node_iter->node;
+	/*
+	 * Since the iterator is shared, the local pointer of the node might be
+	 * set by other backends, we need to make sure to use the local pointer.
+	 */
+	if (iter->shared)
+		RT_PTR_SET_LOCAL(iter->tree, &node);
+#endif
 
 	Assert(node.local != NULL);
 
@@ -2192,8 +2288,8 @@ RT_NODE_ITERATE_NEXT(RT_ITER * iter, int level)
 	}
 
 	/* Update the key */
-	iter->key &= ~(((uint64) RT_CHUNK_MASK) << (level * RT_SPAN));
-	iter->key |= (((uint64) key_chunk) << (level * RT_SPAN));
+	iter->ctl->key &= ~(((uint64) RT_CHUNK_MASK) << (level * RT_SPAN));
+	iter->ctl->key |= (((uint64) key_chunk) << (level * RT_SPAN));
 
 	return slot;
 }
@@ -2207,18 +2303,29 @@ RT_ITERATE_NEXT(RT_ITER * iter, uint64 *key_p)
 {
 	RT_PTR_ALLOC *slot = NULL;
 
-	while (iter->cur_level <= iter->top_level)
+#ifdef RT_SHMEM
+	/* Prevent the shared iterator from being updated concurrently */
+	if (iter->shared)
+		LWLockAcquire(&((RT_ITER_CONTROL_SHARED *) iter->ctl)->lock, LW_EXCLUSIVE);
+#endif
+
+	while (iter->ctl->cur_level <= iter->ctl->top_level)
 	{
 		RT_CHILD_PTR node;
 
-		slot = RT_NODE_ITERATE_NEXT(iter, iter->cur_level);
+		slot = RT_NODE_ITERATE_NEXT(iter, iter->ctl->cur_level);
 
-		if (iter->cur_level == 0 && slot != NULL)
+		if (iter->ctl->cur_level == 0 && slot != NULL)
 		{
 			/* Found a value at the leaf node */
-			*key_p = iter->key;
+			*key_p = iter->ctl->key;
 			node.alloc = *slot;
 
+#ifdef RT_SHMEM
+			if (iter->shared)
+				LWLockRelease(&((RT_ITER_CONTROL_SHARED *) iter->ctl)->lock);
+#endif
+
 			if (RT_CHILDPTR_IS_VALUE(*slot))
 				return (RT_VALUE_TYPE *) slot;
 			else
@@ -2234,17 +2341,23 @@ RT_ITERATE_NEXT(RT_ITER * iter, uint64 *key_p)
 			node.alloc = *slot;
 			RT_PTR_SET_LOCAL(iter->tree, &node);
 
-			iter->cur_level--;
-			iter->node_iters[iter->cur_level].node = node;
-			iter->node_iters[iter->cur_level].idx = 0;
+			iter->ctl->cur_level--;
+			iter->ctl->node_iters[iter->ctl->cur_level].node = node;
+			iter->ctl->node_iters[iter->ctl->cur_level].idx = 0;
 		}
 		else
 		{
 			/* Not found the child slot, move up the tree */
-			iter->cur_level++;
+			iter->ctl->cur_level++;
 		}
+
 	}
 
+#ifdef RT_SHMEM
+	if (iter->shared)
+		LWLockRelease(&((RT_ITER_CONTROL_SHARED *) iter->ctl)->lock);
+#endif
+
 	/* We've visited all nodes, so the iteration finished */
 	return NULL;
 }
@@ -2255,9 +2368,45 @@ RT_ITERATE_NEXT(RT_ITER * iter, uint64 *key_p)
 RT_SCOPE void
 RT_END_ITERATE(RT_ITER * iter)
 {
+#ifdef RT_SHMEM
+	RT_ITER_CONTROL_SHARED *ctl = (RT_ITER_CONTROL_SHARED *) iter->ctl;;
+
+	if (iter->shared &&
+		pg_atomic_sub_fetch_u32(&ctl->refcnt, 1) == 0)
+		dsa_free(iter->tree->dsa, ctl->handle);
+#endif
 	pfree(iter);
 }
 
+#ifdef	RT_SHMEM
+RT_SCOPE	RT_ITER_HANDLE
+RT_GET_ITER_HANDLE(RT_ITER * iter)
+{
+	Assert(iter->shared);
+	return ((RT_ITER_CONTROL_SHARED *) iter->ctl)->handle;
+
+}
+
+RT_SCOPE	RT_ITER *
+RT_ATTACH_ITERATE_SHARED(RT_RADIX_TREE * tree, RT_ITER_HANDLE handle)
+{
+	RT_ITER    *iter;
+	RT_ITER_CONTROL_SHARED *ctl;
+
+	iter = (RT_ITER *) MemoryContextAllocZero(tree->iter_context,
+											  sizeof(RT_ITER));
+	iter->tree = tree;
+	ctl = (RT_ITER_CONTROL_SHARED *) dsa_get_address(tree->dsa, handle);
+	iter->ctl = (RT_ITER_CONTROL *) ctl;
+	iter->shared = true;
+
+	/* For every iterator, increase the refcnt by 1 */
+	pg_atomic_add_fetch_u32(&ctl->refcnt, 1);
+
+	return iter;
+}
+#endif
+
 /***************** DELETION *****************/
 
 #ifdef RT_USE_DELETE
@@ -2957,7 +3106,11 @@ RT_DUMP_NODE(RT_NODE * node)
 #undef RT_PTR_ALLOC
 #undef RT_INVALID_PTR_ALLOC
 #undef RT_HANDLE
+#undef RT_ITER_HANDLE
+#undef RT_ITER_CONTROL
+#undef RT_ITER_HANDLE
 #undef RT_ITER
+#undef RT_SHARED_ITER
 #undef RT_NODE
 #undef RT_NODE_ITER
 #undef RT_NODE_KIND_4
@@ -2994,6 +3147,11 @@ RT_DUMP_NODE(RT_NODE * node)
 #undef RT_LOCK_SHARE
 #undef RT_UNLOCK
 #undef RT_GET_HANDLE
+#undef RT_BEGIN_ITERATE_SHARED
+#undef RT_ATTACH_ITERATE_SHARED
+#undef RT_GET_ITER_HANDLE
+#undef RT_ATTACH_ITER
+#undef RT_GET_ITER_HANDLE
 #undef RT_FIND
 #undef RT_SET
 #undef RT_BEGIN_ITERATE
@@ -3050,5 +3208,6 @@ RT_DUMP_NODE(RT_NODE * node)
 #undef RT_SHRINK_NODE_256
 #undef RT_NODE_DELETE
 #undef RT_NODE_INSERT
+#undef RT_INITIALIZE_ITER
 #undef RT_NODE_ITERATE_NEXT
 #undef RT_VERIFY_NODE
-- 
2.43.5

