From bf8c93b1855a5983e94c8830424b6f814c6b513e Mon Sep 17 00:00:00 2001
From: Diego Fronza <diego.fronza@percona.com>
Date: Fri, 15 Oct 2021 11:56:28 -0300
Subject: [PATCH] PG-261: Fix regression tests on distribution packages.

The code added in pgss_store() to handle an assertion failure when
GetUserId() was being called introduced a problem with regression tests
in some builds, specifically our PG11 and PG12 distributions for Ubuntu.

The problem was detected when calling some json functions that would
trigger parallel workers, to solve the problem now we check if our hooks
are being called by parallel workers, in this case we can avoid doing
work, it also fixes the issue mentioned above as we won't call
GetUserId() anymore in an invalid context.
---
 pg_stat_monitor.c | 26 ++++++++++++++++++--------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/pg_stat_monitor.c b/pg_stat_monitor.c
index 2f51c08..a045cc4 100644
--- a/pg_stat_monitor.c
+++ b/pg_stat_monitor.c
@@ -14,7 +14,9 @@
  *
  *-------------------------------------------------------------------------
  */
+
 #include "postgres.h"
+#include "access/parallel.h"
 #include <regex.h>
 #ifdef BENCHMARK
 #include <time.h> /* clock() */
@@ -367,6 +369,9 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
 	if (!IsSystemInitialized())
 		return;
 
+	if (IsParallelWorker())
+		return;
+
 	/*
 	 * Clear queryId for prepared statements related utility, as those will
 	 * inherit from the underlying statement's one (except DEALLOCATE which is
@@ -423,6 +428,9 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query)
 	if (!IsSystemInitialized())
 		return;
 
+	if (IsParallelWorker())
+		return;
+
 	/*
 	 * Utility statements get queryId zero.  We do this even in cases where
 	 * the statement contains an optimizable statement for which a queryId
@@ -485,6 +493,9 @@ pgss_ExecutorStart(QueryDesc *queryDesc, int eflags)
 	else
 		standard_ExecutorStart(queryDesc, eflags);
 
+	if (IsParallelWorker())
+		return;
+
 	/*
 	 * If query has queryId zero, don't track it.  This prevents double
 	 * counting of optimizable statements that are directly contained in
@@ -655,7 +666,7 @@ pgss_ExecutorEnd(QueryDesc *queryDesc)
 		MemoryContextSwitchTo(mct);
 	}
 
-	if (queryId != UINT64CONST(0) && queryDesc->totaltime)
+	if (queryId != UINT64CONST(0) && queryDesc->totaltime && !IsParallelWorker())
 	{
 		/*
 		 * Make sure stats accumulation is done.  (Note: it's okay if several
@@ -770,7 +781,7 @@ pgss_planner_hook(Query *parse, const char *query_string, int cursorOptions, Par
 {
 	PlannedStmt			*result;
 
-	if (PGSM_TRACK_PLANNING && query_string && parse->queryId != UINT64CONST(0))
+	if (PGSM_TRACK_PLANNING && query_string && parse->queryId != UINT64CONST(0) && !IsParallelWorker())
 	{
 		PlanInfo	plan_info;
 		instr_time	start;
@@ -942,7 +953,7 @@ static void pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	if (PGSM_TRACK_UTILITY &&
 		!IsA(parsetree, ExecuteStmt) &&
 		!IsA(parsetree, PrepareStmt) &&
-		!IsA(parsetree, DeallocateStmt))
+		!IsA(parsetree, DeallocateStmt) && !IsParallelWorker())
 	{
 		instr_time	start;
 		instr_time	duration;
@@ -1453,7 +1464,6 @@ pgss_store(uint64 queryid,
 	uint64          bucketid;
 	uint64          prev_bucket_id;
     uint64          userid;
-    int             con;
 	uint64          planid;
 	uint64          appid;
 	char            comments[512] = "";
@@ -1468,10 +1478,7 @@ pgss_store(uint64 queryid,
 		return;
 
 	Assert(query != NULL);
-	if (kind == PGSS_ERROR)
-		GetUserIdAndSecContext((unsigned int *)&userid, &con);
-	else
-		userid =  GetUserId();
+	userid =  GetUserId();
 
 	application_name_len = pg_get_application_name(application_name);
 	planid = plan_info ? plan_info->planid: 0;
@@ -3255,6 +3262,9 @@ pgsm_emit_log_hook(ErrorData *edata)
 	if (!IsSystemInitialized() || edata == NULL)
 		goto exit;
 
+	if (IsParallelWorker())
+		return;
+
 	if ((edata->elevel == ERROR || edata->elevel == WARNING || edata->elevel == INFO || edata->elevel == DEBUG1))
 	{
 		uint64 queryid = 0;