From e84e68de4054ee88f44a65213d1d5d17c6479a37 Mon Sep 17 00:00:00 2001
From: "houzj.fnst" <houzj.fnst@cn.fujitsu.com>
Date: Wed, 28 Apr 2021 19:22:12 +0800
Subject: [PATCH] check-parallel-safety-in-fmgr_info_cxt_security

---
 src/backend/utils/fmgr/fmgr.c | 62 ++++++++++++++++++++++++++++++++---
 1 file changed, 57 insertions(+), 5 deletions(-)

diff --git a/src/backend/utils/fmgr/fmgr.c b/src/backend/utils/fmgr/fmgr.c
index 3dfe6e5825..79ed3951b2 100644
--- a/src/backend/utils/fmgr/fmgr.c
+++ b/src/backend/utils/fmgr/fmgr.c
@@ -16,6 +16,8 @@
 #include "postgres.h"
 
 #include "access/detoast.h"
+#include "access/parallel.h"
+#include "access/xact.h"
 #include "catalog/pg_language.h"
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
@@ -52,6 +54,7 @@ typedef struct
 } CFuncHashTabEntry;
 
 static HTAB *CFuncHash = NULL;
+static bool safety_checking = false;
 
 
 static void fmgr_info_cxt_security(Oid functionId, FmgrInfo *finfo, MemoryContext mcxt,
@@ -152,6 +155,7 @@ fmgr_info_cxt_security(Oid functionId, FmgrInfo *finfo, MemoryContext mcxt,
 	Datum		prosrcdatum;
 	bool		isnull;
 	char	   *prosrc;
+	char		parallel_safety;
 
 	/*
 	 * fn_oid *must* be filled in last.  Some code assumes that if fn_oid is
@@ -163,6 +167,48 @@ fmgr_info_cxt_security(Oid functionId, FmgrInfo *finfo, MemoryContext mcxt,
 	finfo->fn_mcxt = mcxt;
 	finfo->fn_expr = NULL;		/* caller may set this later */
 
+	procedureTuple = NULL;
+	procedureStruct = NULL;
+
+	/* Do not need to check parallel safety in Init or Bootstrap mode */
+	if ((!safety_checking && IsNormalProcessingMode()) &&
+		(IsInParallelMode() || IsParallelWorker()))
+	{
+		/*
+		 * If cannot find functionId in syscache, it need use some built-in
+		 * function to scan the pg_proc.  Since functions used to scan systable
+		 * must be parallel safe, Setting the safety check flag here to skip
+		 * the recursive safety check.
+		 */
+		safety_checking = true;
+
+		PG_TRY();
+		{
+			procedureTuple = SearchSysCache1(PROCOID, ObjectIdGetDatum(functionId));
+		}
+		PG_FINALLY();
+		{
+			/* Reset the flag */
+			safety_checking = false;
+		}
+		PG_END_TRY();
+
+		if (!HeapTupleIsValid(procedureTuple))
+			elog(ERROR, "cache lookup failed for function %u", functionId);
+		procedureStruct = (Form_pg_proc) GETSTRUCT(procedureTuple);
+
+		parallel_safety = procedureStruct->proparallel;
+
+		/* Check the function's parallel safety */
+		if (((IsParallelWorker() &&
+			parallel_safety == PROPARALLEL_RESTRICTED) ||
+			parallel_safety == PROPARALLEL_UNSAFE))
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_TRANSACTION_STATE),
+					errmsg("parallel-safety execution violation of function \"%s\" (%c)",
+							get_func_name(functionId), parallel_safety)));
+	}
+
 	if ((fbp = fmgr_isbuiltin(functionId)) != NULL)
 	{
 		/*
@@ -174,14 +220,20 @@ fmgr_info_cxt_security(Oid functionId, FmgrInfo *finfo, MemoryContext mcxt,
 		finfo->fn_stats = TRACK_FUNC_ALL;	/* ie, never track */
 		finfo->fn_addr = fbp->func;
 		finfo->fn_oid = functionId;
+
+		if (procedureTuple != NULL)
+			ReleaseSysCache(procedureTuple);
+
 		return;
 	}
 
-	/* Otherwise we need the pg_proc entry */
-	procedureTuple = SearchSysCache1(PROCOID, ObjectIdGetDatum(functionId));
-	if (!HeapTupleIsValid(procedureTuple))
-		elog(ERROR, "cache lookup failed for function %u", functionId);
-	procedureStruct = (Form_pg_proc) GETSTRUCT(procedureTuple);
+	if (procedureStruct == NULL)
+	{
+		procedureTuple = SearchSysCache1(PROCOID, ObjectIdGetDatum(functionId));
+		if (!HeapTupleIsValid(procedureTuple))
+			elog(ERROR, "cache lookup failed for function %u", functionId);
+		procedureStruct = (Form_pg_proc) GETSTRUCT(procedureTuple);
+	}
 
 	finfo->fn_nargs = procedureStruct->pronargs;
 	finfo->fn_strict = procedureStruct->proisstrict;
-- 
2.18.4