Streamlined Lock Distribution

* Packaged the logic in a single method which can be invoked by different operations which need to acquire distributed locks.
* Error checking, cleaned up some duplicated code, removed some redundant validation, fixed bug in reading nowait in `lock_relation_if_exists`

src/backend/distributed/commands/truncate.c

@@ -43,7 +43,6 @@
 static void ErrorIfUnsupportedTruncateStmt(TruncateStmt *truncateStatement);
 static void ExecuteTruncateStmtSequentialIfNecessary(TruncateStmt *command);
 static void EnsurePartitionTableNotReplicatedForTruncate(TruncateStmt *truncateStatement);
-static void LockTruncatedRelationMetadataInWorkers(TruncateStmt *truncateStatement);
 static List * TruncateTaskList(Oid relationId);
 
 
@@ -243,7 +242,8 @@ PreprocessTruncateStatement(TruncateStmt *truncateStatement)
 	ErrorIfUnsupportedTruncateStmt(truncateStatement);
 	EnsurePartitionTableNotReplicatedForTruncate(truncateStatement);
 	ExecuteTruncateStmtSequentialIfNecessary(truncateStatement);
-	LockTruncatedRelationMetadataInWorkers(truncateStatement);
+	AcquireDistributedLockOnRelations(truncateStatement->relations, AccessExclusiveLock,
+									  DIST_LOCK_REFERENCING_TABLES);
 }
 
 
@@ -340,65 +340,3 @@ ExecuteTruncateStmtSequentialIfNecessary(TruncateStmt *command)
 		}
 	}
 }
-
-
-/*
- * LockTruncatedRelationMetadataInWorkers determines if distributed
- * lock is necessary for truncated relations, and acquire locks.
- *
- * LockTruncatedRelationMetadataInWorkers handles distributed locking
- * of truncated tables before standard utility takes over.
- *
- * Actual distributed truncation occurs inside truncate trigger.
- *
- * This is only for distributed serialization of truncate commands.
- * The function assumes that there is no foreign key relation between
- * non-distributed and distributed relations.
- */
-static void
-LockTruncatedRelationMetadataInWorkers(TruncateStmt *truncateStatement)
-{
-	List *distributedRelationList = NIL;
-
-	/* nothing to do if there is no metadata at worker nodes */
-	if (!ClusterHasKnownMetadataWorkers())
-	{
-		return;
-	}
-
-	RangeVar *rangeVar = NULL;
-	foreach_ptr(rangeVar, truncateStatement->relations)
-	{
-		Oid relationId = RangeVarGetRelid(rangeVar, NoLock, false);
-		Oid referencingRelationId = InvalidOid;
-
-		if (!IsCitusTable(relationId))
-		{
-			continue;
-		}
-
-		if (list_member_oid(distributedRelationList, relationId))
-		{
-			continue;
-		}
-
-		distributedRelationList = lappend_oid(distributedRelationList, relationId);
-
-		CitusTableCacheEntry *cacheEntry = GetCitusTableCacheEntry(relationId);
-		Assert(cacheEntry != NULL);
-
-		List *referencingTableList = cacheEntry->referencingRelationsViaForeignKey;
-		foreach_oid(referencingRelationId, referencingTableList)
-		{
-			distributedRelationList = list_append_unique_oid(distributedRelationList,
-															 referencingRelationId);
-		}
-	}
-
-	if (distributedRelationList != NIL)
-	{
-		bool nowait = false;
-		AcquireDistributedLockOnRelations(distributedRelationList, AccessExclusiveLock,
-										  nowait);
-	}
-}

src/backend/distributed/commands/utility_hook.c

