From a97b1bf327ad665c9f71c43063a3a9e6d364716d Mon Sep 17 00:00:00 2001
From: Hou Zhijie <HouZhijie@foxmail.com>
Date: Fri, 30 Jul 2021 10:04:32 +0800
Subject: [PATCH] cache-parallel-dml-safety

The planner is updated to perform additional parallel-safety checks For
non-partitioned table if pg_class.relparalleldml is DEFAULT('d'), and cache the
parallel safety for the relation.

Whenever any function's parallel-safety is changed, invalidate the cached
parallel-safety for all relations in relcache for a particular database.

For partitioned table, If pg_class.relparalleldml is DEFAULT('d'), assume that
the table is UNSAFE to be modified in parallel mode.

If pg_class.relparalleldml is SAFE/RESTRICTED/UNSAFE, respect the specified
parallel dml safety instead of checking it again.

---
 src/backend/catalog/pg_proc.c        | 13 +++++
 src/backend/commands/functioncmds.c  | 18 ++++++-
 src/backend/optimizer/util/clauses.c | 78 ++++++++++++++++++++++------
 src/backend/utils/cache/inval.c      | 53 +++++++++++++++++++
 src/backend/utils/cache/relcache.c   | 19 +++++++
 src/include/storage/sinval.h         |  8 +++
 src/include/utils/inval.h            |  2 +
 src/include/utils/rel.h              |  1 +
 src/include/utils/relcache.h         |  2 +
 9 files changed, 176 insertions(+), 18 deletions(-)

diff --git a/src/backend/catalog/pg_proc.c b/src/backend/catalog/pg_proc.c
index 1454d2fb67..04585dc3ef 100644
--- a/src/backend/catalog/pg_proc.c
+++ b/src/backend/catalog/pg_proc.c
@@ -39,6 +39,7 @@
 #include "tcop/tcopprot.h"
 #include "utils/acl.h"
 #include "utils/builtins.h"
+#include "utils/inval.h"
 #include "utils/lsyscache.h"
 #include "utils/regproc.h"
 #include "utils/rel.h"
@@ -367,6 +368,9 @@ ProcedureCreate(const char *procedureName,
 		Datum		proargnames;
 		bool		isnull;
 		const char *dropcmd;
+		char		old_proparallel;
+
+		old_proparallel = oldproc->proparallel;
 
 		if (!replace)
 			ereport(ERROR,
@@ -559,6 +563,15 @@ ProcedureCreate(const char *procedureName,
 		tup = heap_modify_tuple(oldtup, tupDesc, values, nulls, replaces);
 		CatalogTupleUpdate(rel, &tup->t_self, tup);
 
+		/*
+		 * If the function's parallel safety changed, the tables that depend
+		 * on this function won't be safe to be modified in parallel mode
+		 * anymore. So, we need to invalidate the parallel dml flag in
+		 * relcache.
+		 */
+		if (old_proparallel != parallel)
+			CacheInvalidateParallelDML();
+
 		ReleaseSysCache(oldtup);
 		is_update = true;
 	}
diff --git a/src/backend/commands/functioncmds.c b/src/backend/commands/functioncmds.c
index 79d875ab10..57d9ca52e5 100644
--- a/src/backend/commands/functioncmds.c
+++ b/src/backend/commands/functioncmds.c
@@ -70,6 +70,7 @@
 #include "utils/builtins.h"
 #include "utils/fmgroids.h"
 #include "utils/guc.h"
+#include "utils/inval.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
 #include "utils/rel.h"
@@ -1504,7 +1505,22 @@ AlterFunction(ParseState *pstate, AlterFunctionStmt *stmt)
 								repl_val, repl_null, repl_repl);
 	}
 	if (parallel_item)
-		procForm->proparallel = interpret_func_parallel(parallel_item);
+	{
+		char proparallel;
+
+		proparallel = interpret_func_parallel(parallel_item);
+
+		/*
+		 * If the function's parallel safety changed, the tables that depends
+		 * on this function won't be safe to be modified in parallel mode
+		 * anymore. So, we need to invalidate the parallel dml flag in
+		 * relcache.
+		 */
+		if (proparallel != procForm->proparallel)
+			CacheInvalidateParallelDML();
+
+		procForm->proparallel = proparallel;
+	}
 
 	/* Do the update */
 	CatalogTupleUpdate(rel, &tup->t_self, tup);
diff --git a/src/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index 749cb0dacd..f65c2fc961 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -187,7 +187,7 @@ static Node *substitute_actual_srf_parameters_mutator(Node *node,
 													  substitute_actual_srf_parameters_context *context);
 static bool max_parallel_hazard_test(char proparallel, max_parallel_hazard_context *context);
 static safety_object *make_safety_object(Oid objid, Oid classid, char proparallel);
