From 1c4f33ce05fdfdbeab6e7f35d064aaa97373536b Mon Sep 17 00:00:00 2001
From: Ashutosh Bapat <ashutosh.bapat@enterprisedb.com>
Date: Wed, 12 Jul 2023 14:34:14 +0530
Subject: [PATCH] EXPLAIN reports memory consumed for planning a query

The memory consumed is reported as "Memory" property under "Planning"
section in EXPLAIN output when option MEMORY is specified.

A separate memory context is allocated for measuring memory consumption
to eliminate any effect of previous activity on the calculation.  The
memory context statistics is noted before and after calling
pg_plan_query() from ExplainOneQuery(). The difference in the two
statistics is used to calculate the memory consumption.

We report two values "used" and "allocated".  The difference between
memory statistics counters totalspace and freespace gives the memory
that remains palloc'ed at the end of planning. It is reported as memory
"used". This does not account for chunks of memory that were palloc'ed
but subsequently pfree'ed. But such memory remains allocated in the
memory context is given by the totalspace counter. The value of this
counter is reported as memory "allocated".

Author: Ashutosh Bapat
Reviewed-by: David Rowley, Andrey Lepikhov, Jian He, Andy Fan
Reviewed-by: Alvaro Herrera
Discussion: https://www.postgresql.org/message-id/CAExHW5sZA=5LJ_ZPpRO-w09ck8z9p7eaYAqq3Ks9GDfhrxeWBw@mail.gmail.com
---
 contrib/auto_explain/auto_explain.c   |   2 +
 doc/src/sgml/ref/explain.sgml         |  14 +++
 src/backend/commands/explain.c        | 124 ++++++++++++++++++++++----
 src/backend/commands/prepare.c        |  29 +++++-
 src/backend/utils/mmgr/mcxt.c         |  13 +++
 src/include/commands/explain.h        |   4 +-
 src/include/utils/memutils.h          |   2 +
 src/test/regress/expected/explain.out |  68 ++++++++++++++
 src/test/regress/sql/explain.sql      |   6 ++
 9 files changed, 245 insertions(+), 17 deletions(-)

diff --git a/contrib/auto_explain/auto_explain.c b/contrib/auto_explain/auto_explain.c
index c7aacd7812..dcb5d6c9e5 100644
--- a/contrib/auto_explain/auto_explain.c
+++ b/contrib/auto_explain/auto_explain.c
@@ -396,6 +396,8 @@ explain_ExecutorEnd(QueryDesc *queryDesc)
 			es->wal = (es->analyze && auto_explain_log_wal);
 			es->timing = (es->analyze && auto_explain_log_timing);
 			es->summary = es->analyze;
+			/* No support for MEMORY option */
+			/* es->memory = false; */
 			es->format = auto_explain_log_format;
 			es->settings = auto_explain_log_settings;
 
diff --git a/doc/src/sgml/ref/explain.sgml b/doc/src/sgml/ref/explain.sgml
index 5ba6486da1..26809d68d5 100644
--- a/doc/src/sgml/ref/explain.sgml
+++ b/doc/src/sgml/ref/explain.sgml
@@ -44,6 +44,7 @@ EXPLAIN [ ( <replaceable class="parameter">option</replaceable> [, ...] ) ] <rep
     WAL [ <replaceable class="parameter">boolean</replaceable> ]
     TIMING [ <replaceable class="parameter">boolean</replaceable> ]
     SUMMARY [ <replaceable class="parameter">boolean</replaceable> ]
+    MEMORY [ <replaceable class="parameter">boolean</replaceable> ]
     FORMAT { TEXT | XML | JSON | YAML }
 </synopsis>
  </refsynopsisdiv>
@@ -250,6 +251,19 @@ ROLLBACK;
     </listitem>
    </varlistentry>
 
+   <varlistentry>
+    <term><literal>MEMORY</literal></term>
+    <listitem>
+     <para>
+      Include information on memory consumption by the query planning phase.
+      Specifically, include the precise amount of storage used by planner
+      in-memory structures, as well as total memory considering allocation
+      overhead.
+      The parameter defaults to <literal>FALSE</literal>.
+     </para>
+    </listitem>
+   </varlistentry>
+
    <varlistentry>
     <term><literal>FORMAT</literal></term>
     <listitem>
diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index 843472e6dd..e912f7dae0 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -119,9 +119,12 @@ static void show_instrumentation_count(const char *qlabel, int which,
 static void show_foreignscan_info(ForeignScanState *fsstate, ExplainState *es);
 static void show_eval_params(Bitmapset *bms_params, ExplainState *es);
 static const char *explain_get_index_name(Oid indexId);
-static void show_buffer_usage(ExplainState *es, const BufferUsage *usage,
+static bool show_buffer_usage(ExplainState *es, const BufferUsage *usage,
 							  bool planning);
 static void show_wal_usage(ExplainState *es, const WalUsage *usage);
+static bool show_planner_memory(ExplainState *es,
+								const MemoryContextCounters *mem_counts,
+								bool show_planning);
 static void ExplainIndexScanDetails(Oid indexid, ScanDirection indexorderdir,
 									ExplainState *es);
 static void ExplainScanTarget(Scan *plan, ExplainState *es);
@@ -202,6 +205,8 @@ ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
 			summary_set = true;
 			es->summary = defGetBoolean(opt);
 		}
+		else if (strcmp(opt->defname, "memory") == 0)
+			es->memory = defGetBoolean(opt);
 		else if (strcmp(opt->defname, "format") == 0)
 		{
 			char	   *p = defGetString(opt);
@@ -397,6 +402,24 @@ ExplainOneQuery(Query *query, int cursorOptions,
 					planduration;
 		BufferUsage bufusage_start,
 					bufusage;
+		MemoryContextCounters mem_counts;
+		MemoryContext planner_ctx = NULL;
+		MemoryContext saved_ctx = NULL;
+
+		if (es->memory)
+		{
+			/*
+			 * Create a new memory context to measure planner's memory
+			 * consumption accurately. We should use the same type of memory
+			 * context as the planner would use. That's usually AllocSet but
+			 * ensure that.
+			 */
+			Assert(IsA(CurrentMemoryContext, AllocSetContext));
+			planner_ctx = AllocSetContextCreate(CurrentMemoryContext,
+												"explain analyze planner context",
+												ALLOCSET_DEFAULT_SIZES);
+			saved_ctx = MemoryContextSwitchTo(planner_ctx);
+		}
 
 		if (es->buffers)
 			bufusage_start = pgBufferUsage;
@@ -408,6 +431,12 @@ ExplainOneQuery(Query *query, int cursorOptions,
 		INSTR_TIME_SET_CURRENT(planduration);
 		INSTR_TIME_SUBTRACT(planduration, planstart);
 
+		if (es->memory)
+		{
+			MemoryContextSwitchTo(saved_ctx);
+			MemoryContextMemConsumed(planner_ctx, &mem_counts);
+		}
+
 		/* calc differences of buffer counters. */
 		if (es->buffers)
 		{
@@ -417,7 +446,8 @@ ExplainOneQuery(Query *query, int cursorOptions,
 
 		/* run it (if needed) and produce output */
 		ExplainOnePlan(plan, into, es, queryString, params, queryEnv,
-					   &planduration, (es->buffers ? &bufusage : NULL));
+					   &planduration, (es->buffers ? &bufusage : NULL),
+					   (es->memory ? &mem_counts : NULL));
 	}
 }
 
@@ -527,7 +557,8 @@ void
 ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 			   const char *queryString, ParamListInfo params,
 			   QueryEnvironment *queryEnv, const instr_time *planduration,
-			   const BufferUsage *bufusage)
+			   const BufferUsage *bufusage,
+			   const MemoryContextCounters *mem_counts)
 {
 	DestReceiver *dest;
 	QueryDesc  *queryDesc;
@@ -615,11 +646,22 @@ ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 	/* Create textual dump of plan tree */
 	ExplainPrintPlan(es, queryDesc);
 
-	/* Show buffer usage in planning */
-	if (bufusage)
+	/* Show buffer and/or memory usage in planning */
+	if (bufusage || mem_counts)
 	{
+		bool		showed_planning = false;
+
 		ExplainOpenGroup("Planning", "Planning", true, es);
-		show_buffer_usage(es, bufusage, true);
+
+		if (bufusage)
+			showed_planning = show_buffer_usage(es, bufusage, true);
+
+		if (mem_counts)
+			showed_planning |= show_planner_memory(es, mem_counts, !showed_planning);
+
+		if (showed_planning)
+			es->indent--;
+
 		ExplainCloseGroup("Planning", "Planning", true, es);
 	}
 
@@ -3546,10 +3588,20 @@ explain_get_index_name(Oid indexId)
 
 /*
  * Show buffer usage details.
+ *
+ * When reporting in TEXT format this function opens "planning" section if only there are counts to
+ * be reported for planning. The function returns true if the section was
+ * opened, false otherwise. The function does not close the section. The caller is expected to use the return value to
+ * close the section if required.
+ *
+ * When reporting in non-TEXT format the caller is expected to open and close the section,
+ * if required. Hence the function always returns false for non-TEXT formats.
  */
-static void
+static bool
 show_buffer_usage(ExplainState *es, const BufferUsage *usage, bool planning)
 {
+	bool		show_planning = false;
+
 	if (es->format == EXPLAIN_FORMAT_TEXT)
 	{
 		bool		has_shared = (usage->shared_blks_hit > 0 ||
@@ -3568,12 +3620,10 @@ show_buffer_usage(ExplainState *es, const BufferUsage *usage, bool planning)
 										!INSTR_TIME_IS_ZERO(usage->local_blk_write_time));
 		bool		has_temp_timing = (!INSTR_TIME_IS_ZERO(usage->temp_blk_read_time) ||
 									   !INSTR_TIME_IS_ZERO(usage->temp_blk_write_time));
-		bool		show_planning = (planning && (has_shared ||
-												  has_local || has_temp ||
-												  has_shared_timing ||
-												  has_local_timing ||
-												  has_temp_timing));
 
+		show_planning = (planning && (has_shared || has_local || has_temp ||
+									  has_shared_timing || has_local_timing ||
+									  has_temp_timing));
 		if (show_planning)
 		{
 			ExplainIndentText(es);
@@ -3678,9 +3728,6 @@ show_buffer_usage(ExplainState *es, const BufferUsage *usage, bool planning)
 			}
 			appendStringInfoChar(es->str, '\n');
 		}
-
-		if (show_planning)
-			es->indent--;
 	}
 	else
 	{
@@ -3726,6 +3773,8 @@ show_buffer_usage(ExplainState *es, const BufferUsage *usage, bool planning)
 								 3, es);
 		}
 	}
