From 602e175eb0b0d7338948caa8f92f69cd4a8fa0c7 Mon Sep 17 00:00:00 2001
From: "houzj.fnst" <houzj.fnst@cn.fujitsu.com>
Date: Thu, 29 Apr 2021 08:53:26 +0800
Subject: [PATCH] approach 2 check-parallel-safety-in-fmgr_info_cxt_security

---
 src/backend/utils/Gen_fmgrtab.pl | 33 +++++++++++++++++++-------------
 src/backend/utils/fmgr/fmgr.c    | 27 ++++++++++++++++++++++++--
 src/include/utils/fmgrtab.h      |  8 ++++++--
 3 files changed, 51 insertions(+), 17 deletions(-)

diff --git a/src/backend/utils/Gen_fmgrtab.pl b/src/backend/utils/Gen_fmgrtab.pl
index 881568defd..6d95d2e4d4 100644
--- a/src/backend/utils/Gen_fmgrtab.pl
+++ b/src/backend/utils/Gen_fmgrtab.pl
@@ -72,15 +72,16 @@ foreach my $row (@{ $catalog_data{pg_proc} })
 
 	push @fmgr,
 	  {
-		oid    => $bki_values{oid},
-		name   => $bki_values{proname},
-		lang   => $bki_values{prolang},
-		kind   => $bki_values{prokind},
-		strict => $bki_values{proisstrict},
-		retset => $bki_values{proretset},
-		nargs  => $bki_values{pronargs},
-		args   => $bki_values{proargtypes},
-		prosrc => $bki_values{prosrc},
+		oid      => $bki_values{oid},
+		name     => $bki_values{proname},
+		lang     => $bki_values{prolang},
+		kind     => $bki_values{prokind},
+		strict   => $bki_values{proisstrict},
+		parallel => $bki_values{proparallel},
+		retset   => $bki_values{proretset},
+		nargs    => $bki_values{pronargs},
+		args     => $bki_values{proargtypes},
+		prosrc   => $bki_values{prosrc},
 	  };
 
 	# Count so that we can detect overloaded pronames.
@@ -208,9 +209,13 @@ foreach my $s (sort { $a->{oid} <=> $b->{oid} } @fmgr)
 
 # Create the fmgr_builtins table, collect data for fmgr_builtin_oid_index
 print $tfh "\nconst FmgrBuiltin fmgr_builtins[] = {\n";
-my %bmap;
-$bmap{'t'} = 'true';
-$bmap{'f'} = 'false';
+my %bmap_strict;
+$bmap_strict{'t'} = 1 << 0;
+$bmap_strict{'f'} = 0;
+my %bmap_retset;
+$bmap_retset{'t'} = 1 << 1;
+$bmap_retset{'f'} = 0;
+
 my @fmgr_builtin_oid_index;
 my $last_builtin_oid = 0;
 my $fmgr_count       = 0;
@@ -220,9 +225,11 @@ foreach my $s (sort { $a->{oid} <=> $b->{oid} } @fmgr)
 	# We do not need entries for aggregate functions
 	next if $s->{kind} eq 'a';
 
+	my $bitflag = $bmap_strict{$s->{strict}} | $bmap_retset{$s->{retset}};
 	print $tfh ",\n" if ($fmgr_count > 0);
 	print $tfh
-	  "  { $s->{oid}, $s->{nargs}, $bmap{$s->{strict}}, $bmap{$s->{retset}}, \"$s->{prosrc}\", $s->{prosrc} }";
+#	  "  { $s->{oid}, $s->{nargs}, $bmap{$s->{strict}}, $bmap{$s->{retset}}, \"$s->{prosrc}\", $s->{prosrc} }";
+	  "  { $s->{oid}, $s->{nargs}, $bitflag, \'$s->{parallel}\', \"$s->{prosrc}\", $s->{prosrc} }";
 
 	$fmgr_builtin_oid_index[ $s->{oid} ] = $fmgr_count++;
 	$last_builtin_oid = $s->{oid};
