0001-Report-memory-used-for-planning-a-query-in--20231031.patch

From 073eb3456c704a9b13bc53942b92ee64fbb3be5e 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 1/4] Report memory used for planning a query in EXPLAIN

Helper patch which is squashed version of patches submitted in https://www.postgresql.org/message-id/CAExHW5sCJX7696sF-OnugAiaXS=Ag95=-m1cSrjcmyYj8Pduuw@mail.gmail.com. Not to be committed with this patch set.

Ashutosh Bapat
---
 src/backend/commands/explain.c        | 79 ++++++++++++++++++++++++++-
 src/backend/commands/prepare.c        | 24 +++++++-
 src/backend/utils/mmgr/mcxt.c         | 12 ++++
 src/include/commands/explain.h        | 12 +++-
 src/include/utils/memutils.h          |  2 +
 src/test/regress/expected/explain.out | 26 +++++++--
 src/tools/pgindent/typedefs.list      |  1 +
 7 files changed, 148 insertions(+), 8 deletions(-)

diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index 8570b14f62..e7df8e4455 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -122,6 +122,8 @@ static const char *explain_get_index_name(Oid indexId);
 static void show_buffer_usage(ExplainState *es, const BufferUsage *usage,
 							  bool planning);
 static void show_wal_usage(ExplainState *es, const WalUsage *usage);
+static void show_planning_memory(ExplainState *es,
+								 const MemUsage *usage);
 static void ExplainIndexScanDetails(Oid indexid, ScanDirection indexorderdir,
 									ExplainState *es);
 static void ExplainScanTarget(Scan *plan, ExplainState *es);
@@ -397,9 +399,27 @@ ExplainOneQuery(Query *query, int cursorOptions,
 					planduration;
 		BufferUsage bufusage_start,
 					bufusage;
+		MemoryContextCounters mem_counts_start;
+		MemoryContextCounters mem_counts_end;
+		MemUsage	mem_usage;
+		MemoryContext planner_ctx;
+		MemoryContext saved_ctx;
+
+		/*
+		 * Create a new memory context to accurately measure memory malloc'ed
+		 * by the planner. For further accuracy 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);
 
 		if (es->buffers)
 			bufusage_start = pgBufferUsage;
+		MemoryContextMemConsumed(planner_ctx, &mem_counts_start);
+		saved_ctx = MemoryContextSwitchTo(planner_ctx);
 		INSTR_TIME_SET_CURRENT(planstart);
 
 		/* plan the query */
@@ -407,6 +427,9 @@ ExplainOneQuery(Query *query, int cursorOptions,
 
 		INSTR_TIME_SET_CURRENT(planduration);
 		INSTR_TIME_SUBTRACT(planduration, planstart);
+		MemoryContextSwitchTo(saved_ctx);
+		MemoryContextMemConsumed(planner_ctx, &mem_counts_end);
+		calc_mem_usage(&mem_usage, &mem_counts_end, &mem_counts_start);
 
 		/* calc differences of buffer counters. */
 		if (es->buffers)
@@ -417,7 +440,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),
+					   &mem_usage);
 	}
 }
 
@@ -527,7 +551,7 @@ 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 MemUsage *mem_usage)
 {
 	DestReceiver *dest;
 	QueryDesc  *queryDesc;
@@ -630,6 +654,13 @@ ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 		ExplainPropertyFloat("Planning Time", "ms", 1000.0 * plantime, 3, es);
 	}
 
+	if (es->summary && mem_usage)
+	{
+		ExplainOpenGroup("Planning Memory", "Planning Memory", true, es);
+		show_planning_memory(es, mem_usage);
+		ExplainCloseGroup("Planning Memory", "Planning Memory", true, es);
+	}
+
 	/* Print info about runtime of triggers */
 	if (es->analyze)
 		ExplainPrintTriggers(es, queryDesc);
@@ -3736,6 +3767,50 @@ show_wal_usage(ExplainState *es, const WalUsage *usage)
 	}
 }
 
