v31-0010-Radix-tree-optionally-tracks-memory-usage-when-R.patch

From 7da5e7808ba51aed7ad22b9758b3200cbfcd7d19 Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Wed, 8 Mar 2023 15:08:19 +0900
Subject: [PATCH v31 10/14] Radix tree optionally tracks memory usage, when
 RT_MEASURE_MEMORY_USAGE.

---
 contrib/bench_radix_tree/bench_radix_tree.c   |  1 +
 src/backend/utils/mmgr/dsa.c                  | 12 ---
 src/include/lib/radixtree.h                   | 93 +++++++++++++++++--
 src/include/utils/dsa.h                       |  1 -
 .../modules/test_radixtree/test_radixtree.c   |  1 +
 5 files changed, 85 insertions(+), 23 deletions(-)

diff --git a/contrib/bench_radix_tree/bench_radix_tree.c b/contrib/bench_radix_tree/bench_radix_tree.c
index 6e5149e2c4..8a0c754a2c 100644
--- a/contrib/bench_radix_tree/bench_radix_tree.c
+++ b/contrib/bench_radix_tree/bench_radix_tree.c
@@ -34,6 +34,7 @@ PG_MODULE_MAGIC;
 #define RT_DECLARE
 #define RT_DEFINE
 #define RT_USE_DELETE
+#define RT_MEASURE_MEMORY_USAGE
 #define RT_VALUE_TYPE uint64
 // WIP: compiles with warnings because rt_attach is defined but not used
 // #define RT_SHMEM
diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index 80555aefff..f5a62061a3 100644
--- a/src/backend/utils/mmgr/dsa.c
+++ b/src/backend/utils/mmgr/dsa.c
@@ -1024,18 +1024,6 @@ 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 8bea606c62..f7812eb12a 100644
--- a/src/include/lib/radixtree.h
+++ b/src/include/lib/radixtree.h
@@ -84,7 +84,6 @@
  * RT_BEGIN_ITERATE	- Begin iterating through all key-value pairs
  * RT_ITERATE_NEXT	- Return next key-value pair, if any
  * RT_END_ITERATE	- End iteration
- * RT_MEMORY_USAGE	- Get the memory usage
  *
  * Interface for Shared Memory
  * ---------
@@ -97,6 +96,8 @@
  * ---------
  *
  * RT_DELETE		- Delete a key-value pair. Declared/define if RT_USE_DELETE is defined
+ * RT_MEMORY_USAGE	- Get the memory usage. Declared/define if
+ *					  RT_MEASURE_MEMORY_USAGE is defined.
  *
  *
  * Copyright (c) 2023, PostgreSQL Global Development Group
@@ -138,7 +139,9 @@
 #ifdef RT_USE_DELETE
 #define RT_DELETE RT_MAKE_NAME(delete)
 #endif
+#ifdef RT_MEASURE_MEMORY_USAGE
 #define RT_MEMORY_USAGE RT_MAKE_NAME(memory_usage)
+#endif
 #ifdef RT_DEBUG
 #define RT_DUMP RT_MAKE_NAME(dump)
 #define RT_DUMP_NODE RT_MAKE_NAME(dump_node)
@@ -150,6 +153,9 @@
 #define RT_NEW_ROOT RT_MAKE_NAME(new_root)
 #define RT_ALLOC_NODE RT_MAKE_NAME(alloc_node)
 #define RT_INIT_NODE RT_MAKE_NAME(init_node)
+#ifdef RT_MEASURE_MEMORY_USAGE
+#define RT_FANOUT_GET_NODE_SIZE RT_MAKE_NAME(fanout_get_node_size)
+#endif
 #define RT_FREE_NODE RT_MAKE_NAME(free_node)
 #define RT_FREE_RECURSE RT_MAKE_NAME(free_recurse)
 #define RT_EXTEND_UP RT_MAKE_NAME(extend_up)
@@ -255,7 +261,9 @@ RT_SCOPE RT_ITER * RT_BEGIN_ITERATE(RT_RADIX_TREE *tree);
 RT_SCOPE bool RT_ITERATE_NEXT(RT_ITER *iter, uint64 *key_p, RT_VALUE_TYPE *value_p);
 RT_SCOPE void RT_END_ITERATE(RT_ITER *iter);
 
+#ifdef RT_MEASURE_MEMORY_USAGE
 RT_SCOPE uint64 RT_MEMORY_USAGE(RT_RADIX_TREE *tree);
+#endif
 
 #ifdef RT_DEBUG
 RT_SCOPE void RT_DUMP(RT_RADIX_TREE *tree);
@@ -624,6 +632,10 @@ typedef struct RT_RADIX_TREE_CONTROL
 	uint64		max_val;
 	uint64		num_keys;
 
+#ifdef RT_MEASURE_MEMORY_USAGE
+	int64		mem_used;
+#endif
+
 	/* statistics */
 #ifdef RT_DEBUG
 	int32		cnt[RT_SIZE_CLASS_COUNT];
@@ -1089,6 +1101,11 @@ RT_ALLOC_NODE(RT_RADIX_TREE *tree, RT_SIZE_CLASS size_class, bool is_leaf)
 													  allocsize);
 #endif
 
+#ifdef RT_MEASURE_MEMORY_USAGE
+	/* update memory usage */
+	tree->ctl->mem_used += allocsize;
+#endif
+
 #ifdef RT_DEBUG
 	/* update the statistics */
 	tree->ctl->cnt[size_class]++;
