Allow plugins to set a 64-bit plan identifier

There are extensions that wish to compute a plan identifier
to distinguish different plans that a particular query may use.
Currently, there is no way for such extensions to make this
plan identifier available to the core for reporting and other
potentially interesting use cases such as plan management. This
patch allows extensions to do this by providing a mechanism to
track this identifier in the backend_status as well as in
PlannedStmt. The reset of the plan identifier is controlled by
core and follows the same code path as the query identifier added
in 4f0b096.

Discussion: https://www.postgresql.org/message-id/flat/CAA5RZ0vyWd4r35uUBUmhngv8XqeiJUkJDDKkLf5LCoWxv-t_pw%40mail.gmail.com

Author: Sami Imseih (AWS)

src/backend/executor/execParallel.c

@@ -175,6 +175,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;

src/backend/optimizer/plan/planner.c

@@ -58,6 +58,7 @@
 #include "parser/parsetree.h"
 #include "partitioning/partdesc.h"
 #include "rewrite/rewriteManip.h"
+#include "utils/backend_status.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
 #include "utils/selfuncs.h"
@@ -291,6 +292,9 @@ planner(Query *parse, const char *query_string, int cursorOptions,
 		result = (*planner_hook) (parse, query_string, cursorOptions, boundParams);
 	else
 		result = standard_planner(parse, query_string, cursorOptions, boundParams);
+
+	pgstat_report_plan_id(result->planId, false);
+
 	return result;
 }
 

src/backend/tcop/postgres.c

@@ -1107,6 +1107,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
@@ -2030,6 +2031,18 @@ exec_bind_message(StringInfo input_message)
 					  cplan,
 					  psrc);
 
+	/* Portal is defined, set the plan ID based on its contents. */
+	foreach(lc, portal->stmts)
+	{
+		PlannedStmt *plan = lfirst_node(PlannedStmt, lc);
+
+		if (plan->planId != UINT64CONST(0))
+		{
+			pgstat_report_plan_id(plan->planId, false);
+			break;
+		}
+	}
+
 	/* Done with the snapshot used for parameter I/O and parsing/planning */
 	if (snapshot_set)
 		PopActiveSnapshot();
@@ -2170,6 +2183,17 @@ exec_execute_message(const char *portal_name, long max_rows)
 		}
 	}
 
+	foreach(lc, portal->stmts)
+	{
+		PlannedStmt *stmt = lfirst_node(PlannedStmt, lc);
+
+		if (stmt->planId != UINT64CONST(0))
+		{
+			pgstat_report_plan_id(stmt->planId, false);
+			break;
+		}
+	}
+
 	cmdtagname = GetCommandTagNameAndLen(portal->commandTag, &cmdtaglen);
 
 	set_ps_display_with_len(cmdtagname, cmdtaglen);

src/backend/utils/activity/backend_status.c

@@ -321,6 +321,7 @@ pgstat_bestart_initial(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
@@ -599,6 +600,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);
 		}
@@ -659,7 +661,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)
 	{
@@ -710,6 +715,44 @@ 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.
+ * --------
+ */
+void
+pgstat_report_plan_id(uint64 plan_id, bool force)
+{
+	volatile PgBackendStatus *beentry = MyBEEntry;
+
+	/*
+	 * if track_activities is disabled, st_plan_id should already have been
+	 * reset
+	 */
+	if (!beentry || !pgstat_track_activities)
+		return;
+
+	/*
+	 * We only report the top-level plan identifiers.  The stored plan_id is
+	 * reset when a backend calls pgstat_report_activity(STATE_RUNNING), or
+	 * with an explicit call to this function using the force flag.  If the
+	 * saved plan identifier is not zero it means that it's not a top-level
+	 * command, so ignore the one provided unless it's an explicit call to
+	 * reset the identifier.
+	 */
+	if (beentry->st_plan_id != 0 && !force)
+		return;
+
+	/*
+	 * Update my status entry, following the protocol of bumping
+	 * st_changecount before and after.  We use a volatile pointer here to
+	 * ensure the compiler doesn't try to get cute.
+	 */
+	PGSTAT_BEGIN_WRITE_ACTIVITY(beentry);
+	beentry->st_plan_id = plan_id;
+	PGSTAT_END_WRITE_ACTIVITY(beentry);
+}
 
 /* ----------
  * pgstat_report_appname() -
@@ -1106,6 +1149,21 @@ pgstat_get_my_query_id(void)
 	return MyBEEntry->st_query_id;
 }
 
+/* ----------
+ * pgstat_get_my_plan_id() -
+ *
+ * Return current backend's plan identifier.
+ */
+uint64
+pgstat_get_my_plan_id(void)
+{
+	if (!MyBEEntry)
+		return 0;
+
+	/* No need for a lock, for roughly the same reasons as above. */
+	return MyBEEntry->st_plan_id;
+}
+
 /* ----------
  * pgstat_get_backend_type_by_proc_number() -
  *

src/include/nodes/plannodes.h

@@ -55,6 +55,9 @@ typedef struct PlannedStmt
 	/* query identifier (copied from Query) */
 	uint64		queryId;
 
+	/* plan identifier (can be set by plugins) */
+	uint64		planId;
+
 	/* is it insert|update|delete|merge RETURNING? */
 	bool		hasReturning;
 

src/include/utils/backend_status.h

@@ -171,6 +171,9 @@ typedef struct PgBackendStatus
 
 	/* query identifier, optionally computed using post_parse_analyze_hook */
 	uint64		st_query_id;
+
+	/* plan identifier, optionally computed using planner_hook */
+	uint64		st_plan_id;
 } PgBackendStatus;
 
 
@@ -319,13 +322,15 @@ 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);
 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);