+/*
+ * Show planner's memory usage details.
+ */
+static void
+show_planning_memory(ExplainState *es, const MemUsage *usage)
+{
+	if (es->format == EXPLAIN_FORMAT_TEXT)
+	{
+		appendStringInfo(es->str,
+						 "Planning Memory: used=%zu bytes allocated=%zu bytes",
+						 usage->mem_used, usage->mem_allocated);
+		appendStringInfoChar(es->str, '\n');
+	}
+	else
+	{
+		ExplainPropertyInteger("Used", "bytes", usage->mem_used, es);
+		ExplainPropertyInteger("Allocated", "bytes", usage->mem_allocated, es);
+	}
+}
+
+/*
+ * Compute memory usage from the start and end memory counts.
+ */
+void
+calc_mem_usage(MemUsage *mem_usage, MemoryContextCounters *mem_counts_end,
+			   MemoryContextCounters *mem_counts_start)
+{
+	Size		mem_used_start;
+	Size		mem_used_end;
+
+	mem_used_start = mem_counts_start->totalspace - mem_counts_start->freespace;
+	mem_used_end = mem_counts_end->totalspace - mem_counts_end->freespace;
+
+	mem_usage->mem_used = mem_used_end - mem_used_start;
+
+	/*
+	 * The net memory used is from total memory allocated and not necessarily
+	 * the net memory allocated between the two given samples. Hence do not
+	 * compute the difference between allocated memory reported in the two
+	 * given samples.
+	 */
+	mem_usage->mem_allocated = mem_counts_end->totalspace;
+}
+
 /*
  * Add some additional details about an IndexScan or IndexOnlyScan
  */
diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index 18f70319fc..eb39823d7a 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -583,9 +583,27 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	instr_time	planduration;
 	BufferUsage bufusage_start,
 				bufusage;
+	MemoryContextCounters mem_counts_start;
+	MemoryContextCounters mem_counts_end;
+	MemUsage	mem_usage;
+	MemoryContext planner_ctx;
+	MemoryContext saved_ctx;
+
+	/*
+	 * Create a new memory context to accurately measure memory malloc'ed by
+	 * the planner. For further accuracy 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);
 
 	if (es->buffers)
 		bufusage_start = pgBufferUsage;
+	MemoryContextMemConsumed(planner_ctx, &mem_counts_start);
+	saved_ctx = MemoryContextSwitchTo(planner_ctx);
 	INSTR_TIME_SET_CURRENT(planstart);
 
 	/* Look it up in the hash table */
@@ -623,6 +641,9 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 
 	INSTR_TIME_SET_CURRENT(planduration);
 	INSTR_TIME_SUBTRACT(planduration, planstart);
+	MemoryContextSwitchTo(saved_ctx);
+	MemoryContextMemConsumed(planner_ctx, &mem_counts_end);
+	calc_mem_usage(&mem_usage, &mem_counts_end, &mem_counts_start);
 
 	/* calc differences of buffer counters. */
 	if (es->buffers)
@@ -640,7 +661,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),
+						   &mem_usage);
 		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 9fc83f11f6..94159a6799 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -747,6 +747,18 @@ MemoryContextStatsDetail(MemoryContext context, int max_children,
 								 grand_totals.totalspace - grand_totals.freespace)));
 }
 
+/*
+ * Return the memory used in the given context and its children.
+ */
+extern void
+MemoryContextMemConsumed(MemoryContext context,
+						 MemoryContextCounters *mem_consumed)
+{
+	memset(mem_consumed, 0, sizeof(*mem_consumed));
+
+	MemoryContextStatsInternal(context, 0, false, 100, mem_consumed, false);
+}
+
 /*
  * MemoryContextStatsInternal
  *		One recursion level for MemoryContextStats
diff --git a/src/include/commands/explain.h b/src/include/commands/explain.h
index 3d3e632a0c..da2aeb647a 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -62,6 +62,12 @@ typedef struct ExplainState
 	ExplainWorkersState *workers_state; /* needed if parallel plan */
 } ExplainState;
 
+typedef struct MemUsage
+{
+	Size		mem_used;
+	Size		mem_allocated;
+} MemUsage;
+
 /* Hook for plugins to get control in ExplainOneQuery() */
 typedef void (*ExplainOneQuery_hook_type) (Query *query,
 										   int cursorOptions,
@@ -92,7 +98,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 MemUsage *mem_usage);
 
 extern void ExplainPrintPlan(ExplainState *es, QueryDesc *queryDesc);
 extern void ExplainPrintTriggers(ExplainState *es, QueryDesc *queryDesc);
@@ -125,5 +132,8 @@ extern void ExplainOpenGroup(const char *objtype, const char *labelname,
 							 bool labeled, ExplainState *es);
 extern void ExplainCloseGroup(const char *objtype, const char *labelname,
 							  bool labeled, ExplainState *es);
