v17-0010-fix-shmem-support.patch

From 56a45a0731abc33b3894d0aa0de06869d894637b Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Thu, 12 Jan 2023 23:22:22 +0900
Subject: [PATCH v17 10/12] fix shmem support

---
 src/include/lib/radixtree.h           | 87 ++++++++++++++++++++++++---
 src/include/lib/radixtree_iter_impl.h |  4 ++
 2 files changed, 82 insertions(+), 9 deletions(-)

diff --git a/src/include/lib/radixtree.h b/src/include/lib/radixtree.h
index 2b58a0cdf5..a2e2e7a190 100644
--- a/src/include/lib/radixtree.h
+++ b/src/include/lib/radixtree.h
@@ -100,6 +100,8 @@
 #define RT_SEARCH RT_MAKE_NAME(search)
 #ifdef RT_SHMEM
 #define RT_ATTACH RT_MAKE_NAME(attach)
+#define RT_DETACH RT_MAKE_NAME(detach)
+#define RT_GET_HANDLE RT_MAKE_NAME(get_handle)
 #endif
 #define RT_SET RT_MAKE_NAME(set)
 #define RT_BEGIN_ITERATE RT_MAKE_NAME(begin_iterate)
@@ -164,6 +166,9 @@
 #define RT_RADIX_TREE RT_MAKE_NAME(radix_tree)
 #define RT_RADIX_TREE_CONTROL RT_MAKE_NAME(radix_tree_control)
 #define RT_ITER RT_MAKE_NAME(iter)
+#ifdef RT_SHMEM
+#define RT_HANDLE RT_MAKE_NAME(handle)
+#endif
 #define RT_NODE RT_MAKE_NAME(node)
 #define RT_NODE_ITER RT_MAKE_NAME(node_iter)
 #define RT_NODE_BASE_4 RT_MAKE_NAME(node_base_4)
@@ -194,9 +199,15 @@
 typedef struct RT_RADIX_TREE RT_RADIX_TREE;
 typedef struct RT_ITER RT_ITER;
 
+#ifdef RT_SHMEM
+typedef dsa_pointer RT_HANDLE;
+#endif
+
 #ifdef RT_SHMEM
 RT_SCOPE RT_RADIX_TREE * RT_CREATE(MemoryContext ctx, dsa_area *dsa);
 RT_SCOPE RT_RADIX_TREE * RT_ATTACH(dsa_area *dsa, dsa_pointer dp);
+RT_SCOPE void RT_DETACH(RT_RADIX_TREE *tree);
+RT_SCOPE RT_HANDLE RT_GET_HANDLE(RT_RADIX_TREE *tree);
 #else
 RT_SCOPE RT_RADIX_TREE * RT_CREATE(MemoryContext ctx);
 #endif
@@ -542,9 +553,19 @@ static const RT_SIZE_CLASS RT_KIND_MIN_SIZE_CLASS[RT_NODE_KIND_COUNT] = {
 	[RT_NODE_KIND_256] = RT_CLASS_256,
 };
 