-
+static char max_parallel_dml_hazard(Query *parse, max_parallel_hazard_context *context);
 
 /*****************************************************************************
  *		Aggregate-function clause manipulation
@@ -654,7 +654,6 @@ contain_volatile_functions_not_nextval_walker(Node *node, void *context)
 char
 max_parallel_hazard(Query *parse)
 {
-	bool max_hazard_found;
 	max_parallel_hazard_context context;
 
 	context.max_hazard = PROPARALLEL_SAFE;
@@ -664,28 +663,73 @@ max_parallel_hazard(Query *parse)
 	context.objects = NIL;
 	context.partition_directory = NULL;
 
-	max_hazard_found = max_parallel_hazard_walker((Node *) parse, &context);
+	if(!max_parallel_hazard_walker((Node *) parse, &context))
+		(void) max_parallel_dml_hazard(parse, &context);
+
+	return context.max_hazard;
+}
+
+/* Check the safety of parallel data modification */
+static char
+max_parallel_dml_hazard(Query *parse,
+						max_parallel_hazard_context *context)
+{
+	RangeTblEntry  *rte;
+	Relation		target_rel;
+	char			hazard;
+
+	if (!IsModifySupportedInParallelMode(parse->commandType))
+		return context->max_hazard;
+
+	/*
+	 * The target table is already locked by the caller (this is done in the
+	 * parse/analyze phase), and remains locked until end-of-transaction.
+	 */
+	rte = rt_fetch(parse->resultRelation, parse->rtable);
+	target_rel = table_open(rte->relid, NoLock);
+
+	/*
+	 * If user set specific parallel dml safety safe/restricted/unsafe, we
+	 * respect what user has set. If not set, for non-partitioned table, check
+	 * the safety automatically, for partitioned table, consider it as unsafe.
+	 */
+	hazard = target_rel->rd_rel->relparalleldml;
+	if (target_rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE &&
+		hazard == PROPARALLEL_DEFAULT)
+		hazard = PROPARALLEL_UNSAFE;
+
+	if (hazard != PROPARALLEL_DEFAULT)
+		(void) max_parallel_hazard_test(hazard, context);
 
-	if (!max_hazard_found &&
-		IsModifySupportedInParallelMode(parse->commandType))
+	/* Do parallel safety check for the target relation */
+	else if (!target_rel->rd_paralleldml)
 	{
-		RangeTblEntry *rte;
-		Relation target_rel;
+		bool max_hazard_found;
+		char pre_max_hazard = context->max_hazard;
+		context->max_hazard = PROPARALLEL_SAFE;
 
-		rte = rt_fetch(parse->resultRelation, parse->rtable);
+		max_hazard_found = target_rel_parallel_hazard_recurse(target_rel,
+															  context,
+															  false,
+															  false);
 
-		/*
-		 * The target table is already locked by the caller (this is done in the
-		 * parse/analyze phase), and remains locked until end-of-transaction.
-		 */
-		target_rel = table_open(rte->relid, NoLock);
+		/* Cache the parallel dml safety of this relation */
+		target_rel->rd_paralleldml = context->max_hazard;
 
-		(void) max_parallel_hazard_test(target_rel->rd_rel->relparalleldml,
-										&context);
-		table_close(target_rel, NoLock);
+		if (!max_hazard_found)
+			(void) max_parallel_hazard_test(pre_max_hazard, context);
 	}
 
-	return context.max_hazard;
+	/*
+	 * If we already cached the parallel dml safety of this relation, we don't
+	 * need to check it again.
+	 */
+	else
+		(void) max_parallel_hazard_test(target_rel->rd_paralleldml, context);
+
+	table_close(target_rel, NoLock);
+
+	return context->max_hazard;
 }
 
 /*
diff --git a/src/backend/utils/cache/inval.c b/src/backend/utils/cache/inval.c
index 9c79775725..9459b3c204 100644
--- a/src/backend/utils/cache/inval.c
+++ b/src/backend/utils/cache/inval.c
@@ -385,6 +385,27 @@ AddCatalogInvalidationMessage(InvalidationListHeader *hdr,
 	AddInvalidationMessage(&hdr->cclist, &msg);
 }
 
+/*
+ * Add a Parallel dml inval entry
+ */
+static void
+AddParallelDMLInvalidationMessage(InvalidationListHeader *hdr)
+{
+	SharedInvalidationMessage msg;
+
+	/* Don't add a duplicate item. */
+	ProcessMessageList(hdr->rclist,
+					   if (msg->rc.id == SHAREDINVALPARALLELDML_ID)
+					   return);
+
+	/* OK, add the item */
+	msg.pd.id = SHAREDINVALPARALLELDML_ID;
+	/* check AddCatcacheInvalidationMessage() for an explanation */
+	VALGRIND_MAKE_MEM_DEFINED(&msg, sizeof(msg));
+
+	AddInvalidationMessage(&hdr->rclist, &msg);
+}
+
 /*
  * Add a relcache inval entry
  */
@@ -539,6 +560,21 @@ RegisterRelcacheInvalidation(Oid dbId, Oid relId)
 		transInvalInfo->RelcacheInitFileInval = true;
 }
 
