From c78f27a61d649b0981fc150c3894a0e1a992bcc0 Mon Sep 17 00:00:00 2001
From: John Naylor <john.naylor@postgresql.org>
Date: Mon, 9 Jan 2023 14:32:39 +0700
Subject: [PATCH v17 9/9] Implement shared memory

---
 src/backend/utils/mmgr/dsa.c                  |  12 +
 src/include/lib/radixtree.h                   | 376 +++++++++++++-----
 src/include/lib/radixtree_delete_impl.h       |   6 +
 src/include/lib/radixtree_insert_impl.h       |  43 +-
 src/include/lib/radixtree_iter_impl.h         |  19 +-
 src/include/lib/radixtree_search_impl.h       |  28 +-
 src/include/utils/dsa.h                       |   1 +
 .../modules/test_radixtree/test_radixtree.c   |  43 ++
 8 files changed, 402 insertions(+), 126 deletions(-)

diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index 604b702a91..50f0aae3ab 100644
--- a/src/backend/utils/mmgr/dsa.c
+++ b/src/backend/utils/mmgr/dsa.c
@@ -1024,6 +1024,18 @@ dsa_set_size_limit(dsa_area *area, size_t limit)
 	LWLockRelease(DSA_AREA_LOCK(area));
 }
 
+size_t
+dsa_get_total_size(dsa_area *area)
+{
+	size_t		size;
+
+	LWLockAcquire(DSA_AREA_LOCK(area), LW_SHARED);
+	size = area->control->total_segment_size;
+	LWLockRelease(DSA_AREA_LOCK(area));
+
+	return size;
+}
+
 /*
  * Aggressively free all spare memory in the hope of returning DSM segments to
  * the operating system.
diff --git a/src/include/lib/radixtree.h b/src/include/lib/radixtree.h
index b3d84da033..2b58a0cdf5 100644
--- a/src/include/lib/radixtree.h
+++ b/src/include/lib/radixtree.h
@@ -42,6 +42,8 @@
  *	  - RT_DEFINE - if defined function definitions are generated
  *	  - RT_SCOPE - in which scope (e.g. extern, static inline) do function
  *		declarations reside
+ *	  - RT_SHMEM - if defined, the radix tree is created in the DSA area
+ *		so that multiple processes can access it simultaneously.
  *
  *	  Optional parameters:
  *	  - RT_DEBUG - if defined add stats tracking and debugging functions
@@ -51,6 +53,9 @@
  *
  * RT_CREATE		- Create a new, empty radix tree
  * RT_FREE			- Free the radix tree
+ * RT_ATTACH		- Attach to the radix tree
+ * RT_DETACH		- Detach from the radix tree
+ * RT_GET_HANDLE	- Return the handle of the radix tree
  * RT_SEARCH		- Search a key-value pair
  * RT_SET			- Set a key-value pair
  * RT_DELETE		- Delete a key-value pair
@@ -80,7 +85,8 @@
 #include "miscadmin.h"
 #include "nodes/bitmapset.h"
 #include "port/pg_bitutils.h"
-#include "port/pg_lfind.h"
+#include "port/simd.h"
+#include "utils/dsa.h"
 #include "utils/memutils.h"
 
 /* helpers */
@@ -92,6 +98,9 @@
 #define RT_CREATE RT_MAKE_NAME(create)
 #define RT_FREE RT_MAKE_NAME(free)
 #define RT_SEARCH RT_MAKE_NAME(search)
+#ifdef RT_SHMEM
+#define RT_ATTACH RT_MAKE_NAME(attach)
+#endif
 #define RT_SET RT_MAKE_NAME(set)
 #define RT_BEGIN_ITERATE RT_MAKE_NAME(begin_iterate)
 #define RT_ITERATE_NEXT RT_MAKE_NAME(iterate_next)
@@ -110,9 +119,11 @@
 #define RT_FREE_NODE RT_MAKE_NAME(free_node)
 #define RT_EXTEND RT_MAKE_NAME(extend)
 #define RT_SET_EXTEND RT_MAKE_NAME(set_extend)
-#define RT_GROW_NODE_KIND RT_MAKE_NAME(grow_node_kind)
+//#define RT_GROW_NODE_KIND RT_MAKE_NAME(grow_node_kind)
 #define RT_COPY_NODE RT_MAKE_NAME(copy_node)
 #define RT_REPLACE_NODE RT_MAKE_NAME(replace_node)
+#define RT_PTR_GET_LOCAL RT_MAKE_NAME(ptr_get_local)
+#define RT_PTR_ALLOC_IS_VALID RT_MAKE_NAME(ptr_stored_is_valid)
 #define RT_NODE_4_SEARCH_EQ RT_MAKE_NAME(node_4_search_eq)
 #define RT_NODE_32_SEARCH_EQ RT_MAKE_NAME(node_32_search_eq)
 #define RT_NODE_4_GET_INSERTPOS RT_MAKE_NAME(node_4_get_insertpos)
@@ -138,6 +149,7 @@
 #define RT_SHIFT_GET_MAX_VAL RT_MAKE_NAME(shift_get_max_val)
 #define RT_NODE_SEARCH_INNER RT_MAKE_NAME(node_search_inner)
 #define RT_NODE_SEARCH_LEAF RT_MAKE_NAME(node_search_leaf)
+#define RT_NODE_UPDATE_INNER RT_MAKE_NAME(node_update_inner)
 #define RT_NODE_DELETE_INNER RT_MAKE_NAME(node_delete_inner)
 #define RT_NODE_DELETE_LEAF RT_MAKE_NAME(node_delete_leaf)
 #define RT_NODE_INSERT_INNER RT_MAKE_NAME(node_insert_inner)
@@ -150,6 +162,7 @@
 
 /* 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 RT_MAKE_NAME(iter)
 #define RT_NODE RT_MAKE_NAME(node)
 #define RT_NODE_ITER RT_MAKE_NAME(node_iter)
@@ -181,8 +194,14 @@
 typedef struct RT_RADIX_TREE RT_RADIX_TREE;
 typedef struct RT_ITER RT_ITER;
 
+#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);
+#else
 RT_SCOPE RT_RADIX_TREE * RT_CREATE(MemoryContext ctx);
+#endif
 RT_SCOPE void RT_FREE(RT_RADIX_TREE *tree);
+
 RT_SCOPE bool RT_SEARCH(RT_RADIX_TREE *tree, uint64 key, uint64 *val_p);
 RT_SCOPE bool RT_SET(RT_RADIX_TREE *tree, uint64 key, uint64 val);
 RT_SCOPE bool RT_DELETE(RT_RADIX_TREE *tree, uint64 key);
@@ -301,9 +320,21 @@ typedef struct RT_NODE
 	uint8		kind;
 } RT_NODE;
 
+
 #define RT_PTR_LOCAL RT_NODE *
 
+#ifdef RT_SHMEM
+#define RT_PTR_ALLOC dsa_pointer
+#else
 #define RT_PTR_ALLOC RT_PTR_LOCAL
+#endif
+
+
+#ifdef RT_SHMEM
+#define RT_INVALID_PTR_ALLOC InvalidDsaPointer
+#else
+#define RT_INVALID_PTR_ALLOC NULL
+#endif
 
 #define NODE_IS_LEAF(n)			(((RT_PTR_LOCAL) (n))->shift == 0)
 #define NODE_IS_EMPTY(n)		(((RT_PTR_LOCAL) (n))->count == 0)
@@ -512,21 +543,33 @@ static const RT_SIZE_CLASS RT_KIND_MIN_SIZE_CLASS[RT_NODE_KIND_COUNT] = {
 };
 
 /* A radix tree with nodes */
