v1-0001-Allow-plugins-to-set-a-64-bit-plan-identifier.patch

From 2bb347266b281e6afaf871867fdfa08350ef49eb Mon Sep 17 00:00:00 2001
From: "Sami Imseih (AWS)"
 <simseih@dev-dsk-simseih-1d-3940b79e.us-east-1.amazon.com>
Date: Wed, 19 Feb 2025 15:56:21 +0000
Subject: [PATCH v1 1/2] Allow plugins to set a 64-bit plan identifier

---
 src/backend/executor/execParallel.c         |  1 +
 src/backend/tcop/postgres.c                 | 13 ++++++
 src/backend/utils/activity/backend_status.c | 46 +++++++++++++++++++++
 src/include/nodes/plannodes.h               |  3 ++
 src/include/utils/backend_status.h          |  5 +++
 5 files changed, 68 insertions(+)

diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index 1bedb80836..e3ee449748 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -174,6 +174,7 @@ ExecSerializePlan(Plan *plan, EState *estate)
 	pstmt = makeNode(PlannedStmt);
 	pstmt->commandType = CMD_SELECT;
 	pstmt->queryId = pgstat_get_my_query_id();
+	pstmt->planId = pgstat_get_my_plan_id();
 	pstmt->hasReturning = false;
 	pstmt->hasModifyingCTE = false;
 	pstmt->canSetTag = true;
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index f2f75aa0f8..b828a982e1 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -1106,6 +1106,7 @@ exec_simple_query(const char *query_string)
 		size_t		cmdtaglen;
 
 		pgstat_report_query_id(0, true);
+		pgstat_report_plan_id(0, true);
 
 		/*
 		 * Get the command name for use in status display (it also becomes the
@@ -2015,6 +2016,17 @@ exec_bind_message(StringInfo input_message)
 	 */
 	cplan = GetCachedPlan(psrc, params, NULL, NULL);
 
+	foreach(lc, cplan->stmt_list)
+	{
+		PlannedStmt *plan = lfirst_node(PlannedStmt, lc);
+
+		if (plan->planId != UINT64CONST(0))
+		{
+			pgstat_report_plan_id(plan->planId, false);
+			break;
+		}
+	}
+
 	/*
 	 * Now we can define the portal.
 	 *
@@ -2165,6 +2177,7 @@ exec_execute_message(const char *portal_name, long max_rows)
 		if (stmt->queryId != UINT64CONST(0))
 		{
 			pgstat_report_query_id(stmt->queryId, false);
+			pgstat_report_plan_id(stmt->planId, false);
 			break;
 		}
 	}
diff --git a/src/backend/utils/activity/backend_status.c b/src/backend/utils/activity/backend_status.c
index 5f68ef26ad..65f2e4e0c7 100644
--- a/src/backend/utils/activity/backend_status.c
+++ b/src/backend/utils/activity/backend_status.c
@@ -379,6 +379,7 @@ pgstat_bestart(void)
 	lbeentry.st_progress_command = PROGRESS_COMMAND_INVALID;
 	lbeentry.st_progress_command_target = InvalidOid;
 	lbeentry.st_query_id = UINT64CONST(0);
+	lbeentry.st_plan_id = UINT64CONST(0);
 
 	/*
 	 * we don't zero st_progress_param here to save cycles; nobody should
@@ -533,6 +534,7 @@ pgstat_report_activity(BackendState state, const char *cmd_str)
 			/* st_xact_start_timestamp and wait_event_info are also disabled */
 			beentry->st_xact_start_timestamp = 0;
 			beentry->st_query_id = UINT64CONST(0);
+			beentry->st_plan_id = UINT64CONST(0);
 			proc->wait_event_info = 0;
 			PGSTAT_END_WRITE_ACTIVITY(beentry);
 		}
@@ -593,7 +595,10 @@ pgstat_report_activity(BackendState state, const char *cmd_str)
 	 * identifier.
 	 */
 	if (state == STATE_RUNNING)