+
+	return show_planning;
 }
 
 /*
@@ -3766,6 +3815,51 @@ show_wal_usage(ExplainState *es, const WalUsage *usage)
 	}
 }
 
+/*
+ * Show planner's memory usage details.
+ *
+ * When reporting in TEXT format this function opens the "planning" section if only it was not opened already. The function returns true if it opened the section,
+ * false otherwise. The caller is expected to use the return value to
+ * close the section if required.
+ *
+ * When reporting in non-TEXT format the caller is expected to open and close the section,
+ * if required. Hence the function always returns false for non-TEXT formats.
+ */
+static bool
+show_planner_memory(ExplainState *es,
+					const MemoryContextCounters *mem_counts, bool show_planning)
+{
+	if (es->format == EXPLAIN_FORMAT_TEXT)
+	{
+		if (show_planning)
+		{
+			ExplainIndentText(es);
+			appendStringInfoString(es->str, "Planning:\n");
+			es->indent++;
+		}
+
+		ExplainIndentText(es);
+		appendStringInfo(es->str,
+						 "Memory: used=%lld bytes, allocated=%lld bytes",
+						 (long long) (mem_counts->totalspace - mem_counts->freespace),
+						 (long long) mem_counts->totalspace);
+		appendStringInfoChar(es->str, '\n');
+	}
+	else
+	{
+		show_planning = false;
+
+		ExplainPropertyInteger("Memory Used", "bytes",
+							   mem_counts->totalspace - mem_counts->freespace,
+							   es);
+		ExplainPropertyInteger("Memory Allocated", "bytes",
+							   mem_counts->totalspace, es);
+	}
+
+	return show_planning;
+}
+
+
 /*
  * Add some additional details about an IndexScan or IndexOnlyScan
  */
diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index 703976f633..ab62a22bb6 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -583,6 +583,26 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	instr_time	planduration;
 	BufferUsage bufusage_start,
 				bufusage;