+extern void calc_mem_usage(MemUsage *mem_usage,
+						   MemoryContextCounters *mem_counts_end,
+						   MemoryContextCounters *mem_counts_start);
 
 #endif							/* EXPLAIN_H */
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 21640d62a6..d254a044c1 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -92,6 +92,8 @@ extern void MemoryContextStatsDetail(MemoryContext context, int max_children,
 									 bool print_to_stderr);
 extern void MemoryContextAllowInCriticalSection(MemoryContext context,
 												bool allow);
+extern void MemoryContextMemConsumed(MemoryContext context,
+									 MemoryContextCounters *mem_consumed);
 
 #ifdef MEMORY_CONTEXT_CHECKING
 extern void MemoryContextCheck(MemoryContext context);
diff --git a/src/test/regress/expected/explain.out b/src/test/regress/expected/explain.out
index 1aca77491b..863d3d79d2 100644
--- a/src/test/regress/expected/explain.out
+++ b/src/test/regress/expected/explain.out
@@ -65,8 +65,9 @@ select explain_filter('explain (analyze) select * from int8_tbl i8');
 -----------------------------------------------------------------------------------------------
  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)
  Planning Time: N.N ms
+ Planning Memory: used=N bytes allocated=N bytes
  Execution Time: N.N ms
-(3 rows)
+(4 rows)
 
 select explain_filter('explain (analyze, verbose) select * from int8_tbl i8');
                                             explain_filter                                            
@@ -74,16 +75,18 @@ select explain_filter('explain (analyze, verbose) select * from int8_tbl i8');
  Seq Scan on public.int8_tbl i8  (cost=N.N..N.N rows=N width=N) (actual time=N.N..N.N rows=N loops=N)
    Output: q1, q2
  Planning Time: N.N ms
+ Planning Memory: used=N bytes allocated=N bytes
  Execution Time: N.N ms
-(4 rows)
+(5 rows)
 
 select explain_filter('explain (analyze, buffers, format text) 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)
  Planning Time: N.N ms
+ Planning Memory: used=N bytes allocated=N bytes
  Execution Time: N.N ms
-(3 rows)
+(4 rows)
 
 select explain_filter('explain (analyze, buffers, format xml) select * from int8_tbl i8');
                      explain_filter                     
@@ -128,6 +131,10 @@ select explain_filter('explain (analyze, buffers, format xml) select * from int8
        <Temp-Written-Blocks>N</Temp-Written-Blocks>    +
      </Planning>                                       +
      <Planning-Time>N.N</Planning-Time>                +
+     <Planning-Memory>                                 +
+       <Used>N</Used>                                  +
+       <Allocated>N</Allocated>                        +
+     </Planning-Memory>                                +
      <Triggers>                                        +
      </Triggers>                                       +
      <Execution-Time>N.N</Execution-Time>              +
@@ -174,6 +181,9 @@ select explain_filter('explain (analyze, buffers, format yaml) select * from int
      Temp Read Blocks: N      +
      Temp Written Blocks: N   +
    Planning Time: N.N         +
+   Planning Memory:           +
+     Used: N                  +
+     Allocated: N             +
    Triggers:                  +
    Execution Time: N.N
 (1 row)
@@ -280,6 +290,10 @@ select explain_filter('explain (analyze, buffers, format json) select * from int
        "Temp I/O Write Time": N.N  +
      },                            +
      "Planning Time": N.N,         +
+     "Planning Memory": {          +
+       "Used": N,                  +
+       "Allocated": N              +
+     },                            +
      "Triggers": [                 +
      ],                            +
      "Execution Time": N.N         +
@@ -531,7 +545,11 @@ select jsonb_pretty(
          "Triggers": [                                      +
          ],                                                 +
          "Planning Time": 0.0,                              +
-         "Execution Time": 0.0                              +
+         "Execution Time": 0.0,                             +
+         "Planning Memory": {                               +
+             "Used": 0,                                     +
+             "Allocated": 0                                 +
+         }                                                  +
      }                                                      +
  ]
 (1 row)
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index f2af84d7ca..78ab04bef8 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1544,6 +1544,7 @@ MemoryContextData
 MemoryContextMethodID
 MemoryContextMethods
 MemoryStatsPrintFunc
+MemUsage
 MergeAction
 MergeActionState
 MergeAppend
-- 
2.25.1