From a23bf8803d6d2554edbbf4c88749c7c44792db71 Mon Sep 17 00:00:00 2001
From: naisila <nicypp@gmail.com>
Date: Wed, 10 Jul 2024 16:12:51 +0200
Subject: [PATCH] EXPLAIN (MEMORY) is added, make changes to ExplainOnePlan

Relevant PG commit:
5de890e3610d5a12cdaea36413d967cf5c544e20
https://github.com/postgres/postgres/commit/5de890e3610d5a12cdaea36413d967cf5c544e20
---
 .../distributed/planner/multi_explain.c       | 83 +++++++++++++++++++
 1 file changed, 83 insertions(+)

diff --git a/src/backend/distributed/planner/multi_explain.c b/src/backend/distributed/planner/multi_explain.c
index b70203b90..531b9b362 100644
--- a/src/backend/distributed/planner/multi_explain.c
+++ b/src/backend/distributed/planner/multi_explain.c
@@ -373,6 +373,21 @@ ExplainSubPlans(DistributedPlan *distributedPlan, ExplainState *es)
 		BufferUsage bufusage_start,
 					bufusage;
 
+#if PG_VERSION_NUM >= PG_VERSION_17
+		MemoryContextCounters mem_counters;
+		MemoryContext planner_ctx = NULL;
+		MemoryContext saved_ctx = NULL;
+
+		if (es->memory)
+		{
+			/* copy paste from postgres code */
+			planner_ctx = AllocSetContextCreate(CurrentMemoryContext,
+												"explain analyze planner context",
+												ALLOCSET_DEFAULT_SIZES);
+			saved_ctx = MemoryContextSwitchTo(planner_ctx);
+		}
+#endif
+
 		if (es->buffers)
 		{
 			bufusage_start = pgBufferUsage;
@@ -430,8 +445,20 @@ ExplainSubPlans(DistributedPlan *distributedPlan, ExplainState *es)
 
 		ExplainOpenGroup("PlannedStmt", "PlannedStmt", false, es);
 
+#if PG_VERSION_NUM >= PG_VERSION_17
+		if (es->memory)
+		{
+			MemoryContextSwitchTo(saved_ctx);
+			MemoryContextMemConsumed(planner_ctx, &mem_counters);
+		}
+
+		ExplainOnePlan(plan, into, es, queryString, params, NULL, &planduration,
+					   (es->buffers ? &bufusage : NULL),
+					   (es->memory ? &mem_counters : NULL));
+#else
 		ExplainOnePlan(plan, into, es, queryString, params, NULL, &planduration,
 					   (es->buffers ? &bufusage : NULL));
+#endif
 
 		ExplainCloseGroup("PlannedStmt", "PlannedStmt", false, es);
 		ExplainCloseGroup("Subplan", NULL, true, es);
@@ -1251,6 +1278,21 @@ CitusExplainOneQuery(Query *query, int cursorOptions, IntoClause *into,
 	BufferUsage bufusage_start,
 				bufusage;
 
+#if PG_VERSION_NUM >= PG_VERSION_17
+	MemoryContextCounters mem_counters;
+	MemoryContext planner_ctx = NULL;
+	MemoryContext saved_ctx = NULL;
+
+	if (es->memory)
+	{
+		/* copy paste from postgres code */
+		planner_ctx = AllocSetContextCreate(CurrentMemoryContext,
+											"explain analyze planner context",
+											ALLOCSET_DEFAULT_SIZES);
+		saved_ctx = MemoryContextSwitchTo(planner_ctx);
+	}
+#endif
+
 	if (es->buffers)
 	{
 		bufusage_start = pgBufferUsage;
@@ -1284,9 +1326,23 @@ CitusExplainOneQuery(Query *query, int cursorOptions, IntoClause *into,
 		BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start);
 	}
 
+#if PG_VERSION_NUM >= PG_VERSION_17
+	if (es->memory)
+	{
+		MemoryContextSwitchTo(saved_ctx);
+		MemoryContextMemConsumed(planner_ctx, &mem_counters);
+	}
+
+	/* run it (if needed) and produce output */
+	ExplainOnePlan(plan, into, es, queryString, params, queryEnv,
+				   &planduration, (es->buffers ? &bufusage : NULL),
+				   (es->memory ? &mem_counters : NULL));
+#else
+
 	/* run it (if needed) and produce output */
 	ExplainOnePlan(plan, into, es, queryString, params, queryEnv,
 				   &planduration, (es->buffers ? &bufusage : NULL));
+#endif
 }
 
 
@@ -1699,6 +1755,21 @@ ExplainOneQuery(Query *query, int cursorOptions,
 		BufferUsage bufusage_start,
 			    bufusage;
 
+#if PG_VERSION_NUM >= PG_VERSION_17
+		MemoryContextCounters mem_counters;
+		MemoryContext planner_ctx = NULL;
+		MemoryContext saved_ctx = NULL;
+
+		if (es->memory)
+		{
+			/* copy paste from postgres code */
+			planner_ctx = AllocSetContextCreate(CurrentMemoryContext,
+												"explain analyze planner context",
+												ALLOCSET_DEFAULT_SIZES);
+			saved_ctx = MemoryContextSwitchTo(planner_ctx);
+		}
+#endif
+
 		if (es->buffers)
 			bufusage_start = pgBufferUsage;
 		INSTR_TIME_SET_CURRENT(planstart);
@@ -1716,9 +1787,21 @@ ExplainOneQuery(Query *query, int cursorOptions,
 			BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start);
 		}
 
+#if PG_VERSION_NUM >= PG_VERSION_17
+		if (es->memory)
+		{
+			MemoryContextSwitchTo(saved_ctx);
+			MemoryContextMemConsumed(planner_ctx, &mem_counters);
+		}
+		/* run it (if needed) and produce output */
+		ExplainOnePlan(plan, into, es, queryString, params, queryEnv,
+					   &planduration, (es->buffers ? &bufusage : NULL),
+					   (es->memory ? &mem_counters : NULL));
+#else
 		/* run it (if needed) and produce output */
 		ExplainOnePlan(plan, into, es, queryString, params, queryEnv,
 					   &planduration, (es->buffers ? &bufusage : NULL));
+#endif
 	}
 }