From 56c458643f58723c59ed28477f6d129374a59e6c Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Wed, 25 Jan 2023 11:04:41 +0900
Subject: [PATCH v23 11/18] Add a safeguard for concurrent iteration in
 RT_SHMEM case.

---
 src/include/lib/radixtree.h | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/src/include/lib/radixtree.h b/src/include/lib/radixtree.h
index 003e8215aa..0277d5e6fb 100644
--- a/src/include/lib/radixtree.h
+++ b/src/include/lib/radixtree.h
@@ -602,6 +602,9 @@ typedef struct RT_RADIX_TREE_CONTROL
 	uint64		max_val;
 	uint64		num_keys;
 
+	/* is iteration in progress? */
+	bool		iter_active;
+
 	/* statistics */
 #ifdef RT_DEBUG
 	int32		cnt[RT_SIZE_CLASS_COUNT];
@@ -638,10 +641,7 @@ typedef struct RT_RADIX_TREE
  * advancing the current index within the node or when moving to the next node
  * at the same level.
  *
- * XXX: Currently we allow only one process to do iteration. Therefore, rt_node_iter
- * has the local pointers to nodes, rather than RT_PTR_ALLOC.
- * We need either a safeguard to disallow other processes to begin the iteration
- * while one process is doing or to allow multiple processes to do the iteration.
+ * In RT_SHMEM case, only one process is allowed to do iteration.
  */
 typedef struct RT_NODE_ITER
 {
@@ -1582,6 +1582,9 @@ RT_SET(RT_RADIX_TREE *tree, uint64 key, RT_VALUE_TYPE value)
 	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
 #endif
 
+	if (unlikely(tree->ctl->iter_active))
+		elog(ERROR, "cannot add new key-value to radix tree while iteration is in progress");
+
 	/* Empty tree, create the root */
 	if (!RT_PTR_ALLOC_IS_VALID(tree->ctl->root))
 		RT_NEW_ROOT(tree, key);
@@ -1683,6 +1686,9 @@ RT_DELETE(RT_RADIX_TREE *tree, uint64 key)
 	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
 #endif
 
+	if (unlikely(tree->ctl->iter_active))
+		elog(ERROR, "cannot delete key to radix tree while iteration is in progress");
+
 	if (!RT_PTR_ALLOC_IS_VALID(tree->ctl->root) || key > tree->ctl->max_val)
 		return false;
 
@@ -1822,10 +1828,14 @@ RT_BEGIN_ITERATE(RT_RADIX_TREE *tree)
 	RT_PTR_LOCAL root;
 	int			top_level;
 
+	if (unlikely(tree->ctl->iter_active))
+		elog(ERROR, "cannot begin iteration while another iteration is in progress");
+
 	old_ctx = MemoryContextSwitchTo(tree->context);
 
 	iter = (RT_ITER *) palloc0(sizeof(RT_ITER));
 	iter->tree = tree;
+	tree->ctl->iter_active = true;
 
 	/* empty tree */
 	if (!iter->tree->ctl->root)
@@ -1853,6 +1863,8 @@ RT_BEGIN_ITERATE(RT_RADIX_TREE *tree)
 RT_SCOPE bool
 RT_ITERATE_NEXT(RT_ITER *iter, uint64 *key_p, RT_VALUE_TYPE *value_p)
 {
+	Assert(iter->tree->ctl->iter_active);
+
 	/* Empty tree */
 	if (!iter->tree->ctl->root)
 		return false;
@@ -1905,6 +1917,7 @@ RT_ITERATE_NEXT(RT_ITER *iter, uint64 *key_p, RT_VALUE_TYPE *value_p)
 RT_SCOPE void
 RT_END_ITERATE(RT_ITER *iter)
 {
+	iter->tree->ctl->iter_active = false;
 	pfree(iter);
 }
 
-- 
2.31.1