+	MemoryContextCounters mem_counts;
+	MemoryContext planner_ctx = NULL;
+	MemoryContext saved_ctx = NULL;
+
+	if (es->memory)
+	{
+		/*
+		 * In order to measure planner's memory consumption accurately, create
+		 * a separate AllocSet memory context.
+		 *
+		 * XXX if planner is called under some other memory context type, this
+		 * code overrides that.  Do we need support to create context of
+		 * programmatically determined type?
+		 */
+		Assert(IsA(CurrentMemoryContext, AllocSetContext));
+		planner_ctx = AllocSetContextCreate(CurrentMemoryContext,
+											"explain analyze planner context",
+											ALLOCSET_DEFAULT_SIZES);
+		saved_ctx = MemoryContextSwitchTo(planner_ctx);
+	}
 
 	if (es->buffers)
 		bufusage_start = pgBufferUsage;
@@ -624,6 +644,12 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	INSTR_TIME_SET_CURRENT(planduration);
 	INSTR_TIME_SUBTRACT(planduration, planstart);
 
+	if (es->memory)
+	{
+		MemoryContextSwitchTo(saved_ctx);
+		MemoryContextMemConsumed(planner_ctx, &mem_counts);
+	}
+
 	/* calc differences of buffer counters. */
 	if (es->buffers)
 	{
@@ -640,7 +666,8 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 
 		if (pstmt->commandType != CMD_UTILITY)
 			ExplainOnePlan(pstmt, into, es, query_string, paramLI, queryEnv,
-						   &planduration, (es->buffers ? &bufusage : NULL));
+						   &planduration, (es->buffers ? &bufusage : NULL),
+						   (es->memory ? &mem_counts : NULL));
 		else
 			ExplainOneUtility(pstmt->utilityStmt, into, es, query_string,
 							  paramLI, queryEnv);
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 1336944084..ad7409a02c 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -687,6 +687,19 @@ MemoryContextMemAllocated(MemoryContext context, bool recurse)
 	return total;
 }
 
+/*
+ * Return the memory consumption statistics about the given context and its
+ * children.
+ */
+void
+MemoryContextMemConsumed(MemoryContext context,
+						 MemoryContextCounters *consumed)
+{
+	memset(consumed, 0, sizeof(*consumed));
+
+	MemoryContextStatsInternal(context, 0, false, 0, consumed, false);
+}
+
 /*
  * MemoryContextStats
  *		Print statistics about the named context and all its descendants.
diff --git a/src/include/commands/explain.h b/src/include/commands/explain.h
index 1b44d483d6..c03006250f 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -45,6 +45,7 @@ typedef struct ExplainState
 	bool		wal;			/* print WAL usage */
 	bool		timing;			/* print detailed node timing */
 	bool		summary;		/* print total planning and execution timing */
+	bool		memory;			/* print planner's memory usage information */
 	bool		settings;		/* print modified settings */
 	bool		generic;		/* generate a generic plan */
 	ExplainFormat format;		/* output format */
@@ -92,7 +93,8 @@ extern void ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into,
 						   ExplainState *es, const char *queryString,
 						   ParamListInfo params, QueryEnvironment *queryEnv,
 						   const instr_time *planduration,
-						   const BufferUsage *bufusage);
+						   const BufferUsage *bufusage,
+						   const MemoryContextCounters *mem_counts);
 
 extern void ExplainPrintPlan(ExplainState *es, QueryDesc *queryDesc);
 extern void ExplainPrintTriggers(ExplainState *es, QueryDesc *queryDesc);
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index ca7858d6b6..7fd41d20ca 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -84,6 +84,8 @@ extern Size GetMemoryChunkSpace(void *pointer);
 extern MemoryContext MemoryContextGetParent(MemoryContext context);
 extern bool MemoryContextIsEmpty(MemoryContext context);
 extern Size MemoryContextMemAllocated(MemoryContext context, bool recurse);
+extern void MemoryContextMemConsumed(MemoryContext context,
+									 MemoryContextCounters *consumed);
 extern void MemoryContextStats(MemoryContext context);
 extern void MemoryContextStatsDetail(MemoryContext context, int max_children,
 									 bool print_to_stderr);
