From 3ad0eea74302f3660c53449ce88dcc041f9c5bbd Mon Sep 17 00:00:00 2001
From: Will Mortensen <will@extrahop.com>
Date: Thu, 21 Dec 2023 22:12:49 -0800
Subject: [PATCH v6 2/3] Allow specifying single lockmode in WaitForLockers()

Allow waiting for a single specified lockmode, rather than all lockmodes
that conflict with a specified mode.
---
 src/backend/catalog/index.c      |  4 ++--
 src/backend/commands/indexcmds.c | 12 ++++++------
 src/backend/commands/tablecmds.c |  3 ++-
 src/backend/storage/lmgr/lmgr.c  | 29 ++++++++++++++++-------------
 src/include/storage/lmgr.h       |  6 ++++--
 5 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/src/backend/catalog/index.c b/src/backend/catalog/index.c
index 88f7994b5a..3706e2df46 100644
--- a/src/backend/catalog/index.c
+++ b/src/backend/catalog/index.c
@@ -2290,7 +2290,7 @@ index_drop(Oid indexId, bool concurrent, bool concurrent_lock_mode)
 		 * here, even though it will only be used when we're called by REINDEX
 		 * CONCURRENTLY and not when called by DROP INDEX CONCURRENTLY.
 		 */
-		WaitForLockers(heaplocktag, AccessExclusiveLock, true);
+		WaitForLockers(heaplocktag, AccessExclusiveLock, true, true);
 
 		/* Finish invalidation of index and mark it as dead */
 		index_concurrently_set_dead(heapId, indexId);
@@ -2306,7 +2306,7 @@ index_drop(Oid indexId, bool concurrent, bool concurrent_lock_mode)
 		 * Wait till every transaction that saw the old index state has
 		 * finished.  See above about progress reporting.
 		 */
-		WaitForLockers(heaplocktag, AccessExclusiveLock, true);
+		WaitForLockers(heaplocktag, AccessExclusiveLock, true, true);
 
 		/*
 		 * Re-open relations to allow us to complete our actions.
diff --git a/src/backend/commands/indexcmds.c b/src/backend/commands/indexcmds.c
index 340248a3f2..bf518a2719 100644
--- a/src/backend/commands/indexcmds.c
+++ b/src/backend/commands/indexcmds.c
@@ -1660,7 +1660,7 @@ DefineIndex(Oid tableId,
 	 * exclusive lock on our table.  The lock code will detect deadlock and
 	 * error out properly.
 	 */
-	WaitForLockers(heaplocktag, ShareLock, true);
+	WaitForLockers(heaplocktag, ShareLock, true, true);
 
 	/*
 	 * At this moment we are sure that there are no transactions with the
@@ -1707,7 +1707,7 @@ DefineIndex(Oid tableId,
 	 */
 	pgstat_progress_update_param(PROGRESS_CREATEIDX_PHASE,
 								 PROGRESS_CREATEIDX_PHASE_WAIT_2);
-	WaitForLockers(heaplocktag, ShareLock, true);
+	WaitForLockers(heaplocktag, ShareLock, true, true);
 
 	/*
 	 * Now take the "reference snapshot" that will be used by validate_index()
@@ -3924,7 +3924,7 @@ ReindexRelationConcurrently(const ReindexStmt *stmt, Oid relationOid, const Rein
 
 	pgstat_progress_update_param(PROGRESS_CREATEIDX_PHASE,
 								 PROGRESS_CREATEIDX_PHASE_WAIT_1);
-	WaitForLockersMultiple(lockTags, ShareLock, true);
+	WaitForLockersMultiple(lockTags, ShareLock, true, true);
 	CommitTransactionCommand();
 
 	foreach(lc, newIndexIds)
@@ -3983,7 +3983,7 @@ ReindexRelationConcurrently(const ReindexStmt *stmt, Oid relationOid, const Rein
 
 	pgstat_progress_update_param(PROGRESS_CREATEIDX_PHASE,
 								 PROGRESS_CREATEIDX_PHASE_WAIT_2);
-	WaitForLockersMultiple(lockTags, ShareLock, true);
+	WaitForLockersMultiple(lockTags, ShareLock, true, true);
 	CommitTransactionCommand();
 
 	foreach(lc, newIndexIds)
@@ -4141,7 +4141,7 @@ ReindexRelationConcurrently(const ReindexStmt *stmt, Oid relationOid, const Rein
 
 	pgstat_progress_update_param(PROGRESS_CREATEIDX_PHASE,
 								 PROGRESS_CREATEIDX_PHASE_WAIT_4);
-	WaitForLockersMultiple(lockTags, AccessExclusiveLock, true);
+	WaitForLockersMultiple(lockTags, AccessExclusiveLock, true, true);
 
 	foreach(lc, indexIds)
 	{
@@ -4175,7 +4175,7 @@ ReindexRelationConcurrently(const ReindexStmt *stmt, Oid relationOid, const Rein
 
 	pgstat_progress_update_param(PROGRESS_CREATEIDX_PHASE,
 								 PROGRESS_CREATEIDX_PHASE_WAIT_5);
-	WaitForLockersMultiple(lockTags, AccessExclusiveLock, true);
+	WaitForLockersMultiple(lockTags, AccessExclusiveLock, true, true);
 
 	PushActiveSnapshot(GetTransactionSnapshot());
 
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 2822b2bb44..f2d510522e 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -19444,7 +19444,8 @@ ATExecDetachPartition(List **wqueue, AlteredTableInfo *tab, Relation rel,
 		 * partition itself, since we will acquire AccessExclusiveLock below.
 		 */
 		SET_LOCKTAG_RELATION(tag, MyDatabaseId, parentrelid);