@@ -1165,6 +1182,54 @@ RT_SWITCH_NODE_KIND(RT_RADIX_TREE *tree, RT_PTR_ALLOC allocnode, RT_PTR_LOCAL no
 	return newnode;
 }
 
+#ifdef RT_MEASURE_MEMORY_USAGE
+/* Return the node size of the given fanout of the size class */
+static inline Size
+RT_FANOUT_GET_NODE_SIZE(int fanout, bool is_leaf)
+{
+	const Size fanout_inner_node_size[] = {
+		[3] = RT_SIZE_CLASS_INFO[RT_CLASS_3].inner_size,
+		[15] = RT_SIZE_CLASS_INFO[RT_CLASS_32_MIN].inner_size,
+		[32] = RT_SIZE_CLASS_INFO[RT_CLASS_32_MAX].inner_size,
+		[125] = RT_SIZE_CLASS_INFO[RT_CLASS_125].inner_size,
+		[256] = RT_SIZE_CLASS_INFO[RT_CLASS_256].inner_size,
+	};
+	const Size fanout_leaf_node_size[] = {
+		[3] = RT_SIZE_CLASS_INFO[RT_CLASS_3].leaf_size,
+		[15] = RT_SIZE_CLASS_INFO[RT_CLASS_32_MIN].leaf_size,
+		[32] = RT_SIZE_CLASS_INFO[RT_CLASS_32_MAX].leaf_size,
+		[125] = RT_SIZE_CLASS_INFO[RT_CLASS_125].leaf_size,
+		[256] = RT_SIZE_CLASS_INFO[RT_CLASS_256].leaf_size,
+	};
+	Size node_size;
+
+	node_size = is_leaf ?
+		fanout_leaf_node_size[fanout] : fanout_inner_node_size[fanout];
+
+#ifdef USE_ASSERT_CHECKING
+	{
+		Size assert_node_size = 0;
+
+		for (int i = 0; i < RT_SIZE_CLASS_COUNT; i++)
+		{
+			RT_SIZE_CLASS_ELEM size_class = RT_SIZE_CLASS_INFO[i];
+
+			if (size_class.fanout == fanout)
+			{
+				assert_node_size = is_leaf ?
+					size_class.leaf_size : size_class.inner_size;
+				break;
+			}
+		}
+
+		Assert(node_size == assert_node_size);
+	}
+#endif
+
+	return node_size;
+}
+#endif
+
 /* Free the given node */
 static void
 RT_FREE_NODE(RT_RADIX_TREE *tree, RT_PTR_ALLOC allocnode)
@@ -1197,11 +1262,22 @@ RT_FREE_NODE(RT_RADIX_TREE *tree, RT_PTR_ALLOC allocnode)
 	}
 #endif
 
+#ifdef RT_MEASURE_MEMORY_USAGE
+	/* update memory usage */
+	{
+		RT_PTR_LOCAL node = RT_PTR_GET_LOCAL(tree, allocnode);
+		tree->ctl->mem_used -= RT_FANOUT_GET_NODE_SIZE(node->fanout,
+													   RT_NODE_IS_LEAF(node));
+		Assert(tree->ctl->mem_used >= 0);
+	}
+#endif
+
 #ifdef RT_SHMEM
 	dsa_free(tree->dsa, allocnode);
 #else
 	pfree(allocnode);
 #endif
+
 }
 
 /* Update the parent's pointer when growing a node */
@@ -1989,27 +2065,23 @@ RT_END_ITERATE(RT_ITER *iter)
 /*
  * Return the statistics of the amount of memory used by the radix tree.
  */
+#ifdef RT_MEASURE_MEMORY_USAGE
 RT_SCOPE uint64
 RT_MEMORY_USAGE(RT_RADIX_TREE *tree)
 {
 	Size		total = 0;
 
-	RT_LOCK_SHARED(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++)
-	{
-		total += MemoryContextMemAllocated(tree->inner_slabs[i], true);
-		total += MemoryContextMemAllocated(tree->leaf_slabs[i], true);
-	}
 #endif
 
+	RT_LOCK_SHARED(tree);
+	total = tree->ctl->mem_used;
 	RT_UNLOCK(tree);
+
 	return total;
 }
+#endif
 
 /*
  * Verify the radix tree node.
@@ -2476,6 +2548,7 @@ RT_DUMP(RT_RADIX_TREE *tree)
 #undef RT_NEW_ROOT
 #undef RT_ALLOC_NODE
 #undef RT_INIT_NODE
+#undef RT_FANOUT_GET_NODE_SIZE
 #undef RT_FREE_NODE
 #undef RT_FREE_RECURSE
 #undef RT_EXTEND_UP
diff --git a/src/include/utils/dsa.h b/src/include/utils/dsa.h
index 2af215484f..3ce4ee300a 100644
--- a/src/include/utils/dsa.h
+++ b/src/include/utils/dsa.h
@@ -121,7 +121,6 @@ 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 5a169854d9..19d286d84b 100644
--- a/src/test/modules/test_radixtree/test_radixtree.c
+++ b/src/test/modules/test_radixtree/test_radixtree.c
@@ -114,6 +114,7 @@ static const test_spec test_specs[] = {
 #define RT_DECLARE
 #define RT_DEFINE
 #define RT_USE_DELETE
+#define RT_MEASURE_MEMORY_USAGE
 #define RT_VALUE_TYPE TestValueType
 /* #define RT_SHMEM */
 #include "lib/radixtree.h"
-- 
2.31.1