From 58919a10cf5542495e7af3afa9a908a923f3ddf6 Mon Sep 17 00:00:00 2001
From: Lukas Fittl <lukas@fittl.com>
Date: Tue, 9 Sep 2025 02:16:59 -0700
Subject: [PATCH v16 01/10] instrumentation: Allocate queryDesc->totaltime in
 ExecutorStart if needed

This introduces a new field, queryDesc->totaltime_options, that extensions
can use to indicate whether they need queryDesc->totaltime populated,
and with which instrumentation options. Extensions should take care to
only add options they need, instead of replacing the options of others.

This replaces the practice of extensions allocating queryDesc->totaltime
themselves, which required them to always use INSTRUMENT_ALL for the
options argument. If they wouldn't have, another extension could silently
be impacted by it. It also unnecessarily made extensions hooks worry
about being sure to allocate in the per-query memory context.

Adjust pg_stat_statements and auto_explain to match, and lower the
requested instrumentation level for auto_explain to INSTRUMENT_TIMER,
since the summary instrumentation it needs is only runtime.

Author: Lukas Fittl <lukas@fittl.com>
Reviewed-by:
Discussion:
---
 contrib/auto_explain/auto_explain.c           | 20 +++------------
 .../pg_stat_statements/pg_stat_statements.c   | 25 ++++++-------------
 src/backend/executor/execMain.c               |  9 +++++++
 src/backend/tcop/pquery.c                     |  1 +
 src/include/executor/execdesc.h               |  4 ++-
 5 files changed, 23 insertions(+), 36 deletions(-)

diff --git a/contrib/auto_explain/auto_explain.c b/contrib/auto_explain/auto_explain.c
index 6ceae1c69ce..4f9d35bc30b 100644
--- a/contrib/auto_explain/auto_explain.c
+++ b/contrib/auto_explain/auto_explain.c
@@ -334,6 +334,9 @@ explain_ExecutorStart(QueryDesc *queryDesc, int eflags)
 
 	if (auto_explain_enabled())
 	{
+		/* We're always interested in runtime */
+		queryDesc->totaltime_options |= INSTRUMENT_TIMER;
+
 		/* Enable per-node instrumentation iff log_analyze is required. */
 		if (auto_explain_log_analyze && (eflags & EXEC_FLAG_EXPLAIN_ONLY) == 0)
 		{
@@ -352,23 +355,6 @@ explain_ExecutorStart(QueryDesc *queryDesc, int eflags)
 		prev_ExecutorStart(queryDesc, eflags);
 	else
 		standard_ExecutorStart(queryDesc, eflags);
-
-	if (auto_explain_enabled())
-	{
-		/*
-		 * Set up to track total elapsed time in ExecutorRun.  Make sure the
-		 * space is allocated in the per-query context so it will go away at
-		 * ExecutorEnd.
-		 */
-		if (queryDesc->totaltime == NULL)
-		{
-			MemoryContext oldcxt;
-
-			oldcxt = MemoryContextSwitchTo(queryDesc->estate->es_query_cxt);
-			queryDesc->totaltime = InstrAlloc(INSTRUMENT_ALL);
-			MemoryContextSwitchTo(oldcxt);
-		}
-	}
 }
 
 /*
diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index b5000bc14b7..c8981dcb5cf 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -993,11 +993,6 @@ pgss_planner(Query *parse,
 static void
 pgss_ExecutorStart(QueryDesc *queryDesc, int eflags)
 {
-	if (prev_ExecutorStart)
-		prev_ExecutorStart(queryDesc, eflags);
-	else
-		standard_ExecutorStart(queryDesc, eflags);
-
 	/*
 	 * If query has queryId zero, don't track it.  This prevents double
 	 * counting of optimizable statements that are directly contained in
@@ -1005,20 +1000,14 @@ pgss_ExecutorStart(QueryDesc *queryDesc, int eflags)
 	 */
 	if (pgss_enabled(nesting_level) && queryDesc->plannedstmt->queryId != INT64CONST(0))
 	{
-		/*
-		 * Set up to track total elapsed time in ExecutorRun.  Make sure the
-		 * space is allocated in the per-query context so it will go away at
-		 * ExecutorEnd.
-		 */
-		if (queryDesc->totaltime == NULL)
-		{
-			MemoryContext oldcxt;
-
-			oldcxt = MemoryContextSwitchTo(queryDesc->estate->es_query_cxt);
-			queryDesc->totaltime = InstrAlloc(INSTRUMENT_ALL);
-			MemoryContextSwitchTo(oldcxt);
-		}
+		/* Request all summary instrumentation, i.e. timing, buffers and WAL */
+		queryDesc->totaltime_options |= INSTRUMENT_ALL;
 	}
+
+	if (prev_ExecutorStart)
+		prev_ExecutorStart(queryDesc, eflags);
+	else
+		standard_ExecutorStart(queryDesc, eflags);
 }
 
 /*
diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index b0f636bf8b6..f71f668883c 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -250,6 +250,15 @@ standard_ExecutorStart(QueryDesc *queryDesc, int eflags)
 	estate->es_instrument = queryDesc->instrument_options;
 	estate->es_jit_flags = queryDesc->plannedstmt->jitFlags;
 
+	/*
+	 * Set up query-level instrumentation if extensions have requested it via
+	 * totaltime_options. Ensure an extension has not allocated totaltime
+	 * itself.
+	 */
+	Assert(queryDesc->totaltime == NULL);
+	if (queryDesc->totaltime_options)
+		queryDesc->totaltime = InstrAlloc(queryDesc->totaltime_options);
+
 	/*
 	 * Set up an AFTER-trigger statement context, unless told not to, or
 	 * unless it's EXPLAIN-only mode (when ExecutorFinish won't be called).
diff --git a/src/backend/tcop/pquery.c b/src/backend/tcop/pquery.c
index d8fc75d0bb9..e27f26ecd83 100644
--- a/src/backend/tcop/pquery.c
+++ b/src/backend/tcop/pquery.c
@@ -86,6 +86,7 @@ CreateQueryDesc(PlannedStmt *plannedstmt,
 	qd->params = params;		/* parameter values passed into query */
 	qd->queryEnv = queryEnv;
 	qd->instrument_options = instrument_options;	/* instrumentation wanted? */
+	qd->totaltime_options = 0;
 
 	/* null these fields until set by ExecutorStart */
 	qd->tupDesc = NULL;
diff --git a/src/include/executor/execdesc.h b/src/include/executor/execdesc.h
index d3a57242844..3f9143efb8c 100644
--- a/src/include/executor/execdesc.h
+++ b/src/include/executor/execdesc.h
@@ -42,6 +42,8 @@ typedef struct QueryDesc
 	ParamListInfo params;		/* param values being passed in */
 	QueryEnvironment *queryEnv; /* query environment passed in */
 	int			instrument_options; /* OR of InstrumentOption flags */
+	int			totaltime_options;	/* OR of InstrumentOption flags for
+									 * totaltime */
 
 	/* These fields are set by ExecutorStart */
 	TupleDesc	tupDesc;		/* descriptor for result tuples */
@@ -51,7 +53,7 @@ typedef struct QueryDesc
 	/* This field is set by ExecutePlan */
 	bool		already_executed;	/* true if previously executed */
 
-	/* This is always set NULL by the core system, but plugins can change it */
+	/* This field is allocated by ExecutorStart if needed */
 	struct Instrumentation *totaltime;	/* total time spent in ExecutorRun */
 } QueryDesc;
 
-- 
2.47.1

