From a2ff55d53033a3bc00d587d1ad4354b54bc2a998 Mon Sep 17 00:00:00 2001
From: Ashutosh Bapat <ashutosh.bapat@enterprisedb.com>
Date: Mon, 19 Jun 2023 10:34:41 +0530
Subject: [PATCH 1/3] Partitionwise join memory consumption measurement

Add function to measure memory consumed by various function used during
partitionwise join. The memory consumed is reported at the end of
standard_planner() execution.
---
 src/backend/optimizer/path/allpaths.c |  8 +++
 src/backend/optimizer/path/joinrels.c | 31 +++++++++++-
 src/backend/optimizer/plan/planner.c  |  9 ++++
 src/backend/optimizer/util/pathnode.c |  5 ++
 src/backend/optimizer/util/relnode.c  | 18 ++++++-
 src/backend/utils/mmgr/mcxt.c         | 73 ++++++++++++++++++++++++++-
 src/include/utils/memutils.h          |  7 +++
 7 files changed, 146 insertions(+), 5 deletions(-)

diff --git a/src/backend/optimizer/path/allpaths.c b/src/backend/optimizer/path/allpaths.c
index 9bdc70c702..d82a2ac509 100644
--- a/src/backend/optimizer/path/allpaths.c
+++ b/src/backend/optimizer/path/allpaths.c
@@ -50,6 +50,8 @@
 #include "port/pg_bitutils.h"
 #include "rewrite/rewriteManip.h"
 #include "utils/lsyscache.h"
+#include "utils/memutils.h"
+#include "nodes/memnodes.h"
 
 
 /* Bitmask flags for pushdown_safety_info.unsafeFlags */
@@ -4302,6 +4304,10 @@ generate_partitionwise_join_paths(PlannerInfo *root, RelOptInfo *rel)
 	int			num_parts;
 	RelOptInfo **part_rels;
 
+	MemoryContextCounters mem_start;
+ 
+	MemoryContextFuncStatsStart(CurrentMemoryContext, &mem_start, __FUNCTION__);
+
 	/* Handle only join relations here. */
 	if (!IS_JOIN_REL(rel))
 		return;
@@ -4366,6 +4372,8 @@ generate_partitionwise_join_paths(PlannerInfo *root, RelOptInfo *rel)
 	/* Build additional paths for this rel from child-join paths. */
 	add_paths_to_append_rel(root, rel, live_children);
 	list_free(live_children);
+
+	MemoryContextFuncStatsEnd(CurrentMemoryContext, &mem_start, __FUNCTION__);
 }
 
 
diff --git a/src/backend/optimizer/path/joinrels.c b/src/backend/optimizer/path/joinrels.c
index 2feab2184f..c340a83c89 100644
--- a/src/backend/optimizer/path/joinrels.c
+++ b/src/backend/optimizer/path/joinrels.c
@@ -21,6 +21,7 @@
 #include "optimizer/paths.h"
 #include "partitioning/partbounds.h"
 #include "utils/memutils.h"