diff --git a/src/test/regress/expected/explain.out b/src/test/regress/expected/explain.out
index 809655e16e..55694505a7 100644
--- a/src/test/regress/expected/explain.out
+++ b/src/test/regress/expected/explain.out
@@ -326,6 +326,74 @@ select explain_filter('explain (generic_plan) select unique1 from tenk1 where th
 select explain_filter('explain (analyze, generic_plan) select unique1 from tenk1 where thousand = $1');
 ERROR:  EXPLAIN options ANALYZE and GENERIC_PLAN cannot be used together
 CONTEXT:  PL/pgSQL function explain_filter(text) line 5 at FOR over EXECUTE statement
+-- MEMORY option
+select explain_filter('explain (memory) select * from int8_tbl i8');
+                     explain_filter                      
+---------------------------------------------------------
+ Seq Scan on int8_tbl i8  (cost=N.N..N.N rows=N width=N)
+   Memory: used=N bytes, allocated=N bytes
+(2 rows)
+
+select explain_filter('explain (memory, analyze) select * from int8_tbl i8');
+                                        explain_filter                                         
+-----------------------------------------------------------------------------------------------
+ Seq Scan on int8_tbl i8  (cost=N.N..N.N rows=N width=N) (actual time=N.N..N.N rows=N loops=N)
+   Memory: used=N bytes, allocated=N bytes
+ Planning Time: N.N ms
+ Execution Time: N.N ms
+(4 rows)
+
+select explain_filter('explain (memory, summary, format yaml) select * from int8_tbl i8');
+        explain_filter         
+-------------------------------
+ - Plan:                      +
+     Node Type: "Seq Scan"    +
+     Parallel Aware: false    +
+     Async Capable: false     +
+     Relation Name: "int8_tbl"+
+     Alias: "i8"              +
+     Startup Cost: N.N        +
+     Total Cost: N.N          +
+     Plan Rows: N             +
+     Plan Width: N            +
+   Planning:                  +
+     Memory Used: N           +
+     Memory Allocated: N      +
+   Planning Time: N.N
+(1 row)
+
+select explain_filter('explain (memory, analyze, format json) select * from int8_tbl i8');
+           explain_filter           
+------------------------------------
+ [                                 +
+   {                               +
+     "Plan": {                     +
+       "Node Type": "Seq Scan",    +
+       "Parallel Aware": false,    +
+       "Async Capable": false,     +
+       "Relation Name": "int8_tbl",+
+       "Alias": "i8",              +
+       "Startup Cost": N.N,        +
+       "Total Cost": N.N,          +
+       "Plan Rows": N,             +
+       "Plan Width": N,            +
+       "Actual Startup Time": N.N, +
+       "Actual Total Time": N.N,   +
+       "Actual Rows": N,           +
+       "Actual Loops": N           +
+     },                            +
+     "Planning": {                 +
+       "Memory Used": N,           +
+       "Memory Allocated": N       +
+     },                            +
+     "Planning Time": N.N,         +
+     "Triggers": [                 +
+     ],                            +
+     "Execution Time": N.N         +
+   }                               +
+ ]
+(1 row)
+
 -- Test EXPLAIN (GENERIC_PLAN) with partition pruning
 -- partitions should be pruned at plan time, based on constants,
 -- but there should be no pruning based on parameter placeholders
diff --git a/src/test/regress/sql/explain.sql b/src/test/regress/sql/explain.sql
index b6b7beab27..f719e92f84 100644
--- a/src/test/regress/sql/explain.sql
+++ b/src/test/regress/sql/explain.sql
@@ -94,6 +94,12 @@ select explain_filter('explain (generic_plan) select unique1 from tenk1 where th
 -- should fail
 select explain_filter('explain (analyze, generic_plan) select unique1 from tenk1 where thousand = $1');
 
+-- MEMORY option
+select explain_filter('explain (memory) select * from int8_tbl i8');
+select explain_filter('explain (memory, analyze) select * from int8_tbl i8');
+select explain_filter('explain (memory, summary, format yaml) select * from int8_tbl i8');
+select explain_filter('explain (memory, analyze, format json) select * from int8_tbl i8');
+
 -- Test EXPLAIN (GENERIC_PLAN) with partition pruning
 -- partitions should be pruned at plan time, based on constants,
 -- but there should be no pruning based on parameter placeholders
-- 
2.25.1