-typedef struct RT_RADIX_TREE
+typedef struct RT_RADIX_TREE_CONTROL
 {
-	MemoryContext context;
-
 	RT_PTR_ALLOC root;
 	uint64		max_val;
 	uint64		num_keys;
 
-	MemoryContextData *inner_slabs[RT_SIZE_CLASS_COUNT];
-	MemoryContextData *leaf_slabs[RT_SIZE_CLASS_COUNT];
-
 	/* statistics */
 #ifdef RT_DEBUG
 	int32		cnt[RT_SIZE_CLASS_COUNT];
 #endif
+} RT_RADIX_TREE_CONTROL;
+
+/* A radix tree with nodes */
+typedef struct RT_RADIX_TREE
+{
+	MemoryContext context;
+
+	/* pointing to either local memory or DSA */
+	RT_RADIX_TREE_CONTROL *ctl;
+
+#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];
+#endif
 } RT_RADIX_TREE;
 
 /*
@@ -542,6 +585,11 @@ typedef struct RT_RADIX_TREE
  * construct the key whenever updating the node iteration information, e.g., when
  * 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.
  */
 typedef struct RT_NODE_ITER
 {
@@ -562,14 +610,35 @@ typedef struct RT_ITER
 } RT_ITER;
 
 
-static bool RT_NODE_INSERT_INNER(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_LOCAL node,
-								 uint64 key, RT_PTR_LOCAL child);
-static bool RT_NODE_INSERT_LEAF(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_LOCAL node,
+static bool RT_NODE_INSERT_INNER(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_ALLOC nodep, RT_PTR_LOCAL node,
+								 uint64 key, RT_PTR_ALLOC child);
+static bool RT_NODE_INSERT_LEAF(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_ALLOC nodep, RT_PTR_LOCAL node,
 								uint64 key, uint64 value);
 
 /* verification (available only with assertion) */
 static void RT_VERIFY_NODE(RT_PTR_LOCAL node);
 
+/* Get the local address of an allocated node */
+static inline RT_PTR_LOCAL
+RT_PTR_GET_LOCAL(RT_RADIX_TREE *tree, RT_PTR_ALLOC node)
+{
+#ifdef RT_SHMEM
+	return dsa_get_address(tree->dsa, (dsa_pointer) node);
+#else
+	return node;
+#endif
+}
+
+static inline bool
+RT_PTR_ALLOC_IS_VALID(RT_PTR_ALLOC ptr)
+{
+#ifdef RT_SHMEM
+	return DsaPointerIsValid(ptr);
+#else
+	return PointerIsValid(ptr);
+#endif
+}
+
 /*
  * Return index of the first element in 'base' that equals 'key'. Return -1
  * if there is no such element.
@@ -801,7 +870,7 @@ static inline bool
 RT_NODE_INNER_256_IS_CHUNK_USED(RT_NODE_INNER_256 *node, uint8 chunk)
 {
 	Assert(!NODE_IS_LEAF(node));
-	return (node->children[chunk] != NULL);
+	return node->children[chunk] != RT_INVALID_PTR_ALLOC;
 }
 
 static inline bool
@@ -855,7 +924,7 @@ static inline void
 RT_NODE_INNER_256_DELETE(RT_NODE_INNER_256 *node, uint8 chunk)
 {
 	Assert(!NODE_IS_LEAF(node));
-	node->children[chunk] = NULL;
+	node->children[chunk] = RT_INVALID_PTR_ALLOC;
 }
 
 static inline void
@@ -897,21 +966,31 @@ RT_SHIFT_GET_MAX_VAL(int shift)
 static RT_PTR_ALLOC
 RT_ALLOC_NODE(RT_RADIX_TREE *tree, RT_SIZE_CLASS size_class, bool inner)
 {
-	RT_PTR_ALLOC newnode;
+	RT_PTR_ALLOC allocnode;
+	size_t allocsize;
 
 	if (inner)
-		newnode = (RT_PTR_ALLOC) MemoryContextAlloc(tree->inner_slabs[size_class],
-												 RT_SIZE_CLASS_INFO[size_class].inner_size);
+		allocsize = RT_SIZE_CLASS_INFO[size_class].inner_size;
 	else
-		newnode = (RT_PTR_ALLOC) MemoryContextAlloc(tree->leaf_slabs[size_class],
-												 RT_SIZE_CLASS_INFO[size_class].leaf_size);
+		allocsize = RT_SIZE_CLASS_INFO[size_class].leaf_size;
+
+#ifdef RT_SHMEM
+	allocnode = dsa_allocate(tree->dsa, allocsize);
+#else
+	if (inner)
+		allocnode = (RT_PTR_ALLOC) MemoryContextAlloc(tree->inner_slabs[size_class],
+													  allocsize);
+	else
+		allocnode = (RT_PTR_ALLOC) MemoryContextAlloc(tree->leaf_slabs[size_class],
+													  allocsize);
+#endif
 
 #ifdef RT_DEBUG
 	/* update the statistics */
-	tree->cnt[size_class]++;
+	tree->ctl->cnt[size_class]++;
 #endif
 
-	return newnode;
+	return allocnode;
 }
 
 /* Initialize the node contents */
@@ -951,13 +1030,15 @@ RT_NEW_ROOT(RT_RADIX_TREE *tree, uint64 key)
 {
 	int			shift = RT_KEY_GET_SHIFT(key);
 	bool		inner = shift > 0;
-	RT_PTR_ALLOC newnode;
+	RT_PTR_ALLOC allocnode;
+	RT_PTR_LOCAL newnode;
 
-	newnode = RT_ALLOC_NODE(tree, RT_CLASS_4_FULL, inner);
+	allocnode = RT_ALLOC_NODE(tree, RT_CLASS_4_FULL, inner);
+	newnode = RT_PTR_GET_LOCAL(tree, allocnode);
 	RT_INIT_NODE(newnode, RT_NODE_KIND_4, RT_CLASS_4_FULL, inner);
 	newnode->shift = shift;
-	tree->max_val = RT_SHIFT_GET_MAX_VAL(shift);
-	tree->root = newnode;
+	tree->ctl->max_val = RT_SHIFT_GET_MAX_VAL(shift);
+	tree->ctl->root = allocnode;
 }
 
 static inline void
@@ -967,7 +1048,7 @@ RT_COPY_NODE(RT_PTR_LOCAL newnode, RT_PTR_LOCAL oldnode)
 	newnode->chunk = oldnode->chunk;
 	newnode->count = oldnode->count;
 }
