From 412c26f58082007e08189d0763cc8a2598cbec26 Mon Sep 17 00:00:00 2001
From: "Imseih (AWS)" <simseih@88665a22795f.ant.amazon.com>
Date: Tue, 24 Jan 2023 12:26:44 -0600
Subject: [PATCH 1/1] correct pg_stat_statements tracking of portals

---
 contrib/pg_stat_statements/pg_stat_statements.c | 13 +++++++++----
 src/backend/executor/execMain.c                 |  6 ++++++
 src/backend/executor/execUtils.c                |  2 ++
 src/include/nodes/execnodes.h                   |  4 +++-
 4 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index ad1fe44496..63bbfd06ba 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -339,7 +339,7 @@ static void pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 static void pgss_store(const char *query, uint64 queryId,
 					   int query_location, int query_len,
 					   pgssStoreKind kind,
-					   double total_time, uint64 rows,
+					   double total_time, uint64 rows, uint64 calls,
 					   const BufferUsage *bufusage,
 					   const WalUsage *walusage,
 					   const struct JitInstrumentation *jitusage,
@@ -855,6 +855,7 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
 				   PGSS_INVALID,
 				   0,
 				   0,
+				   0,
 				   NULL,
 				   NULL,
 				   NULL,
@@ -940,6 +941,7 @@ pgss_planner(Query *parse,
 				   PGSS_PLAN,
 				   INSTR_TIME_GET_MILLISEC(duration),
 				   0,
+				   1,
 				   &bufusage,
 				   &walusage,
 				   NULL,
@@ -1000,6 +1002,7 @@ pgss_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count,
 				 bool execute_once)
 {
 	exec_nested_level++;
+
 	PG_TRY();
 	{
 		if (prev_ExecutorRun)
@@ -1058,7 +1061,8 @@ pgss_ExecutorEnd(QueryDesc *queryDesc)
 				   queryDesc->plannedstmt->stmt_len,
 				   PGSS_EXEC,
 				   queryDesc->totaltime->total * 1000.0,	/* convert to msec */
-				   queryDesc->estate->es_processed,
+				   queryDesc->estate->es_total_processed,
+				   queryDesc->estate->es_calls,
 				   &queryDesc->totaltime->bufusage,
 				   &queryDesc->totaltime->walusage,
 				   queryDesc->estate->es_jit ? &queryDesc->estate->es_jit->instr : NULL,
@@ -1189,6 +1193,7 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 				   PGSS_EXEC,
 				   INSTR_TIME_GET_MILLISEC(duration),
 				   rows,
+				   1,
 				   &bufusage,
 				   &walusage,
 				   NULL,
@@ -1222,7 +1227,7 @@ static void
 pgss_store(const char *query, uint64 queryId,
 		   int query_location, int query_len,
 		   pgssStoreKind kind,
-		   double total_time, uint64 rows,
+		   double total_time, uint64 rows, uint64 calls,
 		   const BufferUsage *bufusage,
 		   const WalUsage *walusage,
 		   const struct JitInstrumentation *jitusage,
@@ -1348,7 +1353,7 @@ pgss_store(const char *query, uint64 queryId,
 		if (IS_STICKY(e->counters))
 			e->counters.usage = USAGE_INIT;
 
-		e->counters.calls[kind] += 1;
+		e->counters.calls[kind] += calls;
 		e->counters.total_time[kind] += total_time;
 
 		if (e->counters.calls[kind] == 1)
diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index a5115b9c1f..3ccf6f26e8 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -341,6 +341,11 @@ standard_ExecutorRun(QueryDesc *queryDesc,
 	operation = queryDesc->operation;
 	dest = queryDesc->dest;
 
+	/*
+	 * update the counter for total_processed and calls
+	 */
+	queryDesc->estate->es_total_processed += queryDesc->estate->es_processed;
+	estate->es_calls++;
 	/*
 	 * startup tuple receiver, if we will be emitting tuples
 	 */
@@ -444,6 +449,7 @@ standard_ExecutorFinish(QueryDesc *queryDesc)
 	MemoryContextSwitchTo(oldcontext);
 
 	estate->es_finished = true;
+	queryDesc->estate->es_total_processed += queryDesc->estate->es_processed;
 }
 
 /* ----------------------------------------------------------------
diff --git a/src/backend/executor/execUtils.c b/src/backend/executor/execUtils.c
index c33a3c0bec..98f9c556b8 100644
--- a/src/backend/executor/execUtils.c
+++ b/src/backend/executor/execUtils.c
@@ -146,6 +146,8 @@ CreateExecutorState(void)
 	estate->es_tupleTable = NIL;
 
 	estate->es_processed = 0;
+	estate->es_calls = 0;
+	estate->es_total_processed = 0;
 
 	estate->es_top_eflags = 0;
 	estate->es_instrument = 0;
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 20f4c8b35f..dd1f389a7d 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -657,7 +657,9 @@ typedef struct EState
 
 	List	   *es_tupleTable;	/* List of TupleTableSlots */
 
-	uint64		es_processed;	/* # of tuples processed */
+	uint64		es_processed;	/* # of tuples processed at the top level only */
+	uint64		es_calls;       /* # of calls */
+	uint64		es_total_processed; /* total # of tuples processed */
 
 	int			es_top_eflags;	/* eflags passed to ExecutorStart */
 	int			es_instrument;	/* OR of InstrumentOption flags */
-- 
2.37.1 (Apple Git-137.1)

