From 314e3e7305da740a8fadb2c481a096cc0ca7fff0 Mon Sep 17 00:00:00 2001
From: Lukas Fittl <lukas@fittl.com>
Date: Tue, 9 Sep 2025 02:26:02 -0700
Subject: [PATCH v3 6/7] Introduce alternate Instrumentation stack mechanism
 relying on PG_FINALLY

The resource owner-based Instrumentation stack cannot handle wrapping
certain utility commands that close and re-open the top-level transaction,
like the CLUSTER command. This is a problem for pg_stat_statements tracking
of utility commands specifically. To support tracking such activity, allow
issuing explicit InstrPushStack/InstrPopStack commands to modify the stack,
with the InstrPopStack in a PG_FINALLY to ensure cleanup on abort.
---
 .../pg_stat_statements/pg_stat_statements.c   | 50 +++++--------------
 src/include/executor/instrument.h             |  3 ++
 2 files changed, 15 insertions(+), 38 deletions(-)

diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index eeabd820d8e..0e57ce65062 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -911,21 +911,13 @@ pgss_planner(Query *parse,
 	{
 		instr_time	start;
 		instr_time	duration;
-		BufferUsage bufusage_start,
-					bufusage;
-		WalUsage	walusage_start,
-					walusage;
+		InstrStack *stack;
 
-		/* We need to track buffer usage as the planner can access them. */
-		bufusage_start = pgBufferUsage;
-
-		/*
-		 * Similarly the planner could write some WAL records in some cases
-		 * (e.g. setting a hint bit with those being WAL-logged)
-		 */
-		walusage_start = pgWalUsage;
 		INSTR_TIME_SET_CURRENT(start);
 
+		/* We need to track buffer/WAL usage as the planner can access them. */
+		stack = InstrPushStack();
+
 		nesting_level++;
 		PG_TRY();
 		{
@@ -938,6 +930,7 @@ pgss_planner(Query *parse,
 		}
 		PG_FINALLY();
 		{
+			InstrPopStack(stack);
 			nesting_level--;
 		}
 		PG_END_TRY();
@@ -945,14 +938,6 @@ pgss_planner(Query *parse,
 		INSTR_TIME_SET_CURRENT(duration);
 		INSTR_TIME_SUBTRACT(duration, start);
 
-		/* calc differences of buffer counters. */
-		memset(&bufusage, 0, sizeof(BufferUsage));
-		BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start);
-
-		/* calc differences of WAL counters. */
-		memset(&walusage, 0, sizeof(WalUsage));
-		WalUsageAccumDiff(&walusage, &pgWalUsage, &walusage_start);
-
 		pgss_store(query_string,
 				   parse->queryId,
 				   parse->stmt_location,
@@ -960,8 +945,8 @@ pgss_planner(Query *parse,
 				   PGSS_PLAN,
 				   INSTR_TIME_GET_MILLISEC(duration),
 				   0,
-				   &bufusage,
-				   &walusage,
+				   &stack->bufusage,
+				   &stack->walusage,
 				   NULL,
 				   NULL,
 				   0,
@@ -1157,14 +1142,10 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 		instr_time	start;
 		instr_time	duration;
 		uint64		rows;
-		BufferUsage bufusage_start,
-					bufusage;
-		WalUsage	walusage_start,
-					walusage;
+		InstrStack *stack;
 
-		bufusage_start = pgBufferUsage;
-		walusage_start = pgWalUsage;
 		INSTR_TIME_SET_CURRENT(start);
+		stack = InstrPushStack();
 
 		nesting_level++;
 		PG_TRY();
@@ -1180,6 +1161,7 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 		}
 		PG_FINALLY();
 		{
+			InstrPopStack(stack);
 			nesting_level--;
 		}
 		PG_END_TRY();
@@ -1208,14 +1190,6 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 					   qc->commandTag == CMDTAG_REFRESH_MATERIALIZED_VIEW)) ?
 			qc->nprocessed : 0;
 
-		/* calc differences of buffer counters. */
-		memset(&bufusage, 0, sizeof(BufferUsage));
-		BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start);
-
-		/* calc differences of WAL counters. */
-		memset(&walusage, 0, sizeof(WalUsage));
-		WalUsageAccumDiff(&walusage, &pgWalUsage, &walusage_start);
-
 		pgss_store(queryString,
 				   saved_queryId,
 				   saved_stmt_location,
@@ -1223,8 +1197,8 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 				   PGSS_EXEC,
 				   INSTR_TIME_GET_MILLISEC(duration),
 				   rows,
-				   &bufusage,
-				   &walusage,
+				   &stack->bufusage,
+				   &stack->walusage,
 				   NULL,
 				   NULL,
 				   0,
diff --git a/src/include/executor/instrument.h b/src/include/executor/instrument.h
index bf766706580..8804ee64311 100644
--- a/src/include/executor/instrument.h
+++ b/src/include/executor/instrument.h
@@ -147,6 +147,9 @@ extern Instrumentation *InstrAlloc(int n, int instrument_options);
 extern void InstrStart(Instrumentation *instr);
 extern void InstrStop(Instrumentation *instr, double nTuples, bool finalize);
 
+extern InstrStack * InstrPushStack(void);
+extern void InstrPopStack(InstrStack * res);
+
 extern NodeInstrumentation * InstrAllocNode(int n, int instrument_options,
 											bool async_mode);
 extern void InstrInitNode(NodeInstrumentation * instr, int instrument_options);
-- 
2.47.1