-
+#if 0
 /*
  * Create a new node with 'new_kind' and the same shift, chunk, and
  * count of 'node'.
@@ -975,30 +1056,33 @@ RT_COPY_NODE(RT_PTR_LOCAL newnode, RT_PTR_LOCAL oldnode)
 static RT_NODE*
 RT_GROW_NODE_KIND(RT_RADIX_TREE *tree, RT_PTR_LOCAL node, uint8 new_kind)
 {
-	RT_PTR_ALLOC newnode;
+	RT_PTR_ALLOC allocnode;
+	RT_PTR_LOCAL newnode;
 	bool inner = !NODE_IS_LEAF(node);
 
-	newnode = RT_ALLOC_NODE(tree, RT_KIND_MIN_SIZE_CLASS[new_kind], inner);
+	allocnode = RT_ALLOC_NODE(tree, RT_KIND_MIN_SIZE_CLASS[new_kind], inner);
+	newnode = RT_PTR_GET_LOCAL(tree, allocnode);
 	RT_INIT_NODE(newnode, new_kind, RT_KIND_MIN_SIZE_CLASS[new_kind], inner);
 	RT_COPY_NODE(newnode, node);
 
 	return newnode;
 }
-
+#endif
 /* Free the given node */
 static void
-RT_FREE_NODE(RT_RADIX_TREE *tree, RT_PTR_ALLOC node)
+RT_FREE_NODE(RT_RADIX_TREE *tree, RT_PTR_ALLOC allocnode)
 {
 	/* If we're deleting the root node, make the tree empty */
-	if (tree->root == node)
+	if (tree->ctl->root == allocnode)
 	{
-		tree->root = NULL;
-		tree->max_val = 0;
+		tree->ctl->root = RT_INVALID_PTR_ALLOC;
+		tree->ctl->max_val = 0;
 	}
 
 #ifdef RT_DEBUG
 	{
 		int i;
+		RT_PTR_LOCAL node = RT_PTR_GET_LOCAL(tree, allocnode);
 
 		/* update the statistics */
 		for (i = 0; i < RT_SIZE_CLASS_COUNT; i++)
@@ -1011,12 +1095,26 @@ RT_FREE_NODE(RT_RADIX_TREE *tree, RT_PTR_ALLOC node)
 		if (i == RT_SIZE_CLASS_COUNT)
 			i = RT_CLASS_256;
 
-		tree->cnt[i]--;
-		Assert(tree->cnt[i] >= 0);
+		tree->ctl->cnt[i]--;
+		Assert(tree->ctl->cnt[i] >= 0);
 	}
 #endif
 
-	pfree(node);
+#ifdef RT_SHMEM
+	dsa_free(tree->dsa, allocnode);
+#else
+	pfree(allocnode);
+#endif
+}
+
+static inline bool
+RT_NODE_UPDATE_INNER(RT_PTR_LOCAL node, uint64 key, RT_PTR_ALLOC new_child)
+{
+#define RT_ACTION_UPDATE
+#define RT_NODE_LEVEL_INNER
+#include "lib/radixtree_search_impl.h"
+#undef RT_NODE_LEVEL_INNER
+#undef RT_ACTION_UPDATE
 }
 
 /*
@@ -1026,19 +1124,25 @@ static void
 RT_REPLACE_NODE(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_ALLOC old_child,
 				RT_PTR_ALLOC new_child, uint64 key)
 {
-	Assert(old_child->chunk == new_child->chunk);
-	Assert(old_child->shift == new_child->shift);
+	RT_PTR_LOCAL old = RT_PTR_GET_LOCAL(tree, old_child);
+
+#ifdef USE_ASSERT_CHECKING
+	RT_PTR_LOCAL new = RT_PTR_GET_LOCAL(tree, new_child);
+
+	Assert(old->chunk == new->chunk);
+	Assert(old->shift == new->shift);
+#endif
 
-	if (parent == old_child)
+	if (parent == old)
 	{
 		/* Replace the root node with the new large node */
-		tree->root = new_child;
+		tree->ctl->root = new_child;
 	}
 	else
 	{
-		bool		replaced PG_USED_FOR_ASSERTS_ONLY;
+		bool replaced PG_USED_FOR_ASSERTS_ONLY;
 
-		replaced = RT_NODE_INSERT_INNER(tree, NULL, parent, key, new_child);
+		replaced = RT_NODE_UPDATE_INNER(parent, key, new_child);
 		Assert(replaced);
 	}
 
@@ -1053,7 +1157,8 @@ static void
 RT_EXTEND(RT_RADIX_TREE *tree, uint64 key)
 {
 	int			target_shift;
-	int			shift = tree->root->shift + RT_NODE_SPAN;
+	RT_PTR_LOCAL root = RT_PTR_GET_LOCAL(tree, tree->ctl->root);
+	int			shift = root->shift + RT_NODE_SPAN;
 
 	target_shift = RT_KEY_GET_SHIFT(key);
 
@@ -1065,22 +1170,23 @@ RT_EXTEND(RT_RADIX_TREE *tree, uint64 key)
 		RT_NODE_INNER_4 *n4;
 
 		allocnode = RT_ALLOC_NODE(tree, RT_CLASS_4_FULL, true);
-		node = (RT_PTR_LOCAL) allocnode;
+		node = RT_PTR_GET_LOCAL(tree, allocnode);
 		RT_INIT_NODE(node, RT_NODE_KIND_4, RT_CLASS_4_FULL, true);
 		node->shift = shift;
 		node->count = 1;
 
 		n4 = (RT_NODE_INNER_4 *) node;
 		n4->base.chunks[0] = 0;
-		n4->children[0] = tree->root;
+		n4->children[0] = tree->ctl->root;
 
-		tree->root->chunk = 0;
-		tree->root = node;
+		/* Update the root */
+		tree->ctl->root = allocnode;
+		root->chunk = 0;
 
 		shift += RT_NODE_SPAN;
 	}
 