+/*
+ * RegisterParallelDMLInvalidation
+ *
+ * As above, but register a invalidation event for paralleldml in all relcache.
+ */
+static void
+RegisterParallelDMLInvalidation()
+{
+	AddParallelDMLInvalidationMessage(&transInvalInfo->CurrentCmdInvalidMsgs);
+
+	(void) GetCurrentCommandId(true);
+
+	transInvalInfo->RelcacheInitFileInval = true;
+}
+
 /*
  * RegisterSnapshotInvalidation
  *
@@ -631,6 +667,11 @@ LocalExecuteInvalidationMessage(SharedInvalidationMessage *msg)
 		else if (msg->sn.dbId == MyDatabaseId)
 			InvalidateCatalogSnapshot();
 	}
+	else if (msg->id == SHAREDINVALPARALLELDML_ID)
+	{
+		/* Invalid all the relcache's parallel dml flag */
+		ParallelDMLInvalidate();
+	}
 	else
 		elog(FATAL, "unrecognized SI message ID: %d", msg->id);
 }
@@ -1307,6 +1348,18 @@ CacheInvalidateRelcacheAll(void)
 	RegisterRelcacheInvalidation(InvalidOid, InvalidOid);
 }
 
+/*
+ * CacheInvalidateParallelDML
+ *		Register invalidation of the whole relcache at the end of command.
+ */
+void
+CacheInvalidateParallelDML(void)
+{
+	PrepareInvalidationState();
+
+	RegisterParallelDMLInvalidation();
+}
+
 /*
  * CacheInvalidateRelcacheByTuple
  *		As above, but relation is identified by passing its pg_class tuple.
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 3f38a69687..e013c4d0dc 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -2934,6 +2934,25 @@ RelationCacheInvalidate(void)
 	list_free(rebuildList);
 }
 
+/*
+ * ParallelDMLInvalidate
+ *	 Invalidate all the relcache's parallel dml flag.
+ */
+void
+ParallelDMLInvalidate(void)
+{
+	HASH_SEQ_STATUS status;
+	RelIdCacheEnt *idhentry;
+	Relation	relation;
+
+	hash_seq_init(&status, RelationIdCache);
+
+	while ((idhentry = (RelIdCacheEnt *) hash_seq_search(&status)) != NULL)
+	{
+		relation = idhentry->reldesc;
+		relation->rd_paralleldml = 0;
+	}
+}
 /*
  * RelationCloseSmgrByOid - close a relcache entry's smgr link
  *
diff --git a/src/include/storage/sinval.h b/src/include/storage/sinval.h
index f03dc23b14..9859a3bea0 100644
--- a/src/include/storage/sinval.h
+++ b/src/include/storage/sinval.h
@@ -110,6 +110,13 @@ typedef struct
 	Oid			relId;			/* relation ID */
 } SharedInvalSnapshotMsg;
 
+#define SHAREDINVALPARALLELDML_ID	(-6)
+
+typedef struct
+{
+	int8		id;				/* type field --- must be first */
+} SharedInvalParallelDMLMsg;
+
 typedef union
 {
 	int8		id;				/* type field --- must be first */
@@ -119,6 +126,7 @@ typedef union
 	SharedInvalSmgrMsg sm;
 	SharedInvalRelmapMsg rm;
 	SharedInvalSnapshotMsg sn;
+	SharedInvalParallelDMLMsg pd;
 } SharedInvalidationMessage;
 
 
diff --git a/src/include/utils/inval.h b/src/include/utils/inval.h
index 770672890b..f1ce1462c1 100644
--- a/src/include/utils/inval.h
+++ b/src/include/utils/inval.h
@@ -64,4 +64,6 @@ extern void CallSyscacheCallbacks(int cacheid, uint32 hashvalue);
 extern void InvalidateSystemCaches(void);
 
 extern void LogLogicalInvalidations(void);
+
+extern void CacheInvalidateParallelDML(void);
 #endif							/* INVAL_H */
diff --git a/src/include/utils/rel.h b/src/include/utils/rel.h
index b4faa1c123..52574e9d40 100644
--- a/src/include/utils/rel.h
+++ b/src/include/utils/rel.h
@@ -63,6 +63,7 @@ typedef struct RelationData
 	bool		rd_indexvalid;	/* is rd_indexlist valid? (also rd_pkindex and
 								 * rd_replidindex) */
 	bool		rd_statvalid;	/* is rd_statlist valid? */
+	char		rd_paralleldml; /* parallel dml safety */
 
 	/*----------
 	 * rd_createSubid is the ID of the highest subtransaction the rel has
diff --git a/src/include/utils/relcache.h b/src/include/utils/relcache.h
index 5ea225ac2d..5813aa50a0 100644
--- a/src/include/utils/relcache.h
+++ b/src/include/utils/relcache.h
@@ -128,6 +128,8 @@ extern void RelationCacheInvalidate(void);
 
 extern void RelationCloseSmgrByOid(Oid relationId);
 
+extern void ParallelDMLInvalidate(void);
+
 #ifdef USE_ASSERT_CHECKING
 extern void AssertPendingSyncs_RelationCache(void);
 #else
-- 
2.27.0

