Get rid of radix tree's general purpose memory context

Previously, this was notionally used only for the entry point of the
tree and as a convenient parent for other contexts.

For shared memory, the creator previously allocated the entry point
in this context, but attaching backends didn't have access to that,
so they just used the caller's context. For the sake of consistency,
allocate every instance of an entry point in the caller's context.

For local memory, allocate the control object in the caller's context
as well. This commit also makes the "leaf context" the notional parent
of the child contexts used for nodes, so it's a bit of a misnomer,
but a future commit will make the node contexts independent rather
than children, so leave it this way for now to avoid code churn.

The memory context parameter for RT_CREATE is now unused in the case
of shared memory, so remove it and adjust callers to match.

In passing, remove unused "context" member from struct TidStore,
which seems to have been an oversight.

Reviewed by Masahiko Sawada

Discussion: https://postgr.es/m/CANWCAZZDCo4k5oURg_pPxM6+WZ1oiG=sqgjmQiELuyP0Vtrwig@mail.gmail.com

Author: j-naylor

src/backend/access/common/tidstore.c

@@ -113,10 +113,10 @@ typedef struct BlocktableEntry
 /* Per-backend state for a TidStore */
 struct TidStore
 {
-	/* MemoryContext where the TidStore is allocated */
-	MemoryContext context;
-
-	/* MemoryContext that the radix tree uses */
+	/*
+	 * MemoryContext for the radix tree when using local memory, NULL for
+	 * shared memory
+	 */
 	MemoryContext rt_context;
 
 	/* Storage for TIDs. Use either one depending on TidStoreIsShared() */
@@ -167,7 +167,6 @@ TidStoreCreateLocal(size_t max_bytes, bool insert_only)
 	size_t		maxBlockSize = ALLOCSET_DEFAULT_MAXSIZE;
 
 	ts = palloc0(sizeof(TidStore));
-	ts->context = CurrentMemoryContext;
 
 	/* choose the maxBlockSize to be no larger than 1/16 of max_bytes */
 	while (16 * maxBlockSize > max_bytes)
@@ -201,8 +200,7 @@ TidStoreCreateLocal(size_t max_bytes, bool insert_only)
 
 /*
  * Similar to TidStoreCreateLocal() but create a shared TidStore on a
- * DSA area. The TID storage will live in the DSA area, and the memory
- * context rt_context will have only meta data of the radix tree.
+ * DSA area.
  *
  * The returned object is allocated in backend-local memory.
  */
@@ -215,11 +213,6 @@ TidStoreCreateShared(size_t max_bytes, int tranche_id)
 	size_t		dsa_max_size = DSA_MAX_SEGMENT_SIZE;
 
 	ts = palloc0(sizeof(TidStore));
-	ts->context = CurrentMemoryContext;
-
-	ts->rt_context = AllocSetContextCreate(CurrentMemoryContext,
-										   "TID storage meta data",
-										   ALLOCSET_SMALL_SIZES);
 
 	/*
 	 * Choose the initial and maximum DSA segment sizes to be no longer than
@@ -235,8 +228,7 @@ TidStoreCreateShared(size_t max_bytes, int tranche_id)
 		dsa_init_size = dsa_max_size;
 
 	area = dsa_create_ext(tranche_id, dsa_init_size, dsa_max_size);
-	ts->tree.shared = shared_ts_create(ts->rt_context, area,
-									   tranche_id);
+	ts->tree.shared = shared_ts_create(area, tranche_id);
 	ts->area = area;
 
 	return ts;
@@ -328,13 +320,13 @@ TidStoreDestroy(TidStore *ts)
 	if (TidStoreIsShared(ts))
 	{
 		shared_ts_free(ts->tree.shared);
-
 		dsa_detach(ts->area);
 	}
 	else
+	{
 		local_ts_free(ts->tree.local);
-
-	MemoryContextDelete(ts->rt_context);
+		MemoryContextDelete(ts->rt_context);
+	}
 
 	pfree(ts);
 }

src/include/lib/radixtree.h

@@ -275,7 +275,7 @@ typedef dsa_pointer RT_HANDLE;
 #endif
 
 #ifdef RT_SHMEM
-RT_SCOPE	RT_RADIX_TREE *RT_CREATE(MemoryContext ctx, dsa_area *dsa, int tranche_id);
+RT_SCOPE	RT_RADIX_TREE *RT_CREATE(dsa_area *dsa, int tranche_id);
 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);
@@ -706,8 +706,6 @@ typedef struct RT_RADIX_TREE_CONTROL
 /* Entry point for allocating and accessing the tree */
 struct RT_RADIX_TREE
 {
-	MemoryContext context;
-
 	/* pointing to either local memory or DSA */
 	RT_RADIX_TREE_CONTROL *ctl;
 
@@ -1809,31 +1807,25 @@ RT_SET(RT_RADIX_TREE * tree, uint64 key, RT_VALUE_TYPE * value_p)
 /***************** SETUP / TEARDOWN *****************/
 
 /*
- * Create the radix tree in the given memory context and return it.
+ * Create the radix tree root in the caller's memory context and return it.
  *
- * All local memory required for a radix tree is allocated in the given
- * memory context and its children. Note that RT_FREE() will delete all
- * allocated space within the given memory context, so the dsa_area should
- * be created in a different context.
+ * The tree's nodes and leaves are allocated in "ctx" and its children for
+ * local memory, or in "dsa" for shared memory.
  */
 RT_SCOPE	RT_RADIX_TREE *
 #ifdef RT_SHMEM
-RT_CREATE(MemoryContext ctx, dsa_area *dsa, int tranche_id)
+RT_CREATE(dsa_area *dsa, int tranche_id)
 #else
 RT_CREATE(MemoryContext ctx)
 #endif
 {
 	RT_RADIX_TREE *tree;
-	MemoryContext old_ctx;
 	RT_CHILD_PTR rootnode;
 #ifdef RT_SHMEM
 	dsa_pointer dp;
 #endif
 
-	old_ctx = MemoryContextSwitchTo(ctx);
-
 	tree = (RT_RADIX_TREE *) palloc0(sizeof(RT_RADIX_TREE));
-	tree->context = ctx;
 
 #ifdef RT_SHMEM
 	tree->dsa = dsa;
@@ -1858,7 +1850,7 @@ RT_CREATE(MemoryContext ctx)
 	}
 
 	/* By default we use the passed context for leaves. */
-	tree->leaf_context = tree->context;
+	tree->leaf_context = ctx;
 
 #ifndef RT_VARLEN_VALUE_SIZE
 
@@ -1880,8 +1872,6 @@ RT_CREATE(MemoryContext ctx)
 	tree->ctl->start_shift = 0;
 	tree->ctl->max_val = RT_SHIFT_GET_MAX_VAL(0);
 
-	MemoryContextSwitchTo(old_ctx);
-
 	return tree;
 }
 
@@ -2054,13 +2044,16 @@ RT_FREE(RT_RADIX_TREE * tree)
 	 */
 	tree->ctl->magic = 0;
 	dsa_free(tree->dsa, tree->ctl->handle);
-#endif
-
+#else
 	/*
-	 * Free all space allocated within the tree's context and delete all child
+	 * Free all space allocated within the leaf context and delete all child
 	 * contexts such as those used for nodes.
 	 */
-	MemoryContextReset(tree->context);
+	MemoryContextReset(tree->leaf_context);
+
+	pfree(tree->ctl);
+#endif
+	pfree(tree);
 }
 
 /***************** ITERATION *****************/