+	{
 		beentry->st_query_id = UINT64CONST(0);
+		beentry->st_plan_id = UINT64CONST(0);
+	}
 
 	if (cmd_str != NULL)
 	{
@@ -644,6 +649,32 @@ pgstat_report_query_id(uint64 query_id, bool force)
 	PGSTAT_END_WRITE_ACTIVITY(beentry);
 }
 
+/* --------
+ * pgstat_report_plan_id() -
+ *
+ * Called to update top-level plan identifier.
+ *
+ * For the same reasons as the query identifiers (see above), we report only
+ * the top-level plan identifiers. Currently, only plugins update a plan
+ * identifier via the standard_planner or ExecutorStart hooks.
+ *
+ * --------
+ */
+void
+pgstat_report_plan_id(uint64 plan_id, bool force)
+{
+	volatile PgBackendStatus *beentry = MyBEEntry;
+
+	if (!beentry || !pgstat_track_activities)
+		return;
+
+	if (beentry->st_plan_id != 0 && !force)
+		return;
+
+	PGSTAT_BEGIN_WRITE_ACTIVITY(beentry);
+	beentry->st_plan_id = plan_id;
+	PGSTAT_END_WRITE_ACTIVITY(beentry);
+}
 
 /* ----------
  * pgstat_report_appname() -
@@ -1040,6 +1071,21 @@ pgstat_get_my_query_id(void)
 	return MyBEEntry->st_query_id;
 }
 
+/* ----------
+ * pgstat_get_my_plan_id() -
+ *
+ * Similar to pgstat_get_my_query_id, see above. Currently, only plugins
+ * have a use for this routine.
+ */
+uint64
+pgstat_get_my_plan_id(void)
+{
+	if (!MyBEEntry)
+		return 0;
+
+	return MyBEEntry->st_plan_id;
+}
+
 /* ----------
  * pgstat_get_backend_type_by_proc_number() -
  *
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index bf1f25c0db..1f84c63fcc 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -55,6 +55,9 @@ typedef struct PlannedStmt
 	/* query identifier (copied from Query) */
 	uint64		queryId;
 
+	/* plan identifier (set by plugins) */
+	uint64		planId;
+
 	/* is it insert|update|delete|merge RETURNING? */
 	bool		hasReturning;
 
diff --git a/src/include/utils/backend_status.h b/src/include/utils/backend_status.h
index d3d4ff6c5c..b2be0ad774 100644
--- a/src/include/utils/backend_status.h
+++ b/src/include/utils/backend_status.h
@@ -170,6 +170,9 @@ typedef struct PgBackendStatus
 
 	/* query identifier, optionally computed using post_parse_analyze_hook */
 	uint64		st_query_id;
+
+	/* plan identifier, optionally computed by plugins */
+	uint64		st_plan_id;
 } PgBackendStatus;
 
 
@@ -316,6 +319,7 @@ extern void pgstat_clear_backend_activity_snapshot(void);
 /* Activity reporting functions */
 extern void pgstat_report_activity(BackendState state, const char *cmd_str);
 extern void pgstat_report_query_id(uint64 query_id, bool force);
+extern void pgstat_report_plan_id(uint64 plan_id, bool force);
 extern void pgstat_report_tempfile(size_t filesize);
 extern void pgstat_report_appname(const char *appname);
 extern void pgstat_report_xact_timestamp(TimestampTz tstamp);
@@ -323,6 +327,7 @@ extern const char *pgstat_get_backend_current_activity(int pid, bool checkUser);
 extern const char *pgstat_get_crashed_backend_activity(int pid, char *buffer,
 													   int buflen);
 extern uint64 pgstat_get_my_query_id(void);
+extern uint64 pgstat_get_my_plan_id(void);
 extern BackendType pgstat_get_backend_type_by_proc_number(ProcNumber procNumber);
 
 
-- 
2.47.1