+#include "nodes/memnodes.h"
 
 
 static void make_rels_by_clause_joins(PlannerInfo *root,
@@ -895,6 +896,13 @@ populate_joinrel_with_paths(PlannerInfo *root, RelOptInfo *rel1,
 							RelOptInfo *rel2, RelOptInfo *joinrel,
 							SpecialJoinInfo *sjinfo, List *restrictlist)
 {
+	MemoryContextCounters mem_start;
+	char *label;
+
+	label = joinrel->reloptkind == RELOPT_OTHER_JOINREL ? "child_join_path_creation" : "parent_join_path_creation";
+
+	MemoryContextFuncStatsStart(CurrentMemoryContext, &mem_start, label);
+
 	/*
 	 * Consider paths using each rel as both outer and inner.  Depending on
 	 * the join type, a provably empty outer or inner rel might mean the join
@@ -1045,6 +1053,8 @@ populate_joinrel_with_paths(PlannerInfo *root, RelOptInfo *rel1,
 
 	/* Apply partitionwise join technique, if possible. */
 	try_partitionwise_join(root, rel1, rel2, joinrel, sjinfo, restrictlist);
+
+	MemoryContextFuncStatsEnd(CurrentMemoryContext, &mem_start, label);
 }
 
 
@@ -1487,6 +1497,10 @@ try_partitionwise_join(PlannerInfo *root, RelOptInfo *rel1, RelOptInfo *rel2,
 	ListCell   *lcr1 = NULL;
 	ListCell   *lcr2 = NULL;
 	int			cnt_parts;
+	MemoryContextCounters start_mem;
+
+	/* Start measuring memory */
+	MemoryContextFuncStatsStart(CurrentMemoryContext, &start_mem, __FUNCTION__);
 
 	/* Guard against stack overflow due to overly deep partition hierarchy. */
 	check_stack_depth();
@@ -1546,6 +1560,7 @@ try_partitionwise_join(PlannerInfo *root, RelOptInfo *rel1, RelOptInfo *rel2,
 		Relids		child_joinrelids;
 		AppendRelInfo **appinfos;
 		int			nappinfos;
+		MemoryContextCounters restrict_mem;
 
 		if (joinrel->partbounds_merged)
 		{
@@ -1648,6 +1663,8 @@ try_partitionwise_join(PlannerInfo *root, RelOptInfo *rel1, RelOptInfo *rel2,
 		/* Find the AppendRelInfo structures */
 		appinfos = find_appinfos_by_relids(root, child_joinrelids, &nappinfos);
 
+		MemoryContextFuncStatsStart(CurrentMemoryContext, &restrict_mem, "restrictlist translation");
+
 		/*
 		 * Construct restrictions applicable to the child join from those
 		 * applicable to the parent join.
@@ -1656,6 +1673,9 @@ try_partitionwise_join(PlannerInfo *root, RelOptInfo *rel1, RelOptInfo *rel2,
 			(List *) adjust_appendrel_attrs(root,
 											(Node *) parent_restrictlist,
 											nappinfos, appinfos);
+
+		MemoryContextFuncStatsEnd(CurrentMemoryContext, &restrict_mem, "restrictlist translation");
+
 		pfree(appinfos);
 
 		child_joinrel = joinrel->part_rels[cnt_parts];
@@ -1676,6 +1696,10 @@ try_partitionwise_join(PlannerInfo *root, RelOptInfo *rel1, RelOptInfo *rel2,
 									child_joinrel, child_sjinfo,
 									child_restrictlist);
 	}
+
+
+	/* Stop measuring memory and print the stats. */
+	MemoryContextFuncStatsEnd(CurrentMemoryContext, &start_mem, __FUNCTION__);
 }
 
 /*
@@ -1687,12 +1711,15 @@ static SpecialJoinInfo *
 build_child_join_sjinfo(PlannerInfo *root, SpecialJoinInfo *parent_sjinfo,
 						Relids left_relids, Relids right_relids)
 {
-	SpecialJoinInfo *sjinfo = makeNode(SpecialJoinInfo);
 	AppendRelInfo **left_appinfos;
 	int			left_nappinfos;
 	AppendRelInfo **right_appinfos;
 	int			right_nappinfos;
+	MemoryContextCounters mem_start;
+	SpecialJoinInfo *sjinfo;
 
+	MemoryContextFuncStatsStart(CurrentMemoryContext, &mem_start, __FUNCTION__);
+	sjinfo = makeNode(SpecialJoinInfo);
 	memcpy(sjinfo, parent_sjinfo, sizeof(SpecialJoinInfo));
 	left_appinfos = find_appinfos_by_relids(root, left_relids,
 											&left_nappinfos);
@@ -1718,6 +1745,8 @@ build_child_join_sjinfo(PlannerInfo *root, SpecialJoinInfo *parent_sjinfo,
 	pfree(left_appinfos);
 	pfree(right_appinfos);
 
+	MemoryContextFuncStatsEnd(CurrentMemoryContext, &mem_start, __FUNCTION__);
+
 	return sjinfo;
 }
 
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 1e4dd27dba..0d742c2587 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -67,6 +67,8 @@
 #include "utils/rel.h"
 #include "utils/selfuncs.h"
 #include "utils/syscache.h"
+#include "utils/memutils.h"
+#include "nodes/memnodes.h"
 
 /* GUC parameters */
 double		cursor_tuple_fraction = DEFAULT_CURSOR_TUPLE_FRACTION;
@@ -295,6 +297,9 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 	Plan	   *top_plan;
 	ListCell   *lp,
 			   *lr;
+	MemoryContextCounters mem_start;
+
+	MemoryContextFuncStatsStart(CurrentMemoryContext, &mem_start, __FUNCTION__);
 
 	/*
 	 * Set up global state for this planner invocation.  This data is needed
@@ -565,6 +570,10 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 	if (glob->partition_directory != NULL)
 		DestroyPartitionDirectory(glob->partition_directory);
 
+	MemoryContextFuncStatsEnd(CurrentMemoryContext, &mem_start, __FUNCTION__);
+
+	MemoryContextFuncStatsReport();
+
 	return result;
 }
 
diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c
index 5f5596841c..af5eea2e31 100644
--- a/src/backend/optimizer/util/pathnode.c
+++ b/src/backend/optimizer/util/pathnode.c
@@ -4063,6 +4063,9 @@ do { \
 	ParamPathInfo *new_ppi;
 	ParamPathInfo *old_ppi;
 	Relids		required_outer;
+	MemoryContextCounters	mem_start;
+
+	MemoryContextFuncStatsStart(CurrentMemoryContext, &mem_start, __FUNCTION__);
 
 	/*
 	 * If the path is not parameterized by parent of the given relation, it
@@ -4311,6 +4314,8 @@ do { \
 		ADJUST_CHILD_ATTRS(new_path->pathtarget->exprs);
 	}
 
+	MemoryContextFuncStatsEnd(CurrentMemoryContext, &mem_start, __FUNCTION__);
+
 	return new_path;
 }
 
diff --git a/src/backend/optimizer/util/relnode.c b/src/backend/optimizer/util/relnode.c
index 15e3910b79..e305bc894a 100644
--- a/src/backend/optimizer/util/relnode.c
+++ b/src/backend/optimizer/util/relnode.c
@@ -32,6 +32,8 @@
 #include "parser/parse_relation.h"
 #include "utils/hsearch.h"
 #include "utils/lsyscache.h"
+#include "nodes/memnodes.h"
+#include "utils/memutils.h"
 
 
 typedef struct JoinHashEntry
@@ -861,16 +863,22 @@ build_child_join_rel(PlannerInfo *root, RelOptInfo *outer_rel,
 					 RelOptInfo *inner_rel, RelOptInfo *parent_joinrel,
 					 List *restrictlist, SpecialJoinInfo *sjinfo)
 {
-	RelOptInfo *joinrel = makeNode(RelOptInfo);
+	RelOptInfo *joinrel;
 	AppendRelInfo **appinfos;
 	int			nappinfos;
 
+	MemoryContextCounters mem_start;
+	MemoryContextCounters tlist_mem;
+	MemoryContextCounters jlist_mem;
+
+	MemoryContextFuncStatsStart(CurrentMemoryContext, &mem_start, __FUNCTION__);
+
 	/* Only joins between "other" relations land here. */
 	Assert(IS_OTHER_REL(outer_rel) && IS_OTHER_REL(inner_rel));
 
 	/* The parent joinrel should have consider_partitionwise_join set. */
 	Assert(parent_joinrel->consider_partitionwise_join);
-
+	joinrel = makeNode(RelOptInfo);
 	joinrel->reloptkind = RELOPT_OTHER_JOINREL;
 	joinrel->relids = bms_union(outer_rel->relids, inner_rel->relids);
 	joinrel->relids = add_outer_joins_to_relids(root, joinrel->relids, sjinfo,
@@ -939,14 +947,18 @@ build_child_join_rel(PlannerInfo *root, RelOptInfo *outer_rel,
 	appinfos = find_appinfos_by_relids(root, joinrel->relids, &nappinfos);
 
 	/* Set up reltarget struct */
+	MemoryContextFuncStatsStart(CurrentMemoryContext, &tlist_mem, "targetlist");
 	build_child_join_reltarget(root, parent_joinrel, joinrel,
 							   nappinfos, appinfos);
+	MemoryContextFuncStatsEnd(CurrentMemoryContext, &tlist_mem, "targetlist");
 
 	/* Construct joininfo list. */
+	MemoryContextFuncStatsStart(CurrentMemoryContext, &jlist_mem, "joininfo");
 	joinrel->joininfo = (List *) adjust_appendrel_attrs(root,
 														(Node *) parent_joinrel->joininfo,
 														nappinfos,
 														appinfos);
+	MemoryContextFuncStatsEnd(CurrentMemoryContext, &jlist_mem, "joininfo");
 
 	/*
 	 * Lateral relids referred in child join will be same as that referred in
@@ -991,6 +1003,8 @@ build_child_join_rel(PlannerInfo *root, RelOptInfo *outer_rel,
 
 	pfree(appinfos);
 
+	MemoryContextFuncStatsEnd(CurrentMemoryContext, &mem_start, __FUNCTION__);
+
 	return joinrel;
 }
 
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 9fc83f11f6..54320536e1 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -150,7 +150,7 @@ MemoryContext CurTransactionContext = NULL;
 MemoryContext PortalContext = NULL;
 
 static void MemoryContextCallResetCallbacks(MemoryContext context);
-static void MemoryContextStatsInternal(MemoryContext context, int level,
+void MemoryContextStatsInternal(MemoryContext context, int level,
 									   bool print, int max_children,
 									   MemoryContextCounters *totals,
 									   bool print_to_stderr);
@@ -754,7 +754,7 @@ MemoryContextStatsDetail(MemoryContext context, int max_children,
  * Print this context if print is true, but in any case accumulate counts into
  * *totals (if given).
  */
-static void
+void
 MemoryContextStatsInternal(MemoryContext context, int level,
 						   bool print, int max_children,
 						   MemoryContextCounters *totals,
@@ -838,6 +838,75 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 	}
 }
 
+HTAB *mem_cons_ht = NULL;
+typedef struct
+{
+	char	label[NAMEDATALEN];
+	Size	consumed_mem;
+} MemConsEntry;
+
+void
+MemoryContextFuncStatsStart(MemoryContext context,
+							MemoryContextCounters *start_counts,
+							const char *label)
+{
+	if (!mem_cons_ht)
+	{
+		HASHCTL ctl;
+
+		ctl.keysize = NAMEDATALEN;
+		ctl.entrysize = sizeof(MemConsEntry);
+
+		mem_cons_ht = hash_create("planner memory consumption", 10, &ctl, HASH_ELEM | HASH_STRINGS);
+	}
+
+	memset(start_counts, 0, sizeof(*start_counts));
+	MemoryContextStatsInternal(context, 0, false, 100, start_counts, false);
+}
+
+void
+MemoryContextFuncStatsEnd(MemoryContext context,
+						  MemoryContextCounters *start_counts,
+						  const char *label)
+{
+	MemoryContextCounters end_counts;
+	Size	start_used_space = start_counts->totalspace - start_counts->freespace;
+	Size	end_used_space;
+	bool	found;
+	MemConsEntry   *entry;
+
+	memset(&end_counts, 0, sizeof(end_counts));
+	MemoryContextStatsInternal(context, 0, false, 100, &end_counts, false);
+	end_used_space = end_counts.totalspace - end_counts.freespace;
+
+/*
+	elog(NOTICE, "%s,%s,%zu,%zu,%ld", label, context->name,
+		 start_used_space, end_used_space, end_used_space - start_used_space);
+*/
+	entry = hash_search(mem_cons_ht, label, HASH_ENTER, &found);
+	if (!found)
+	{
+		strncpy(entry->label, label, NAMEDATALEN);
+		entry->consumed_mem = 0;
+	}
+	Assert(strncmp(entry->label, label, NAMEDATALEN) == 0);
+	entry->consumed_mem = entry->consumed_mem + (end_used_space - start_used_space);
+}
+
+void
+MemoryContextFuncStatsReport(void)
+{
+	HASH_SEQ_STATUS status;
+	MemConsEntry   *entry;
+
+	/* Walk through the hash table printing memory consumed by the label. */
+	hash_seq_init(&status, mem_cons_ht);
+	while ((entry = (MemConsEntry *) hash_seq_search(&status)) != NULL)
+		elog(NOTICE, "%s, %ld", entry->label, entry->consumed_mem);
+	hash_destroy(mem_cons_ht);
+	mem_cons_ht = NULL;
+}
+
 /*
  * MemoryContextStatsPrint
  *		Print callback used by MemoryContextStatsInternal
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 21640d62a6..0027a27981 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -92,6 +92,13 @@ extern void MemoryContextStatsDetail(MemoryContext context, int max_children,
 									 bool print_to_stderr);
 extern void MemoryContextAllowInCriticalSection(MemoryContext context,
 												bool allow);
+extern void MemoryContextFuncStatsStart(MemoryContext context,
+										MemoryContextCounters *start_counts,
+										const char *func_label);
+extern void MemoryContextFuncStatsEnd(MemoryContext context,
+									  MemoryContextCounters *start_counts,
+									  const char *func_label);
+extern void MemoryContextFuncStatsReport(void);
 
 #ifdef MEMORY_CONTEXT_CHECKING
 extern void MemoryContextCheck(MemoryContext context);
-- 
2.25.1

