From 63ed126ad4ffad793292c8688d6fb1f9f29fdd16 Mon Sep 17 00:00:00 2001
From: Sait Talha Nisanci <s.talhanisanci@gmail.com>
Date: Fri, 24 Jul 2020 16:44:50 +0300
Subject: [PATCH] Set buffer usage with explain

It seems that currently we process even postgres tables in explain
commands. This is because we register a hook for explain and we don't
have any check to see if the query has any citus table.

With this commit, we now send the buffer usage as well to the relevant
API. There is some duplicate in the code but it is because of the
existing structure, we can refactor this separately.
---
 .../distributed/planner/multi_explain.c       | 60 ++++++++++++++++++-
 src/include/distributed/version_compat.h      |  6 +-
 2 files changed, 60 insertions(+), 6 deletions(-)

diff --git a/src/backend/distributed/planner/multi_explain.c b/src/backend/distributed/planner/multi_explain.c
index f386f3123..414531c56 100644
--- a/src/backend/distributed/planner/multi_explain.c
+++ b/src/backend/distributed/planner/multi_explain.c
@@ -272,7 +272,16 @@ ExplainSubPlans(DistributedPlan *distributedPlan, ExplainState *es)
 		ParamListInfo params = NULL;
 		char *queryString = NULL;
 		instr_time planduration;
+		#if PG_VERSION_NUM >= PG_VERSION_13
 
+		BufferUsage bufusage_start,
+					bufusage;
+
+		if (es->buffers)
+		{
+			bufusage_start = pgBufferUsage;
+		}
+		#endif
 		if (es->format == EXPLAIN_FORMAT_TEXT)
 		{
 			char *resultId = GenerateResultId(planId, subPlan->subPlanId);
@@ -314,7 +323,18 @@ ExplainSubPlans(DistributedPlan *distributedPlan, ExplainState *es)
 
 		INSTR_TIME_SET_ZERO(planduration);
 
-		ExplainOnePlanCompat(plan, into, es, queryString, params, NULL, &planduration);
+		#if PG_VERSION_NUM >= PG_VERSION_13
+
+		/* calc differences of buffer counters. */
+		if (es->buffers)
+		{
+			memset(&bufusage, 0, sizeof(BufferUsage));
+			BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start);
+		}
+		#endif
+
+		ExplainOnePlanCompat(plan, into, es, queryString, params, NULL, &planduration,
+							 (es->buffers ? &bufusage : NULL));
 
 		if (es->format == EXPLAIN_FORMAT_TEXT)
 		{
@@ -1123,6 +1143,15 @@ CitusExplainOneQuery(Query *query, int cursorOptions, IntoClause *into,
 	/* rest is copied from ExplainOneQuery() */
 	instr_time planstart,
 			   planduration;
+	#if PG_VERSION_NUM >= PG_VERSION_13
+	BufferUsage bufusage_start,
+				bufusage;
+
+	if (es->buffers)
+	{
+		bufusage_start = pgBufferUsage;
+	}
+	#endif
 
 	INSTR_TIME_SET_CURRENT(planstart);
 
@@ -1132,9 +1161,19 @@ CitusExplainOneQuery(Query *query, int cursorOptions, IntoClause *into,
 	INSTR_TIME_SET_CURRENT(planduration);
 	INSTR_TIME_SUBTRACT(planduration, planstart);
 
+	#if PG_VERSION_NUM >= PG_VERSION_13
+
+	/* calc differences of buffer counters. */
+	if (es->buffers)
+	{
+		memset(&bufusage, 0, sizeof(BufferUsage));
+		BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start);
+	}
+	#endif
+
 	/* run it (if needed) and produce output */
 	ExplainOnePlanCompat(plan, into, es, queryString, params, queryEnv,
-						 &planduration);
+						 &planduration, (es->buffers ? &bufusage : NULL));
 }
 
 
@@ -1454,7 +1493,13 @@ ExplainOneQuery(Query *query, int cursorOptions,
 	{
 		instr_time	planstart,
 					planduration;
+		#if PG_VERSION_NUM >= PG_VERSION_13
+		BufferUsage bufusage_start,
+			    bufusage;
 
+		if (es->buffers)
+			bufusage_start = pgBufferUsage;
+		#endif
 		INSTR_TIME_SET_CURRENT(planstart);
 
 		/* plan the query */
@@ -1463,9 +1508,18 @@ ExplainOneQuery(Query *query, int cursorOptions,
 		INSTR_TIME_SET_CURRENT(planduration);
 		INSTR_TIME_SUBTRACT(planduration, planstart);
 
+		#if PG_VERSION_NUM >= PG_VERSION_13
+
+		/* calc differences of buffer counters. */
+		if (es->buffers)
+		{
+			memset(&bufusage, 0, sizeof(BufferUsage));
+			BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start);
+		}
+		#endif
 		/* run it (if needed) and produce output */
 		ExplainOnePlanCompat(plan, into, es, queryString, params, queryEnv,
-					   &planduration);
+					   &planduration, (es->buffers ? &bufusage : NULL));
 	}
 }
 
diff --git a/src/include/distributed/version_compat.h b/src/include/distributed/version_compat.h
index 4bda13761..d140a66b7 100644
--- a/src/include/distributed/version_compat.h
+++ b/src/include/distributed/version_compat.h
@@ -36,8 +36,8 @@
 #define standard_planner_compat(a, c, d) standard_planner(a, NULL, c, d)
 #define getOwnedSequencesCompat(a, b) getOwnedSequences(a)
 #define CMDTAG_SELECT_COMPAT CMDTAG_SELECT
-#define ExplainOnePlanCompat(a, b, c, d, e, f, g) ExplainOnePlan(a, b, c, d, e, f, g, \
-																 NULL)
+#define ExplainOnePlanCompat(a, b, c, d, e, f, g, h) ExplainOnePlan(a, b, c, d, e, f, g, \
+																	h)
 #define SetListCellPtr(a, b) ((a)->ptr_value = (b))
 #define RangeTableEntryFromNSItem(a) ((a)->p_rte)
 #define QueryCompletionCompat QueryCompletion
@@ -49,7 +49,7 @@
 #define standard_planner_compat(a, c, d) standard_planner(a, c, d)
 #define CMDTAG_SELECT_COMPAT "SELECT"
 #define getOwnedSequencesCompat(a, b) getOwnedSequences(a, b)
-#define ExplainOnePlanCompat(a, b, c, d, e, f, g) ExplainOnePlan(a, b, c, d, e, f, g)
+#define ExplainOnePlanCompat(a, b, c, d, e, f, g, h) ExplainOnePlan(a, b, c, d, e, f, g)
 #define SetListCellPtr(a, b) ((a)->data.ptr_value = (b))
 #define RangeTableEntryFromNSItem(a) (a)
 #define QueryCompletionCompat char