@@ -2674,7 +2667,7 @@ RT_MEMORY_USAGE(RT_RADIX_TREE * tree)
 	Assert(tree->ctl->magic == RT_RADIX_TREE_MAGIC);
 	total = dsa_get_total_size(tree->dsa);
 #else
-	total = MemoryContextMemAllocated(tree->context, true);
+	total = MemoryContextMemAllocated(tree->leaf_context, true);
 #endif
 
 	return total;

src/test/modules/test_radixtree/test_radixtree.c

@@ -120,25 +120,22 @@ PG_FUNCTION_INFO_V1(test_radixtree);
 static void
 test_empty(void)
 {
-	MemoryContext radixtree_ctx;
 	rt_radix_tree *radixtree;
 	rt_iter    *iter;
 	uint64		key;
 #ifdef TEST_SHARED_RT
 	int			tranche_id = LWLockNewTrancheId();
 	dsa_area   *dsa;
-#endif
-
-	radixtree_ctx = AllocSetContextCreate(CurrentMemoryContext,
-										  "test_radix_tree",
-										  ALLOCSET_SMALL_SIZES);
 
-#ifdef TEST_SHARED_RT
 	LWLockRegisterTranche(tranche_id, "test_radix_tree");
 	dsa = dsa_create(tranche_id);
-
-	radixtree = rt_create(radixtree_ctx, dsa, tranche_id);
+	radixtree = rt_create(dsa, tranche_id);
 #else
+	MemoryContext radixtree_ctx;
+
+	radixtree_ctx = AllocSetContextCreate(CurrentMemoryContext,
+										  "test_radix_tree",
+										  ALLOCSET_SMALL_SIZES);
 	radixtree = rt_create(radixtree_ctx);
 #endif
 
@@ -165,26 +162,23 @@ test_empty(void)
 static void
 test_basic(rt_node_class_test_elem *test_info, int shift, bool asc)
 {
-	MemoryContext radixtree_ctx;
 	rt_radix_tree *radixtree;
 	rt_iter    *iter;
 	uint64	   *keys;
 	int			children = test_info->nkeys;
 #ifdef TEST_SHARED_RT
 	int			tranche_id = LWLockNewTrancheId();
 	dsa_area   *dsa;
-#endif
 
-	radixtree_ctx = AllocSetContextCreate(CurrentMemoryContext,
-										  "test_radix_tree",
-										  ALLOCSET_SMALL_SIZES);
-
-#ifdef TEST_SHARED_RT
 	LWLockRegisterTranche(tranche_id, "test_radix_tree");
 	dsa = dsa_create(tranche_id);
-
-	radixtree = rt_create(radixtree_ctx, dsa, tranche_id);
+	radixtree = rt_create(dsa, tranche_id);
 #else
+	MemoryContext radixtree_ctx;
+
+	radixtree_ctx = AllocSetContextCreate(CurrentMemoryContext,
+										  "test_radix_tree",
+										  ALLOCSET_SMALL_SIZES);
 	radixtree = rt_create(radixtree_ctx);
 #endif
 
@@ -300,7 +294,6 @@ key_cmp(const void *a, const void *b)
 static void
 test_random(void)
 {
-	MemoryContext radixtree_ctx;
 	rt_radix_tree *radixtree;
 	rt_iter    *iter;
 	pg_prng_state state;
@@ -313,18 +306,16 @@ test_random(void)
 #ifdef TEST_SHARED_RT
 	int			tranche_id = LWLockNewTrancheId();
 	dsa_area   *dsa;
-#endif
 
-	radixtree_ctx = AllocSetContextCreate(CurrentMemoryContext,
-										  "test_radix_tree",
-										  ALLOCSET_SMALL_SIZES);
-
-#ifdef TEST_SHARED_RT
 	LWLockRegisterTranche(tranche_id, "test_radix_tree");
 	dsa = dsa_create(tranche_id);
-
-	radixtree = rt_create(radixtree_ctx, dsa, tranche_id);
+	radixtree = rt_create(dsa, tranche_id);
 #else
+	MemoryContext radixtree_ctx;
+
+	radixtree_ctx = AllocSetContextCreate(CurrentMemoryContext,
+										  "test_radix_tree",
+										  ALLOCSET_SMALL_SIZES);
 	radixtree = rt_create(radixtree_ctx);
 #endif