-	tree->max_val = RT_SHIFT_GET_MAX_VAL(target_shift);
+	tree->ctl->max_val = RT_SHIFT_GET_MAX_VAL(target_shift);
 }
 
 /*
@@ -1089,10 +1195,12 @@ RT_EXTEND(RT_RADIX_TREE *tree, uint64 key)
  */
 static inline void
 RT_SET_EXTEND(RT_RADIX_TREE *tree, uint64 key, uint64 value, RT_PTR_LOCAL parent,
-			  RT_PTR_LOCAL node)
+			  RT_PTR_ALLOC nodep, RT_PTR_LOCAL node)
 {
 	int			shift = node->shift;
 
+	Assert(RT_PTR_GET_LOCAL(tree, nodep) == node);
+
 	while (shift >= RT_NODE_SPAN)
 	{
 		RT_PTR_ALLOC allocchild;
@@ -1101,19 +1209,20 @@ RT_SET_EXTEND(RT_RADIX_TREE *tree, uint64 key, uint64 value, RT_PTR_LOCAL parent
 		bool		inner = newshift > 0;
 
 		allocchild = RT_ALLOC_NODE(tree, RT_CLASS_4_FULL, inner);
-		newchild = (RT_PTR_LOCAL) allocchild;
+		newchild = RT_PTR_GET_LOCAL(tree, allocchild);
 		RT_INIT_NODE(newchild, RT_NODE_KIND_4, RT_CLASS_4_FULL, inner);
 		newchild->shift = newshift;
 		newchild->chunk = RT_GET_KEY_CHUNK(key, node->shift);
-		RT_NODE_INSERT_INNER(tree, parent, node, key, newchild);
+		RT_NODE_INSERT_INNER(tree, parent, nodep, node, key, allocchild);
 
 		parent = node;
 		node = newchild;
+		nodep = allocchild;
 		shift -= RT_NODE_SPAN;
 	}
 
-	RT_NODE_INSERT_LEAF(tree, parent, node, key, value);
-	tree->num_keys++;
+	RT_NODE_INSERT_LEAF(tree, parent, nodep, node, key, value);
+	tree->ctl->num_keys++;
 }
 
 /*
@@ -1172,8 +1281,8 @@ RT_NODE_DELETE_LEAF(RT_PTR_LOCAL node, uint64 key)
 
 /* Insert the child to the inner node */
 static bool
-RT_NODE_INSERT_INNER(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_LOCAL node, uint64 key,
-					 RT_PTR_ALLOC child)
+RT_NODE_INSERT_INNER(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_ALLOC nodep, RT_PTR_LOCAL node,
+					uint64 key, RT_PTR_ALLOC child)
 {
 #define RT_NODE_LEVEL_INNER
 #include "lib/radixtree_insert_impl.h"
@@ -1182,7 +1291,7 @@ RT_NODE_INSERT_INNER(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_LOCAL node
 
 /* Insert the value to the leaf node */
 static bool
-RT_NODE_INSERT_LEAF(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_LOCAL node,
+RT_NODE_INSERT_LEAF(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_ALLOC nodep, RT_PTR_LOCAL node,
 					uint64 key, uint64 value)
 {
 #define RT_NODE_LEVEL_LEAF
@@ -1194,18 +1303,26 @@ RT_NODE_INSERT_LEAF(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent, RT_PTR_LOCAL node,
  * Create the radix tree in the given memory context and return it.
  */
 RT_SCOPE RT_RADIX_TREE *
+#ifdef RT_SHMEM
+RT_CREATE(MemoryContext ctx, dsa_area *dsa)
+#else
 RT_CREATE(MemoryContext ctx)
+#endif
 {
 	RT_RADIX_TREE *tree;
 	MemoryContext old_ctx;
 
 	old_ctx = MemoryContextSwitchTo(ctx);
 
-	tree = palloc(sizeof(RT_RADIX_TREE));
+	tree = (RT_RADIX_TREE *) palloc0(sizeof(RT_RADIX_TREE));
 	tree->context = ctx;
-	tree->root = NULL;
-	tree->max_val = 0;
-	tree->num_keys = 0;
+
+#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);
+#else
+	tree->ctl = (RT_RADIX_TREE_CONTROL *) palloc0(sizeof(RT_RADIX_TREE_CONTROL));
 
 	/* Create the slab allocator for each size class */
 	for (int i = 0; i < RT_SIZE_CLASS_COUNT; i++)
@@ -1218,27 +1335,52 @@ RT_CREATE(MemoryContext ctx)
 												RT_SIZE_CLASS_INFO[i].name,
 												RT_SIZE_CLASS_INFO[i].leaf_blocksize,
 												RT_SIZE_CLASS_INFO[i].leaf_size);
-#ifdef RT_DEBUG
-		tree->cnt[i] = 0;
-#endif
 	}
+#endif
+
+	tree->ctl->root = RT_INVALID_PTR_ALLOC;
 
 	MemoryContextSwitchTo(old_ctx);
 
 	return tree;
 }
 
+#ifdef RT_SHMEM
+RT_RADIX_TREE *
+RT_ATTACH(dsa_area *dsa, dsa_pointer dp)
+{
+	RT_RADIX_TREE *tree;
+
+	/* 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);
+
+	/* XXX: do we need to set a callback on exit to detach dsa? */
+
+	return tree;
+}
+#endif
+
 /*
  * Free the given radix tree.
  */
 RT_SCOPE void
 RT_FREE(RT_RADIX_TREE *tree)
 {
+#ifdef RT_SHMEM
+	dsa_free(tree->dsa, tree->ctl_dp); // XXX
+	dsa_detach(tree->dsa);
+#else
+	pfree(tree->ctl);
+
 	for (int i = 0; i < RT_SIZE_CLASS_COUNT; i++)
 	{
 		MemoryContextDelete(tree->inner_slabs[i]);
 		MemoryContextDelete(tree->leaf_slabs[i]);
 	}
+#endif
 
 	pfree(tree);
 }
@@ -1252,46 +1394,50 @@ RT_SET(RT_RADIX_TREE *tree, uint64 key, uint64 value)
 {
 	int			shift;
 	bool		updated;
-	RT_PTR_LOCAL node;
 	RT_PTR_LOCAL parent;
+	RT_PTR_ALLOC nodep;
+	RT_PTR_LOCAL  node;
 
 	/* Empty tree, create the root */
-	if (!tree->root)
+	if (!RT_PTR_ALLOC_IS_VALID(tree->ctl->root))
 		RT_NEW_ROOT(tree, key);
 
 	/* Extend the tree if necessary */
-	if (key > tree->max_val)
+	if (key > tree->ctl->max_val)
 		RT_EXTEND(tree, key);
 
-	Assert(tree->root);
+	//Assert(tree->ctl->root);
 
-	shift = tree->root->shift;
-	node = parent = tree->root;
+	nodep = tree->ctl->root;
+	parent = RT_PTR_GET_LOCAL(tree, nodep);
+	shift = parent->shift;
 
 	/* Descend the tree until a leaf node */
 	while (shift >= 0)
 	{
-		RT_PTR_LOCAL child;
+		RT_PTR_ALLOC child;
+
+		node = RT_PTR_GET_LOCAL(tree, nodep);
 
 		if (NODE_IS_LEAF(node))
 			break;
 
 		if (!RT_NODE_SEARCH_INNER(node, key, &child))
 		{
-			RT_SET_EXTEND(tree, key, value, parent, node);
+			RT_SET_EXTEND(tree, key, value, parent, nodep, node);
 			return false;
 		}
 
 		parent = node;
-		node = child;
+		nodep = child;
 		shift -= RT_NODE_SPAN;
 	}
 
-	updated = RT_NODE_INSERT_LEAF(tree, parent, node, key, value);
+	updated = RT_NODE_INSERT_LEAF(tree, parent, nodep, node, key, value);
 
 	/* Update the statistics */
 	if (!updated)
-		tree->num_keys++;
+		tree->ctl->num_keys++;
 
 	return updated;
 }
@@ -1309,11 +1455,11 @@ RT_SEARCH(RT_RADIX_TREE *tree, uint64 key, uint64 *value_p)
 
 	Assert(value_p != NULL);
 
-	if (!tree->root || key > tree->max_val)
+	if (!RT_PTR_ALLOC_IS_VALID(tree->ctl->root) || key > tree->ctl->max_val)
 		return false;
 
-	node = tree->root;
-	shift = tree->root->shift;
+	node = RT_PTR_GET_LOCAL(tree, tree->ctl->root);
+	shift = node->shift;
 
 	/* Descend the tree until a leaf node */
 	while (shift >= 0)
@@ -1326,7 +1472,7 @@ RT_SEARCH(RT_RADIX_TREE *tree, uint64 key, uint64 *value_p)
 		if (!RT_NODE_SEARCH_INNER(node, key, &child))
 			return false;
 
-		node = child;
+		node = RT_PTR_GET_LOCAL(tree, child);
 		shift -= RT_NODE_SPAN;
 	}
 
@@ -1341,37 +1487,40 @@ RT_SCOPE bool
 RT_DELETE(RT_RADIX_TREE *tree, uint64 key)
 {
 	RT_PTR_LOCAL node;
+	RT_PTR_ALLOC allocnode;
 	RT_PTR_ALLOC stack[RT_MAX_LEVEL] = {0};
 	int			shift;
 	int			level;
 	bool		deleted;
 
-	if (!tree->root || key > tree->max_val)
+	if (!RT_PTR_ALLOC_IS_VALID(tree->ctl->root) || key > tree->ctl->max_val)
 		return false;
 
 	/*
 	 * Descend the tree to search the key while building a stack of nodes we
 	 * visited.
 	 */
-	node = tree->root;
-	shift = tree->root->shift;
+	allocnode = tree->ctl->root;
+	node = RT_PTR_GET_LOCAL(tree, allocnode);
+	shift = node->shift;
 	level = -1;
 	while (shift > 0)
 	{
 		RT_PTR_ALLOC child;
 
 		/* Push the current node to the stack */
-		stack[++level] = node;
+		stack[++level] = allocnode;
+		node = RT_PTR_GET_LOCAL(tree, allocnode);
 
 		if (!RT_NODE_SEARCH_INNER(node, key, &child))
 			return false;
 
-		node = child;
+		allocnode = child;
 		shift -= RT_NODE_SPAN;
 	}
 
 	/* Delete the key from the leaf node if exists */
-	Assert(NODE_IS_LEAF(node));
+	node = RT_PTR_GET_LOCAL(tree, allocnode);
 	deleted = RT_NODE_DELETE_LEAF(node, key);
 
 	if (!deleted)
@@ -1381,7 +1530,7 @@ RT_DELETE(RT_RADIX_TREE *tree, uint64 key)
 	}
 
 	/* Found the key to delete. Update the statistics */
-	tree->num_keys--;
+	tree->ctl->num_keys--;
 
 	/*
 	 * Return if the leaf node still has keys and we don't need to delete the
@@ -1391,13 +1540,14 @@ RT_DELETE(RT_RADIX_TREE *tree, uint64 key)
 		return true;
 
 	/* Free the empty leaf node */
-	RT_FREE_NODE(tree, node);
+	RT_FREE_NODE(tree, allocnode);
 
 	/* Delete the key in inner nodes recursively */
 	while (level >= 0)
 	{
-		node = stack[level--];
+		allocnode = stack[level--];
 
+		node = RT_PTR_GET_LOCAL(tree, allocnode);
 		deleted = RT_NODE_DELETE_INNER(node, key);
 		Assert(deleted);
 
@@ -1406,7 +1556,7 @@ RT_DELETE(RT_RADIX_TREE *tree, uint64 key)
 			break;
 
 		/* The node became empty */
-		RT_FREE_NODE(tree, node);
+		RT_FREE_NODE(tree, allocnode);
 	}
 
 	return true;
@@ -1478,6 +1628,7 @@ RT_BEGIN_ITERATE(RT_RADIX_TREE *tree)
 {
 	MemoryContext old_ctx;
 	RT_ITER    *iter;
+	RT_PTR_LOCAL root;
 	int			top_level;
 
 	old_ctx = MemoryContextSwitchTo(tree->context);
@@ -1486,17 +1637,18 @@ RT_BEGIN_ITERATE(RT_RADIX_TREE *tree)
 	iter->tree = tree;
 
 	/* empty tree */
-	if (!iter->tree->root)
+	if (!iter->tree->ctl->root)
 		return iter;
 
-	top_level = iter->tree->root->shift / RT_NODE_SPAN;
+	root = RT_PTR_GET_LOCAL(tree, iter->tree->ctl->root);
+	top_level = root->shift / RT_NODE_SPAN;
 	iter->stack_len = top_level;
 
 	/*
 	 * Descend to the left most leaf node from the root. The key is being
 	 * constructed while descending to the leaf.
 	 */
-	RT_UPDATE_ITER_STACK(iter, iter->tree->root, top_level);
+	RT_UPDATE_ITER_STACK(iter, root, top_level);
 
 	MemoryContextSwitchTo(old_ctx);
 
@@ -1511,7 +1663,7 @@ RT_SCOPE bool
 RT_ITERATE_NEXT(RT_ITER *iter, uint64 *key_p, uint64 *value_p)
 {
 	/* Empty tree */
-	if (!iter->tree->root)
+	if (!iter->tree->ctl->root)
 		return false;
 
 	for (;;)
@@ -1571,7 +1723,7 @@ RT_END_ITERATE(RT_ITER *iter)
 RT_SCOPE uint64
 RT_NUM_ENTRIES(RT_RADIX_TREE *tree)
 {
-	return tree->num_keys;
+	return tree->ctl->num_keys;
 }
 
 /*
@@ -1580,13 +1732,18 @@ RT_NUM_ENTRIES(RT_RADIX_TREE *tree)
 RT_SCOPE uint64
 RT_MEMORY_USAGE(RT_RADIX_TREE *tree)
 {
+	// XXX is this necessary?
 	Size		total = sizeof(RT_RADIX_TREE);
 
+#ifdef RT_SHMEM
+	total = dsa_get_total_size(tree->dsa);
+#else
 	for (int i = 0; i < RT_SIZE_CLASS_COUNT; i++)
 	{
 		total += MemoryContextMemAllocated(tree->inner_slabs[i], true);
 		total += MemoryContextMemAllocated(tree->leaf_slabs[i], true);
 	}
+#endif
 
 	return total;
 }
@@ -1670,13 +1827,13 @@ void
 rt_stats(RT_RADIX_TREE *tree)
 {
 	ereport(NOTICE, (errmsg("num_keys = " UINT64_FORMAT ", height = %d, n4 = %u, n15 = %u, n32 = %u, n125 = %u, n256 = %u",
-						 tree->num_keys,
-						 tree->root->shift / RT_NODE_SPAN,
-						 tree->cnt[RT_CLASS_4_FULL],
-						 tree->cnt[RT_CLASS_32_PARTIAL],
-						 tree->cnt[RT_CLASS_32_FULL],
-						 tree->cnt[RT_CLASS_125_FULL],
-						 tree->cnt[RT_CLASS_256])));
+						 tree->ctl->num_keys,
+						 tree->ctl->root->shift / RT_NODE_SPAN,
+						 tree->ctl->cnt[RT_CLASS_4_FULL],
+						 tree->ctl->cnt[RT_CLASS_32_PARTIAL],
+						 tree->ctl->cnt[RT_CLASS_32_FULL],
+						 tree->ctl->cnt[RT_CLASS_125_FULL],
+						 tree->ctl->cnt[RT_CLASS_256])));
 }
 
 static void
@@ -1848,23 +2005,23 @@ rt_dump_search(RT_RADIX_TREE *tree, uint64 key)
 
 	elog(NOTICE, "-----------------------------------------------------------");
 	elog(NOTICE, "max_val = " UINT64_FORMAT "(0x" UINT64_FORMAT_HEX ")",
-		 tree->max_val, tree->max_val);
+		 tree->ctl->max_val, tree->ctl->max_val);
 
-	if (!tree->root)
+	if (!tree->ctl->root)
 	{
 		elog(NOTICE, "tree is empty");
 		return;
 	}
 
-	if (key > tree->max_val)
+	if (key > tree->ctl->max_val)
 	{
 		elog(NOTICE, "key " UINT64_FORMAT "(0x" UINT64_FORMAT_HEX ") is larger than max val",
 			 key, key);
 		return;
 	}
 
-	node = tree->root;
-	shift = tree->root->shift;
+	node = tree->ctl->root;
+	shift = tree->ctl->root->shift;
 	while (shift >= 0)
 	{
 		RT_PTR_LOCAL child;
@@ -1901,15 +2058,15 @@ rt_dump(RT_RADIX_TREE *tree)
 				RT_SIZE_CLASS_INFO[i].inner_blocksize,
 				RT_SIZE_CLASS_INFO[i].leaf_size,
 				RT_SIZE_CLASS_INFO[i].leaf_blocksize);
-	fprintf(stderr, "max_val = " UINT64_FORMAT "\n", tree->max_val);
+	fprintf(stderr, "max_val = " UINT64_FORMAT "\n", tree->ctl->max_val);
 
-	if (!tree->root)
+	if (!tree->ctl->root)
 	{
 		fprintf(stderr, "empty tree\n");
 		return;
 	}
 
-	rt_dump_node(tree->root, 0, true);
+	rt_dump_node(tree->ctl->root, 0, true);
 }
 #endif
 
@@ -1931,6 +2088,7 @@ rt_dump(RT_RADIX_TREE *tree)
 
 /* type declarations */
 #undef RT_RADIX_TREE
+#undef RT_RADIX_TREE_CONTROL
 #undef RT_ITER
 #undef RT_NODE
 #undef RT_NODE_ITER
@@ -1959,6 +2117,7 @@ rt_dump(RT_RADIX_TREE *tree)
 /* function declarations */
 #undef RT_CREATE
 #undef RT_FREE
+#undef RT_ATTACH
 #undef RT_SET
 #undef RT_BEGIN_ITERATE
 #undef RT_ITERATE_NEXT
@@ -1980,6 +2139,8 @@ rt_dump(RT_RADIX_TREE *tree)
 #undef RT_GROW_NODE_KIND
 #undef RT_COPY_NODE
 #undef RT_REPLACE_NODE
+#undef RT_PTR_GET_LOCAL
+#undef RT_PTR_ALLOC_IS_VALID
 #undef RT_NODE_4_SEARCH_EQ
 #undef RT_NODE_32_SEARCH_EQ
 #undef RT_NODE_4_GET_INSERTPOS
@@ -2005,6 +2166,7 @@ rt_dump(RT_RADIX_TREE *tree)
 #undef RT_SHIFT_GET_MAX_VAL
 #undef RT_NODE_SEARCH_INNER
 #undef RT_NODE_SEARCH_LEAF
+#undef RT_NODE_UPDATE_INNER
 #undef RT_NODE_DELETE_INNER
 #undef RT_NODE_DELETE_LEAF
 #undef RT_NODE_INSERT_INNER
diff --git a/src/include/lib/radixtree_delete_impl.h b/src/include/lib/radixtree_delete_impl.h
index 6eefc63e19..eb87866b90 100644
--- a/src/include/lib/radixtree_delete_impl.h
+++ b/src/include/lib/radixtree_delete_impl.h
@@ -16,6 +16,12 @@
 
 	uint8		chunk = RT_GET_KEY_CHUNK(key, node->shift);
 
+#ifdef RT_NODE_LEVEL_LEAF
+	Assert(NODE_IS_LEAF(node));
+#else
+	Assert(!NODE_IS_LEAF(node));
+#endif
+
 	switch (node->kind)
 	{
 		case RT_NODE_KIND_4:
diff --git a/src/include/lib/radixtree_insert_impl.h b/src/include/lib/radixtree_insert_impl.h
index ff76583402..e4faf54d9d 100644
--- a/src/include/lib/radixtree_insert_impl.h
+++ b/src/include/lib/radixtree_insert_impl.h
@@ -14,11 +14,14 @@
 
 	uint8		chunk = RT_GET_KEY_CHUNK(key, node->shift);
 	bool		chunk_exists = false;
-	RT_NODE		*newnode = NULL;
+	RT_PTR_LOCAL newnode = NULL;
+	RT_PTR_ALLOC allocnode;
 
 #ifdef RT_NODE_LEVEL_LEAF
+	const bool inner = false;
 	Assert(NODE_IS_LEAF(node));
 #else
+	const bool inner = true;
 	Assert(!NODE_IS_LEAF(node));
 #endif
 
@@ -45,9 +48,15 @@
 				if (unlikely(!VAR_NODE_HAS_FREE_SLOT(n4)))
 				{
 					RT_NODE32_TYPE *new32;
+					const uint8 new_kind = RT_NODE_KIND_32;
+					const RT_SIZE_CLASS new_class = RT_KIND_MIN_SIZE_CLASS[new_kind];
 
 					/* grow node from 4 to 32 */
-					newnode = RT_GROW_NODE_KIND(tree, node, RT_NODE_KIND_32);
+					allocnode = RT_ALLOC_NODE(tree, new_class, inner);
+					newnode = RT_PTR_GET_LOCAL(tree, allocnode);
+					RT_INIT_NODE(newnode, new_kind, new_class, inner);
+					RT_COPY_NODE(newnode, node);
+					//newnode = RT_GROW_NODE_KIND(tree, node, RT_NODE_KIND_32);
 					new32 = (RT_NODE32_TYPE *) newnode;
 #ifdef RT_NODE_LEVEL_LEAF
 					RT_CHUNK_VALUES_ARRAY_COPY(n4->base.chunks, n4->values,
@@ -57,7 +66,7 @@
 											  new32->base.chunks, new32->children);
 #endif
 					Assert(parent != NULL);
-					RT_REPLACE_NODE(tree, parent, node, newnode, key);
+					RT_REPLACE_NODE(tree, parent, nodep, allocnode, key);
 					node = newnode;
 				}
 				else
@@ -112,17 +121,19 @@
 					n32->base.n.fanout == class32_min.fanout)
 				{
 					/* grow to the next size class of this kind */
+					const RT_SIZE_CLASS new_class = RT_CLASS_32_FULL;
+
+					allocnode = RT_ALLOC_NODE(tree, new_class, inner);
+					newnode = RT_PTR_GET_LOCAL(tree, allocnode);
 #ifdef RT_NODE_LEVEL_LEAF
-					newnode = RT_ALLOC_NODE(tree, RT_CLASS_32_FULL, false);
 					memcpy(newnode, node, class32_min.leaf_size);
 #else
-					newnode = RT_ALLOC_NODE(tree, RT_CLASS_32_FULL, true);
 					memcpy(newnode, node, class32_min.inner_size);
 #endif
 					newnode->fanout = class32_max.fanout;
 
 					Assert(parent != NULL);
-					RT_REPLACE_NODE(tree, parent, node, newnode, key);
+					RT_REPLACE_NODE(tree, parent, nodep, allocnode, key);
 					node = newnode;
 
 					/* also update pointer for this kind */
@@ -132,11 +143,17 @@
 				if (unlikely(!VAR_NODE_HAS_FREE_SLOT(n32)))
 				{
 					RT_NODE125_TYPE *new125;
+					const uint8 new_kind = RT_NODE_KIND_125;
+					const RT_SIZE_CLASS new_class = RT_KIND_MIN_SIZE_CLASS[new_kind];
 
 					Assert(n32->base.n.fanout == class32_max.fanout);
 
 					/* grow node from 32 to 125 */
-					newnode = RT_GROW_NODE_KIND(tree, node, RT_NODE_KIND_125);
+					allocnode = RT_ALLOC_NODE(tree, new_class, inner);
+					newnode = RT_PTR_GET_LOCAL(tree, allocnode);
+					RT_INIT_NODE(newnode, new_kind, new_class, inner);
+					RT_COPY_NODE(newnode, node);
+					//newnode = RT_GROW_NODE_KIND(tree, node, RT_NODE_KIND_125);
 					new125 = (RT_NODE125_TYPE *) newnode;
 
 					for (int i = 0; i < class32_max.fanout; i++)
@@ -153,7 +170,7 @@
 					new125->base.isset[0] = (bitmapword) (((uint64) 1 << class32_max.fanout) - 1);
 
 					Assert(parent != NULL);
-					RT_REPLACE_NODE(tree, parent, node, newnode, key);
+					RT_REPLACE_NODE(tree, parent, nodep, allocnode, key);
 					node = newnode;
 				}
 				else
@@ -204,9 +221,15 @@
 				if (unlikely(!VAR_NODE_HAS_FREE_SLOT(n125)))
 				{
 					RT_NODE256_TYPE *new256;
+					const uint8 new_kind = RT_NODE_KIND_256;
+					const RT_SIZE_CLASS new_class = RT_KIND_MIN_SIZE_CLASS[new_kind];
 
 					/* grow node from 125 to 256 */
-					newnode = RT_GROW_NODE_KIND(tree, node, RT_NODE_KIND_256);
+					allocnode = RT_ALLOC_NODE(tree, new_class, inner);
+					newnode = RT_PTR_GET_LOCAL(tree, allocnode);
+					RT_INIT_NODE(newnode, new_kind, new_class, inner);
+					RT_COPY_NODE(newnode, node);
+					//newnode = RT_GROW_NODE_KIND(tree, node, RT_NODE_KIND_256);
 					new256 = (RT_NODE256_TYPE *) newnode;
 					for (int i = 0; i < RT_NODE_MAX_SLOTS && cnt < n125->base.n.count; i++)
 					{
@@ -221,7 +244,7 @@
 					}
 
 					Assert(parent != NULL);
-					RT_REPLACE_NODE(tree, parent, node, newnode, key);
+					RT_REPLACE_NODE(tree, parent, nodep, allocnode, key);
 					node = newnode;
 				}
 				else
diff --git a/src/include/lib/radixtree_iter_impl.h b/src/include/lib/radixtree_iter_impl.h
index a153011376..09d2018dc0 100644
--- a/src/include/lib/radixtree_iter_impl.h
+++ b/src/include/lib/radixtree_iter_impl.h
@@ -12,13 +12,18 @@
 #error node level must be either inner or leaf
 #endif
 
+	bool		found = false;
+	uint8		key_chunk;
+
 #ifdef RT_NODE_LEVEL_LEAF
 	uint64		value;
+
+	Assert(NODE_IS_LEAF(node_iter->node));
 #else
-	RT_NODE    *child = NULL;
+	RT_PTR_LOCAL child = NULL;
+
+	Assert(!NODE_IS_LEAF(node_iter->node));
 #endif
-	bool		found = false;
-	uint8		key_chunk;
 
 	switch (node_iter->node->kind)
 	{
@@ -32,7 +37,7 @@
 #ifdef RT_NODE_LEVEL_LEAF
 				value = n4->values[node_iter->current_idx];
 #else
-				child = n4->children[node_iter->current_idx];
+				child = RT_PTR_GET_LOCAL(iter->tree, n4->children[node_iter->current_idx]);
 #endif
 				key_chunk = n4->base.chunks[node_iter->current_idx];
 				found = true;
@@ -49,7 +54,7 @@
 #ifdef RT_NODE_LEVEL_LEAF
 				value = n32->values[node_iter->current_idx];
 #else
-				child = n32->children[node_iter->current_idx];
+				child = RT_PTR_GET_LOCAL(iter->tree, n32->children[node_iter->current_idx]);
 #endif
 				key_chunk = n32->base.chunks[node_iter->current_idx];
 				found = true;
@@ -73,7 +78,7 @@
 #ifdef RT_NODE_LEVEL_LEAF
 				value = RT_NODE_LEAF_125_GET_VALUE(n125, i);
 #else
-				child = RT_NODE_INNER_125_GET_CHILD(n125, i);
+				child = RT_PTR_GET_LOCAL(iter->tree, RT_NODE_INNER_125_GET_CHILD(n125, i));
 #endif
 				key_chunk = i;
 				found = true;
@@ -101,7 +106,7 @@
 #ifdef RT_NODE_LEVEL_LEAF
 				value = RT_NODE_LEAF_256_GET_VALUE(n256, i);
 #else
-				child = RT_NODE_INNER_256_GET_CHILD(n256, i);
+				child = RT_PTR_GET_LOCAL(iter->tree, RT_NODE_INNER_256_GET_CHILD(n256, i));
 #endif
 				key_chunk = i;
 				found = true;
diff --git a/src/include/lib/radixtree_search_impl.h b/src/include/lib/radixtree_search_impl.h
index cbc357dcc8..3e97c31c2c 100644
--- a/src/include/lib/radixtree_search_impl.h
+++ b/src/include/lib/radixtree_search_impl.h
@@ -16,8 +16,13 @@
 
 #ifdef RT_NODE_LEVEL_LEAF
 	uint64		value = 0;
+
+	Assert(NODE_IS_LEAF(node));
 #else
-	RT_PTR_LOCAL child = NULL;
+#ifndef RT_ACTION_UPDATE
+	RT_PTR_ALLOC child = RT_INVALID_PTR_ALLOC;
+#endif
+	Assert(!NODE_IS_LEAF(node));
 #endif
 
 	switch (node->kind)
@@ -32,8 +37,12 @@
 
 #ifdef RT_NODE_LEVEL_LEAF
 				value = n4->values[idx];
+#else
+#ifdef RT_ACTION_UPDATE
+				n4->children[idx] = new_child;
 #else
 				child = n4->children[idx];
+#endif
 #endif
 				break;
 			}
@@ -47,22 +56,31 @@
 
 #ifdef RT_NODE_LEVEL_LEAF
 				value = n32->values[idx];
+#else
+#ifdef RT_ACTION_UPDATE
+				n32->children[idx] = new_child;
 #else
 				child = n32->children[idx];
+#endif
 #endif
 				break;
 			}
 		case RT_NODE_KIND_125:
 			{
 				RT_NODE125_TYPE *n125 = (RT_NODE125_TYPE *) node;
+				int			slotpos = n125->base.slot_idxs[chunk];
 
-				if (!RT_NODE_125_IS_CHUNK_USED((RT_NODE_BASE_125 *) n125, chunk))
+				if (slotpos == RT_NODE_125_INVALID_IDX)
 					return false;
 
 #ifdef RT_NODE_LEVEL_LEAF
 				value = RT_NODE_LEAF_125_GET_VALUE(n125, chunk);
+#else
+#ifdef RT_ACTION_UPDATE
+				n125->children[slotpos] = new_child;
 #else
 				child = RT_NODE_INNER_125_GET_CHILD(n125, chunk);
+#endif
 #endif
 				break;
 			}
@@ -79,19 +97,25 @@
 
 #ifdef RT_NODE_LEVEL_LEAF
 				value = RT_NODE_LEAF_256_GET_VALUE(n256, chunk);
+#else
+#ifdef RT_ACTION_UPDATE
+				RT_NODE_INNER_256_SET(n256, chunk, new_child);
 #else
 				child = RT_NODE_INNER_256_GET_CHILD(n256, chunk);
+#endif
 #endif
 				break;
 			}
 	}
 
+#ifndef RT_ACTION_UPDATE
 #ifdef RT_NODE_LEVEL_LEAF
 	Assert(value_p != NULL);
 	*value_p = value;
 #else
 	Assert(child_p != NULL);
 	*child_p = child;
+#endif
 #endif
 
 	return true;
diff --git a/src/include/utils/dsa.h b/src/include/utils/dsa.h
index 104386e674..c67f936880 100644
--- a/src/include/utils/dsa.h
+++ b/src/include/utils/dsa.h
@@ -117,6 +117,7 @@ extern dsa_handle dsa_get_handle(dsa_area *area);
 extern dsa_pointer dsa_allocate_extended(dsa_area *area, size_t size, int flags);
 extern void dsa_free(dsa_area *area, dsa_pointer dp);
 extern void *dsa_get_address(dsa_area *area, dsa_pointer dp);
+extern size_t dsa_get_total_size(dsa_area *area);
 extern void dsa_trim(dsa_area *area);
 extern void dsa_dump(dsa_area *area);
 
diff --git a/src/test/modules/test_radixtree/test_radixtree.c b/src/test/modules/test_radixtree/test_radixtree.c
index 2256d08100..61d842789d 100644
--- a/src/test/modules/test_radixtree/test_radixtree.c
+++ b/src/test/modules/test_radixtree/test_radixtree.c
@@ -18,6 +18,7 @@
 #include "nodes/bitmapset.h"
 #include "storage/block.h"
 #include "storage/itemptr.h"
+#include "storage/lwlock.h"
 #include "utils/memutils.h"
 #include "utils/timestamp.h"
 
@@ -103,6 +104,8 @@ static const test_spec test_specs[] = {
 #define RT_SCOPE static
 #define RT_DECLARE
 #define RT_DEFINE
+// WIP: compiles with warnings because rt_attach is defined but not used
+// #define RT_SHMEM
 #include "lib/radixtree.h"
 
 
@@ -119,7 +122,15 @@ test_empty(void)
 	uint64		key;
 	uint64		val;
 
+#ifdef RT_SHMEM
+	int			tranche_id = LWLockNewTrancheId();
+	dsa_area   *dsa;
+	dsa = dsa_create(tranche_id);
+
+	radixtree = rt_create(CurrentMemoryContext, dsa);
+#else
 	radixtree = rt_create(CurrentMemoryContext);
+#endif
 
 	if (rt_search(radixtree, 0, &dummy))
 		elog(ERROR, "rt_search on empty tree returned true");
@@ -153,10 +164,20 @@ test_basic(int children, bool test_inner)
 	uint64 *keys;
 	int	shift = test_inner ? 8 : 0;
 
+#ifdef RT_SHMEM
+	int			tranche_id = LWLockNewTrancheId();
+	dsa_area   *dsa;
+	dsa = dsa_create(tranche_id);
+#endif
+
 	elog(NOTICE, "testing basic operations with %s node %d",
 		 test_inner ? "inner" : "leaf", children);
 
+#ifdef RT_SHMEM
+	radixtree = rt_create(CurrentMemoryContext, dsa);
+#else
 	radixtree = rt_create(CurrentMemoryContext);
+#endif
 
 	/* prepare keys in order like 1, 32, 2, 31, 2, ... */
 	keys = palloc(sizeof(uint64) * children);
@@ -297,9 +318,19 @@ test_node_types(uint8 shift)
 {
 	rt_radix_tree *radixtree;
 
+#ifdef RT_SHMEM
+	int			tranche_id = LWLockNewTrancheId();
+	dsa_area   *dsa;
+	dsa = dsa_create(tranche_id);
+#endif
+
 	elog(NOTICE, "testing radix tree node types with shift \"%d\"", shift);
 
+#ifdef RT_SHMEM
+	radixtree = rt_create(CurrentMemoryContext, dsa);
+#else
 	radixtree = rt_create(CurrentMemoryContext);
+#endif
 
 	/*
 	 * Insert and search entries for every node type at the 'shift' level,
@@ -332,6 +363,11 @@ test_pattern(const test_spec * spec)
 	int			patternlen;
 	uint64	   *pattern_values;
 	uint64		pattern_num_values;
+#ifdef RT_SHMEM
+	int			tranche_id = LWLockNewTrancheId();
+	dsa_area   *dsa;
+	dsa = dsa_create(tranche_id);
+#endif
 
 	elog(NOTICE, "testing radix tree with pattern \"%s\"", spec->test_name);
 	if (rt_test_stats)
@@ -357,7 +393,13 @@ test_pattern(const test_spec * spec)
 										  "radixtree test",
 										  ALLOCSET_SMALL_SIZES);
 	MemoryContextSetIdentifier(radixtree_ctx, spec->test_name);
+
+#ifdef RT_SHMEM
+	radixtree = rt_create(radixtree_ctx, dsa);
+#else
 	radixtree = rt_create(radixtree_ctx);
+#endif
+
 
 	/*
 	 * Add values to the set.
@@ -563,6 +605,7 @@ test_pattern(const test_spec * spec)
 		elog(ERROR, "rt_num_entries returned " UINT64_FORMAT ", expected " UINT64_FORMAT "after " UINT64_FORMAT " deletion",
 			 nafter, (nbefore - ndeleted), ndeleted);
 
+	rt_free(radixtree);
 	MemoryContextDelete(radixtree_ctx);
 }
 
-- 
2.39.0