-		WaitForLockersMultiple(list_make1(&tag), AccessExclusiveLock, false);
+		WaitForLockersMultiple(list_make1(&tag), AccessExclusiveLock, true,
+							   false);
 
 		/*
 		 * Now acquire locks in both relations again.  Note they may have been
diff --git a/src/backend/storage/lmgr/lmgr.c b/src/backend/storage/lmgr/lmgr.c
index 4975d4b67d..58dd6f256e 100644
--- a/src/backend/storage/lmgr/lmgr.c
+++ b/src/backend/storage/lmgr/lmgr.c
@@ -893,19 +893,20 @@ XactLockTableWaitErrorCb(void *arg)
 
 /*
  * WaitForLockersMultiple
- *		Wait until no transaction holds locks that conflict with the given
- *		locktags at the given lockmode.
+ *		Wait until no transaction holds locks on the given locktags, either in
+ *		or conflicting with the given lockmode, depending on the value of the
+ *		conflicting argument.
  *
  * To do this, obtain the current list of lockers, and wait on their VXIDs
  * until they are finished.
  *
  * Note we don't try to acquire the locks on the given locktags, only the
- * VXIDs and XIDs of their lock holders; if somebody grabs a conflicting lock
- * on the objects after we obtained our initial list of lockers, we will not
- * wait for them.
+ * VXIDs and XIDs of their lock holders; if somebody grabs a lock on the objects
+ * after we obtained our initial list of lockers, we will not wait for them.
  */
 void
-WaitForLockersMultiple(List *locktags, LOCKMODE lockmode, bool progress)
+WaitForLockersMultiple(List *locktags, LOCKMODE lockmode, bool conflicting,
+					   bool progress)
 {
 	List	   *holders = NIL;
 	ListCell   *lc;
@@ -920,11 +921,13 @@ WaitForLockersMultiple(List *locktags, LOCKMODE lockmode, bool progress)
 	foreach(lc, locktags)
 	{
 		LOCKTAG    *locktag = lfirst(lc);
+		LOCKMASK	lockmask = conflicting ?
+			GetLockConflictMask(locktag, lockmode) : LOCKBIT_ON(lockmode);
 		int			count;
 
 		holders = lappend(holders,
-						  GetLockConflicts(locktag, lockmode,
-										   progress ? &count : NULL));
+						  GetLockers(locktag, lockmask,
+									 progress ? &count : NULL));
 		if (progress)
 			total += count;
 	}
@@ -933,8 +936,8 @@ WaitForLockersMultiple(List *locktags, LOCKMODE lockmode, bool progress)
 		pgstat_progress_update_param(PROGRESS_WAITFOR_TOTAL, total);
 
 	/*
-	 * Note: GetLockConflicts() never reports our own xid, hence we need not
-	 * check for that.  Also, prepared xacts are reported and awaited.
+	 * Note: GetLockers() never reports our own xid, hence we need not check for
+	 * that.  Also, prepared xacts are reported and awaited.
 	 */
 
 	/* Finally wait for each such transaction to complete */
@@ -983,16 +986,16 @@ WaitForLockersMultiple(List *locktags, LOCKMODE lockmode, bool progress)
  * Same as WaitForLockersMultiple, for a single lock tag.
  */
 void
-WaitForLockers(LOCKTAG heaplocktag, LOCKMODE lockmode, bool progress)
+WaitForLockers(LOCKTAG heaplocktag, LOCKMASK lockmask, bool conflicting,
+			   bool progress)
 {
 	List	   *l;
 
 	l = list_make1(&heaplocktag);
-	WaitForLockersMultiple(l, lockmode, progress);
+	WaitForLockersMultiple(l, lockmask, conflicting, progress);
 	list_free(l);
 }
 
-
 /*
  *		LockDatabaseObject
  *
diff --git a/src/include/storage/lmgr.h b/src/include/storage/lmgr.h
index e8bd71ba68..7468a0500f 100644
--- a/src/include/storage/lmgr.h
+++ b/src/include/storage/lmgr.h
@@ -82,8 +82,10 @@ extern void XactLockTableWait(TransactionId xid, Relation rel,
 extern bool ConditionalXactLockTableWait(TransactionId xid);
 
 /* Lock VXIDs, specified by conflicting locktags */
-extern void WaitForLockers(LOCKTAG heaplocktag, LOCKMODE lockmode, bool progress);
-extern void WaitForLockersMultiple(List *locktags, LOCKMODE lockmode, bool progress);
+extern void WaitForLockers(LOCKTAG heaplocktag, LOCKMODE lockmode,
+						   bool conflicting ,bool progress);
+extern void WaitForLockersMultiple(List *locktags, LOCKMODE lockmode,
+								   bool conflicting, bool progress);
 
 /* Lock an XID for tuple insertion (used to wait for an insertion to finish) */
 extern uint32 SpeculativeInsertionLockAcquire(TransactionId xid);
-- 
2.34.1