@@ -471,31 +471,13 @@ ProcessUtilityInternal(PlannedStmt *pstmt,
 	if (IsA(parsetree, LockStmt))
 	{
 		LockStmt *stmt = (LockStmt *) parsetree;
-		List *distributedRelationList = NIL;
 
-		RangeVar *rangeVar = NULL;
+		ereport(NOTICE, errmsg("Processing LOCK command."));
-		foreach_ptr(rangeVar, stmt->relations)
-		{
-			Oid relationId = RangeVarGetRelid(rangeVar, NoLock, false);
 
-			if (!IsCitusTable(relationId))
+		ErrorIfUnsupportedLockStmt(stmt);
-			{
+		uint32 nowaitFlag = stmt->nowait ? DIST_LOCK_NOWAIT : 0;
-				continue;
+		AcquireDistributedLockOnRelations(stmt->relations, stmt->mode,
-			}
+										  DIST_LOCK_VIEWS_RECUR | nowaitFlag);
-
-			if (list_member_oid(distributedRelationList, relationId))
-			{
-				continue;
-			}
-
-			distributedRelationList = lappend_oid(distributedRelationList, relationId);
-		}
-
-		if (distributedRelationList != NIL)
-		{
-			AcquireDistributedLockOnRelations(distributedRelationList, stmt->mode,
-											  stmt->nowait);
-		}
 	}
 
 	/*

src/backend/distributed/sql/citus--10.2-4--11.0-1.sql

@@ -105,3 +105,5 @@ GRANT SELECT ON pg_catalog.pg_dist_object TO public;
 #include "udfs/citus_pid_for_gpid/11.0-1.sql"
 
 #include "udfs/citus_coordinator_nodeid/11.0-1.sql"
+
+#include "udfs/lock_relation_if_exists/11.0-1.sql"

src/backend/distributed/sql/citus--11.0-2--11.1-1.sql

@@ -6,14 +6,3 @@ DROP FUNCTION pg_catalog.worker_hash_partition_table(bigint, integer, text, text
 DROP FUNCTION pg_catalog.worker_merge_files_into_table(bigint, integer, text[], text[]);
 DROP FUNCTION pg_catalog.worker_range_partition_table(bigint, integer, text, text, oid, anyarray);
 DROP FUNCTION pg_catalog.worker_repartition_cleanup(bigint);
-
-SET search_path = 'pg_catalog';
-
-CREATE OR REPLACE FUNCTION lock_relation_if_exists(table_name text, lock_mode text, nowait boolean)
-RETURNS BOOL
-LANGUAGE C STRICT as 'MODULE_PATHNAME',
-$$lock_relation_if_exists$$;
-COMMENT ON FUNCTION lock_relation_if_exists(table_name text, lock_mode text, nowait boolean)
-IS 'locks relation in the lock_mode if the relation exists';
-
-RESET search_path;

src/backend/distributed/sql/udfs/lock_relation_if_exists/11.0-1.sql

@@ -0,0 +1,6 @@
+CREATE OR REPLACE FUNCTION pg_catalog.lock_relation_if_exists(table_name text, lock_mode text, nowait boolean)
+RETURNS BOOL
+LANGUAGE C STRICT as 'MODULE_PATHNAME',
+$$lock_relation_if_exists$$;
+COMMENT ON FUNCTION pg_catalog.lock_relation_if_exists(table_name text, lock_mode text, nowait boolean)
+IS 'locks relation in the lock_mode if the relation exists';

src/backend/distributed/sql/udfs/lock_relation_if_exists/latest.sql

@@ -0,0 +1,6 @@
+CREATE OR REPLACE FUNCTION pg_catalog.lock_relation_if_exists(table_name text, lock_mode text, nowait boolean)
+RETURNS BOOL
+LANGUAGE C STRICT as 'MODULE_PATHNAME',
+$$lock_relation_if_exists$$;
+COMMENT ON FUNCTION pg_catalog.lock_relation_if_exists(table_name text, lock_mode text, nowait boolean)
+IS 'locks relation in the lock_mode if the relation exists';

src/backend/distributed/utils/resource_lock.c

@@ -36,6 +36,7 @@
 #include "distributed/resource_lock.h"
 #include "distributed/shardinterval_utils.h"
 #include "distributed/worker_protocol.h"
+#include "distributed/worker_shard_visibility.h"
 #include "distributed/utils/array_type.h"
 #include "distributed/version_compat.h"
 #include "storage/lmgr.h"
@@ -1006,7 +1007,7 @@ lock_relation_if_exists(PG_FUNCTION_ARGS)
 	char *lockModeCString = text_to_cstring(lockModeText);
 
 	bool nowait = false;
-	if (!PG_ARGISNULL(2))
+	if (PG_NARGS() == 3)
 	{
 		nowait = PG_GETARG_BOOL(2);
 	}
@@ -1088,32 +1089,31 @@ CitusLockTableAclCheck(Oid relationId, LOCKMODE lockmode, Oid userId)
 
 
 /*
- * AcquireDistributedLockOnRelations acquire a distributed lock on worker nodes
+ * AcquireDistributedLockOnRelations_Internal acquire a distributed lock on worker nodes
- * for given list of relations ids. Relation id list and worker node list
+ * for given list of relations ids. Worker node list is sorted so that the lock
- * sorted so that the lock is acquired in the same order regardless of which
+ * is acquired in the same order regardless of which node it was run on.
- * node it was run on. Notice that no lock is acquired on coordinator node.
+ * Notice that no lock is acquired on coordinator node if the coordinator is not
+ * added to the metadata.
  *
- * Notice that the locking functions is sent to all workers regardless of if
+ * Notice that no validation is done on the relationIds, that is the responsibility of
- * it has metadata or not. This is because a worker node only knows itself
+ * AcquireDistributedLockOnRelations.
- * and previous workers that has metadata sync turned on. The node does not
- * know about other nodes that have metadata sync turned on afterwards.
  *
  * A nowait flag is used to require the locks to be available immediately
  * and if that is not the case, an error will be thrown
  */
-void
+static void
-AcquireDistributedLockOnRelations(List *relationIdList, LOCKMODE lockMode, bool nowait)
+AcquireDistributedLockOnRelations_Internal(List *relationIdList, LOCKMODE lockMode, bool
+										   nowait)
 {
 	Oid relationId = InvalidOid;
 	List *workerNodeList = ActivePrimaryNodeList(NoLock);
+	Oid userId = GetUserId();
 
 	const char *lockModeCString = LockModeToLockModeCString(lockMode);
 
 	/*
 	 * We want to acquire locks in the same order across the nodes.
-	 * Although relation ids may change, their ordering will not.
 	 */
-	relationIdList = SortList(relationIdList, CompareOids);
 	workerNodeList = SortList(workerNodeList, CompareWorkerNodes);
 
 	UseCoordinatedTransaction();
@@ -1121,16 +1121,6 @@ AcquireDistributedLockOnRelations(List *relationIdList, LOCKMODE lockMode, bool
 	int32 localGroupId = GetLocalGroupId();
 
 	foreach_oid(relationId, relationIdList)
-	{
-		/*
-		 * We only acquire distributed lock on relation if
-		 * the relation is sync'ed between mx nodes.
-		 *
-		 * Even if users disable metadata sync, we cannot
-		 * allow them not to acquire the remote locks.
-		 * Hence, we have !IsCoordinator() check.
-		 */
-		if (ShouldSyncTableMetadata(relationId) || !IsCoordinator())
 	{
 		char *qualifiedRelationName = generate_qualified_relation_name(relationId);
 		StringInfo lockRelationCommand = makeStringInfo();
@@ -1141,6 +1131,10 @@ AcquireDistributedLockOnRelations(List *relationIdList, LOCKMODE lockMode, bool
 						 quote_literal_cstr(qualifiedRelationName),
 						 lockModeCString);
 
+		/* preemptive permission check so a worker connection is not */
+		/* established if the user is not allowed to acquire the lock */
+		CitusLockTableAclCheck(relationId, lockMode, userId);
+
 		WorkerNode *workerNode = NULL;
 		foreach_ptr(workerNode, workerNodeList)
 		{
@@ -1153,13 +1147,128 @@ AcquireDistributedLockOnRelations(List *relationIdList, LOCKMODE lockMode, bool
 				LockRelationIfExists(
 					cstring_to_text(qualifiedRelationName),
 					lockModeCString,
-						nowait
+					nowait);
-						);
+
 				continue;
 			}
 
 			SendCommandToWorker(nodeName, nodePort, lockRelationCommand->data);
 		}
 	}
+}
+
+
+/*
+ * AcquireDistributedLockOnRelations filters relations before passing them to
+ * AcquireDistributedLockOnRelations_Internal to acquire the locks.
+ *
+ * Only tables, views, and foreign tables can be locked with this function. Other relations
+ * will cause an error.
+ *
+ * Gracefully locks relations so no errors are thrown for things like invalid id-s
+ * or missing relations on worker nodes.
+ *
+ * Considers different types of relations based on a 'configs' parameter:
+ * - DIST_LOCK_DEFAULT: locks citus tables
+ * - DIST_LOCK_VIEWS_RECUR: locks citus tables that locked views are dependent on recursively
+ * - DIST_LOCK_REFERENCING_TABLES: locks tables that refer to locked citus tables with a foreign key
+ * - DIST_LOCK_NOWAIT: throws an error if the lock is not immediately available
+ */
+void
+AcquireDistributedLockOnRelations(List *relationList, LOCKMODE lockMode, uint32 configs)
+{
+	/* nothing to do if there is no metadata at worker nodes */
+	if (!ClusterHasKnownMetadataWorkers())
+	{
+		return;
+	}
+
+	List *distributedRelationList = NIL;
+
+	RangeVar *rangeVar = NULL;
+	List *relations = list_copy(relationList);
+	foreach_ptr(rangeVar, relations)
+	{
+		Oid relationId = RangeVarGetRelid(rangeVar, NoLock, false);
+
+		if ((configs & DIST_LOCK_VIEWS_RECUR) > 0 && get_rel_relkind(relationId) ==
+			RELKIND_VIEW)
+		{
+			ObjectAddress viewAddress = { 0 };
+			Oid schemaOid = RangeVarGetCreationNamespace(rangeVar);
+			ObjectAddressSet(viewAddress, schemaOid, relationId);
+
+			List *distDependencies = GetDistributableDependenciesForObject(&viewAddress);
+
+			ObjectAddress *address = NULL;
+			foreach_ptr(address, distDependencies)
+			{
+				distributedRelationList = list_append_unique_oid(distributedRelationList,
+																 address->objectId);
+			}
+
+			continue;
+		}
+
+		if (!IsCitusTable(relationId))
+		{
+			continue;
+		}
+
+		if (list_member_oid(distributedRelationList, relationId))
+		{
+			continue;
+		}
+
+		distributedRelationList = lappend_oid(distributedRelationList, relationId);
+
+		if ((configs & DIST_LOCK_REFERENCING_TABLES) > 0)
+		{
+			Oid referencingRelationId = InvalidOid;
+			CitusTableCacheEntry *cacheEntry = GetCitusTableCacheEntry(relationId);
+			Assert(cacheEntry != NULL);
+
+			List *referencingTableList = cacheEntry->referencingRelationsViaForeignKey;
+			foreach_oid(referencingRelationId, referencingTableList)
+			{
+				distributedRelationList = list_append_unique_oid(distributedRelationList,
+																 referencingRelationId);
+			}
+		}
+	}
+
+	if (distributedRelationList != NIL)
+	{
+		bool nowait = (configs & DIST_LOCK_NOWAIT) > 0;
+		AcquireDistributedLockOnRelations_Internal(distributedRelationList, lockMode,
+												   nowait);
+	}
+}
+
+
+/*
+ * ErrorIfUnsupportedLockStmt errors out if:
+ * - The lock statement is not in a transaction block
+ * - The relation id-s being locked do not exist
+ * - Locking shard, but citus.enable_manual_changes_to_shards is false
+ */
+void
+ErrorIfUnsupportedLockStmt(LockStmt *stmt)
+{
+	RequireTransactionBlock(true, "LOCK TABLE");
+
+	RangeVar *rangeVar = NULL;
+	foreach_ptr(rangeVar, stmt->relations)
+	{
+		bool missingOk = false;
+		Oid relationId = RangeVarGetRelid(rangeVar, NoLock, missingOk);
+
+		/* Note that allowing the user to lock shards could lead to */
+		/* distributed deadlocks due to shards not being locked when */
+		/* a distributed table is locked. */
+		/* However, because citus.enable_manual_changes_to_shards */
+		/* is a guc which is not visible by default, whoever is using this */
+		/* guc will hopefully know what they're doing and avoid such scenarios. */
+		ErrorIfIllegallyChangingKnownShard(relationId);
 	}
 }

src/include/distributed/listutils.h

@@ -13,6 +13,7 @@
 #define CITUS_LISTUTILS_H
 
 #include "postgres.h"
+#include "storage/lock.h"
 #include "c.h"
 
 #include "nodes/pg_list.h"

src/include/distributed/resource_lock.h

@@ -16,6 +16,7 @@
 #include "distributed/worker_transaction.h"
 #include "nodes/pg_list.h"
 #include "storage/lock.h"
+#include "tcop/utility.h"
 
 
 /*
@@ -110,9 +111,14 @@ typedef enum CitusOperations
 						 ADV_LOCKTAG_CLASS_CITUS_PLACEMENT_CLEANUP)
 
 
-#define LOCK_RELATION_IF_EXISTS "SELECT lock_relation_if_exists(%s, '%s');"
+#define DIST_LOCK_DEFAULT 0
+#define DIST_LOCK_VIEWS_RECUR 1 << 0
+#define DIST_LOCK_REFERENCING_TABLES 1 << 1
+#define DIST_LOCK_NOWAIT 1 << 2
+
+#define LOCK_RELATION_IF_EXISTS "SELECT pg_catalog.lock_relation_if_exists(%s, '%s');"
 #define LOCK_RELATION_IF_EXISTS_NOWAIT \
-	"SELECT lock_relation_if_exists(%s, '%s', nowait => true);"
+	"SELECT pg_catalog.lock_relation_if_exists(%s, '%s', nowait => true);"
 
 /* Lock shard/relation metadata for safe modifications */
 extern void LockShardDistributionMetadata(int64 shardId, LOCKMODE lockMode);
@@ -155,6 +161,7 @@ extern void LockParentShardResourceIfPartition(List *shardIntervalList,
 /* Lock mode translation between text and enum */
 extern LOCKMODE LockModeCStringToLockMode(const char *lockModeName);
 extern const char * LockModeToLockModeCString(LOCKMODE lockMode);
-extern void AcquireDistributedLockOnRelations(List *relationIdList, LOCKMODE lockMode,
+extern void AcquireDistributedLockOnRelations(List *relationList, LOCKMODE lockMode,
-											  bool nowait);
+											  uint32 configs);
+extern void ErrorIfUnsupportedLockStmt(LockStmt *stmt);
 #endif /* RESOURCE_LOCK_H */

src/test/regress/expected/mx_regular_user.out

@@ -89,7 +89,6 @@ INSERT INTO super_user_owned_regular_user_granted VALUES (1, 1), (2, 1) ON CONFL
 ERROR:  permission denied for table super_user_owned_regular_user_granted
 TRUNCATE super_user_owned_regular_user_granted;
 ERROR:  permission denied for table super_user_owned_regular_user_granted
-CONTEXT:  while executing command on localhost:xxxxx
 DELETE FROM super_user_owned_regular_user_granted;
 ERROR:  permission denied for table super_user_owned_regular_user_granted
 UPDATE super_user_owned_regular_user_granted SET a = 1;