diff --git a/src/backend/utils/fmgr/fmgr.c b/src/backend/utils/fmgr/fmgr.c
index 3dfe6e5825..7f7286cb8b 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"
@@ -56,6 +58,7 @@ static HTAB *CFuncHash = NULL;
 
 static void fmgr_info_cxt_security(Oid functionId, FmgrInfo *finfo, MemoryContext mcxt,
 								   bool ignore_security);
+static void fmgr_check_parallel_safety(char parallel_safety, Oid functionId);
 static void fmgr_info_C_lang(Oid functionId, FmgrInfo *finfo, HeapTuple procedureTuple);
 static void fmgr_info_other_lang(Oid functionId, FmgrInfo *finfo, HeapTuple procedureTuple);
 static CFuncHashTabEntry *lookup_C_func(HeapTuple procedureTuple);
@@ -165,12 +168,15 @@ fmgr_info_cxt_security(Oid functionId, FmgrInfo *finfo, MemoryContext mcxt,
 
 	if ((fbp = fmgr_isbuiltin(functionId)) != NULL)
 	{
+		/* Check parallel safety for built-in functions */
+		fmgr_check_parallel_safety(fbp->parallel, functionId);
+
 		/*
 		 * Fast path for builtin functions: don't bother consulting pg_proc
 		 */
 		finfo->fn_nargs = fbp->nargs;
-		finfo->fn_strict = fbp->strict;
-		finfo->fn_retset = fbp->retset;
+		finfo->fn_strict = GETSTRICT(fbp);
+		finfo->fn_retset = GETRETSET(fbp);
 		finfo->fn_stats = TRACK_FUNC_ALL;	/* ie, never track */
 		finfo->fn_addr = fbp->func;
 		finfo->fn_oid = functionId;
@@ -183,6 +189,9 @@ fmgr_info_cxt_security(Oid functionId, FmgrInfo *finfo, MemoryContext mcxt,
 		elog(ERROR, "cache lookup failed for function %u", functionId);
 	procedureStruct = (Form_pg_proc) GETSTRUCT(procedureTuple);
 
+	/* Check parallel safety for other functions */
+	fmgr_check_parallel_safety(procedureStruct->proparallel, functionId);
+
 	finfo->fn_nargs = procedureStruct->pronargs;
 	finfo->fn_strict = procedureStruct->proisstrict;
 	finfo->fn_retset = procedureStruct->proretset;
@@ -264,6 +273,20 @@ fmgr_info_cxt_security(Oid functionId, FmgrInfo *finfo, MemoryContext mcxt,
 	ReleaseSysCache(procedureTuple);
 }
 
+static void
+fmgr_check_parallel_safety(char parallel_safety, Oid functionId)
+{
+	if (IsInParallelMode() &&
+	   ((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)));
+}
+
+
 /*
  * Return module and C function name providing implementation of functionId.
  *
diff --git a/src/include/utils/fmgrtab.h b/src/include/utils/fmgrtab.h
index 21a5f21156..beaf38db63 100644
--- a/src/include/utils/fmgrtab.h
+++ b/src/include/utils/fmgrtab.h
@@ -26,12 +26,16 @@ typedef struct
 {
 	Oid			foid;			/* OID of the function */
 	short		nargs;			/* 0..FUNC_MAX_ARGS, or -1 if variable count */
-	bool		strict;			/* T if function is "strict" */
-	bool		retset;			/* T if function returns a set */
+	char		bitflag;		/* 1 << 0 if function is "strict"
+								 * 1 << 1 if function returns a set */
+	char		parallel;
 	const char *funcName;		/* C name of the function */
 	PGFunction	func;			/* pointer to compiled function */
 } FmgrBuiltin;
 
+#define GETSTRICT(fbp) ((fbp->bitflag & (1 << 0)) ? true : false)
+#define GETRETSET(fbp) ((fbp->bitflag & (1 << 1)) ? true : false)
+
 extern const FmgrBuiltin fmgr_builtins[];
 
 extern const int fmgr_nbuiltins;	/* number of entries in table */
-- 
2.18.4