+#ifdef RT_SHMEM
+/* A magic value used to identify our radix tree */
+#define RT_RADIX_TREE_MAGIC 0x54A48167
+#endif
+
 /* A radix tree with nodes */
 typedef struct RT_RADIX_TREE_CONTROL
 {
+#ifdef RT_SHMEM
+	RT_HANDLE	handle;
+	uint32		magic;
+#endif
+
 	RT_PTR_ALLOC root;
 	uint64		max_val;
 	uint64		num_keys;
@@ -565,7 +586,6 @@ typedef struct RT_RADIX_TREE
 
 #ifdef RT_SHMEM
 	dsa_area   *dsa;
-	dsa_pointer ctl_dp;
 #else
 	MemoryContextData *inner_slabs[RT_SIZE_CLASS_COUNT];
 	MemoryContextData *leaf_slabs[RT_SIZE_CLASS_COUNT];
@@ -1311,6 +1331,9 @@ RT_CREATE(MemoryContext ctx)
 {
 	RT_RADIX_TREE *tree;
 	MemoryContext old_ctx;
+#ifdef RT_SHMEM
+	dsa_pointer dp;
+#endif
 
 	old_ctx = MemoryContextSwitchTo(ctx);
 
@@ -1319,8 +1342,10 @@ RT_CREATE(MemoryContext ctx)
 
 #ifdef RT_SHMEM
 	tree->dsa = dsa;
-	tree->ctl_dp = dsa_allocate0(dsa, sizeof(RT_RADIX_TREE_CONTROL));
-	tree->ctl = (RT_RADIX_TREE_CONTROL *) dsa_get_address(dsa, tree->ctl_dp);
+	dp = dsa_allocate0(dsa, sizeof(RT_RADIX_TREE_CONTROL));
+	tree->ctl = (RT_RADIX_TREE_CONTROL *) dsa_get_address(dsa, dp);
+	tree->ctl->handle = dp;
+	tree->ctl->magic = RT_RADIX_TREE_MAGIC;
 #else
 	tree->ctl = (RT_RADIX_TREE_CONTROL *) palloc0(sizeof(RT_RADIX_TREE_CONTROL));
 
@@ -1346,21 +1371,40 @@ RT_CREATE(MemoryContext ctx)
 }
 
 #ifdef RT_SHMEM
-RT_RADIX_TREE *
-RT_ATTACH(dsa_area *dsa, dsa_pointer dp)
+RT_SCOPE RT_RADIX_TREE *
+RT_ATTACH(dsa_area *dsa, RT_HANDLE handle)
 {
 	RT_RADIX_TREE *tree;
+	dsa_pointer	control;
 
 	/* XXX: memory context support */
 	tree = (RT_RADIX_TREE *) palloc0(sizeof(RT_RADIX_TREE));
 
-	tree->ctl_dp = dp;
-	tree->ctl = (RT_RADIX_TREE_CONTROL *) dsa_get_address(dsa, dp);
+	/* Find the control object in shard memory */
+	control = handle;
+
+	tree->dsa = dsa;
+	tree->ctl = (RT_RADIX_TREE_CONTROL *) dsa_get_address(dsa, control);
+	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
 
 	/* XXX: do we need to set a callback on exit to detach dsa? */
 
 	return tree;
 }
+
+RT_SCOPE void
+RT_DETACH(RT_RADIX_TREE *tree)
+{
+	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
+	pfree(tree);
+}
+
+RT_SCOPE RT_HANDLE
+RT_GET_HANDLE(RT_RADIX_TREE *tree)
+{
+	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
+	return tree->ctl->handle;
+}
 #endif
 
 /*
@@ -1370,8 +1414,15 @@ RT_SCOPE void
 RT_FREE(RT_RADIX_TREE *tree)
 {
 #ifdef RT_SHMEM
-	dsa_free(tree->dsa, tree->ctl_dp); // XXX
-	dsa_detach(tree->dsa);
+	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
+
+	/*
+	 * Vandalize the control block to help catch programming error where
+	 * other backends access the memory formerly occupied by this radix tree.
+	 */
+	tree->ctl->magic = 0;
+	dsa_free(tree->dsa, tree->ctl->handle); // XXX
+	//dsa_detach(tree->dsa);
 #else
 	pfree(tree->ctl);
 
@@ -1398,6 +1449,10 @@ RT_SET(RT_RADIX_TREE *tree, uint64 key, uint64 value)
 	RT_PTR_ALLOC nodep;
 	RT_PTR_LOCAL  node;
 
+#ifdef RT_SHMEM
+	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
+#endif
+
 	/* Empty tree, create the root */
 	if (!RT_PTR_ALLOC_IS_VALID(tree->ctl->root))
 		RT_NEW_ROOT(tree, key);
@@ -1453,6 +1508,9 @@ RT_SEARCH(RT_RADIX_TREE *tree, uint64 key, uint64 *value_p)
 	RT_PTR_LOCAL node;
 	int			shift;
 
+#ifdef RT_SHMEM
+	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
+#endif
 	Assert(value_p != NULL);
 
 	if (!RT_PTR_ALLOC_IS_VALID(tree->ctl->root) || key > tree->ctl->max_val)
@@ -1493,6 +1551,10 @@ RT_DELETE(RT_RADIX_TREE *tree, uint64 key)
 	int			level;
 	bool		deleted;
 
+#ifdef RT_SHMEM
+	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
+#endif
+
 	if (!RT_PTR_ALLOC_IS_VALID(tree->ctl->root) || key > tree->ctl->max_val)
 		return false;
 
@@ -1736,6 +1798,7 @@ RT_MEMORY_USAGE(RT_RADIX_TREE *tree)
 	Size		total = sizeof(RT_RADIX_TREE);
 
 #ifdef RT_SHMEM
+	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
 	total = dsa_get_total_size(tree->dsa);
 #else
 	for (int i = 0; i < RT_SIZE_CLASS_COUNT; i++)
@@ -2085,10 +2148,14 @@ rt_dump(RT_RADIX_TREE *tree)
 #undef VAR_NODE_HAS_FREE_SLOT
 #undef FIXED_NODE_HAS_FREE_SLOT
 #undef RT_SIZE_CLASS_COUNT
+#undef RT_RADIX_TREE_MAGIC
 
 /* type declarations */
 #undef RT_RADIX_TREE
 #undef RT_RADIX_TREE_CONTROL
+#undef RT_PTR_ALLOC
+#undef RT_INVALID_PTR_ALLOC
+#undef RT_HANDLE
 #undef RT_ITER
 #undef RT_NODE
 #undef RT_NODE_ITER
@@ -2118,6 +2185,8 @@ rt_dump(RT_RADIX_TREE *tree)
 #undef RT_CREATE
 #undef RT_FREE
 #undef RT_ATTACH
+#undef RT_DETACH
+#undef RT_GET_HANDLE
 #undef RT_SET
 #undef RT_BEGIN_ITERATE
 #undef RT_ITERATE_NEXT
diff --git a/src/include/lib/radixtree_iter_impl.h b/src/include/lib/radixtree_iter_impl.h
index 09d2018dc0..fd00851732 100644
--- a/src/include/lib/radixtree_iter_impl.h
+++ b/src/include/lib/radixtree_iter_impl.h
@@ -12,6 +12,10 @@
 #error node level must be either inner or leaf
 #endif
 
+#ifdef RT_SHMEM
+	Assert(iter->tree->ctl->magic == RT_RADIX_TREE_MAGIC);
+#endif
+
 	bool		found = false;
 	uint8		key_chunk;
 
-- 
2.31.1