From 4451fbaef21046fddc5aaa64e2819a8bb0e48dd1 Mon Sep 17 00:00:00 2001
From: Will Mortensen <will@extrahop.com>
Date: Thu, 21 Dec 2023 22:08:51 -0800
Subject: [PATCH v6 1/3] Refactor GetLockConflicts() into more general
 GetLockers()

GetLockers() supports getting lockers in modes matching an arbitrary
lockmask, not just those that conflict with a specific lockmode.
GetLockConflicts() is still available as a convenience wrapper, and its
semantics are unchanged.

Also factor out and export GetLockConflictMask() for use in a later
commit.
---
 src/backend/storage/lmgr/lock.c | 91 +++++++++++++++++++++++----------
 src/include/storage/lock.h      |  3 ++
 2 files changed, 67 insertions(+), 27 deletions(-)

diff --git a/src/backend/storage/lmgr/lock.c b/src/backend/storage/lmgr/lock.c
index c70a1adb9a..41b35de019 100644
--- a/src/backend/storage/lmgr/lock.c
+++ b/src/backend/storage/lmgr/lock.c
@@ -2832,45 +2832,78 @@ FastPathGetRelationLockEntry(LOCALLOCK *locallock)
 	return proclock;
 }
 
+/*
+ * GetLockConflictMask
+ *		Return the LOCKMASK of lockmodes that would conflict with the given
+ *		lockmode if taken on locktag.
+ */
+LOCKMASK
+GetLockConflictMask(const LOCKTAG *locktag, LOCKMODE lockmode)
+{
+	LOCKMETHODID lockmethodid = locktag->locktag_lockmethodid;
+	LockMethod	lockMethodTable;
+
+	if (lockmethodid <= 0 || lockmethodid >= lengthof(LockMethods))
+		elog(ERROR, "unrecognized lock method: %d", lockmethodid);
+	lockMethodTable = LockMethods[lockmethodid];
+	if (lockmode <= 0 || lockmode > lockMethodTable->numLockModes)
+		elog(ERROR, "unrecognized lock mode: %d", lockmode);
+	return lockMethodTable->conflictTab[lockmode];
+}
+
 /*
  * GetLockConflicts
+ *
+ * Convenience wrapper for GetLockers conflicting with a single lockmode.
+ */
+VirtualTransactionId *
+GetLockConflicts(const LOCKTAG *locktag, LOCKMODE lockmode, int *countp)
+{
+	return GetLockers(locktag, GetLockConflictMask(locktag, lockmode), countp);
+}
+
+/*
+ * GetLockers
  *		Get an array of VirtualTransactionIds of xacts currently holding locks
- *		that would conflict with the specified lock/lockmode.
- *		xacts merely awaiting such a lock are NOT reported.
+ *		on the specified locktag and matching the specified getMask, which is
+ *		assumed to be valid for locktag. xacts merely awaiting such a lock are
+ *		NOT reported.
  *
  * The result array is palloc'd and is terminated with an invalid VXID.
  * *countp, if not null, is updated to the number of items set.
  *
  * Of course, the result could be out of date by the time it's returned, so
  * use of this function has to be thought about carefully.  Similarly, a
- * PGPROC with no "lxid" will be considered non-conflicting regardless of any
- * lock it holds.  Existing callers don't care about a locker after that
- * locker's pg_xact updates complete.  CommitTransaction() clears "lxid" after
- * pg_xact updates and before releasing locks.
+ * PGPROC with no "lxid" will not be returned regardless of any lock it holds.
+ * Existing callers don't care about a locker after that locker's pg_xact
+ * updates complete.  CommitTransaction() clears "lxid" after pg_xact updates
+ * and before releasing locks.
  *
- * Note we never include the current xact's vxid in the result array,
- * since an xact never blocks itself.
+ * Note we never include the current xact's vxid in the result array, because
+ * existing callers don't care to know about it, since an xact never blocks
+ * itself and can see its own uncommitted changes.
  */
 VirtualTransactionId *
-GetLockConflicts(const LOCKTAG *locktag, LOCKMODE lockmode, int *countp)
+GetLockers(const LOCKTAG *locktag, LOCKMASK getMask, int *countp)
 {
 	static VirtualTransactionId *vxids;
 	LOCKMETHODID lockmethodid = locktag->locktag_lockmethodid;
 	LockMethod	lockMethodTable;
+	int			numLockModes;
 	LOCK	   *lock;
-	LOCKMASK	conflictMask;
 	dlist_iter	proclock_iter;
 	PROCLOCK   *proclock;
 	uint32		hashcode;
 	LWLock	   *partitionLock;
 	int			count = 0;
+	int			i;
+	bool		checkFast = false;
 	int			fast_count = 0;
 
 	if (lockmethodid <= 0 || lockmethodid >= lengthof(LockMethods))
 		elog(ERROR, "unrecognized lock method: %d", lockmethodid);
 	lockMethodTable = LockMethods[lockmethodid];
-	if (lockmode <= 0 || lockmode > lockMethodTable->numLockModes)
-		elog(ERROR, "unrecognized lock mode: %d", lockmode);
+	numLockModes = lockMethodTable->numLockModes;
 
 	/*
 	 * Allocate memory to store results, and fill with InvalidVXID.  We only
@@ -2890,19 +2923,25 @@ GetLockConflicts(const LOCKTAG *locktag, LOCKMODE lockmode, int *countp)
 			palloc0(sizeof(VirtualTransactionId) *
 					(MaxBackends + max_prepared_xacts + 1));
 
-	/* Compute hash code and partition lock, and look up conflicting modes. */
+	/* Compute hash code and partition lock. */
 	hashcode = LockTagHashCode(locktag);
 	partitionLock = LockHashPartitionLock(hashcode);
-	conflictMask = lockMethodTable->conflictTab[lockmode];
 
 	/*
 	 * Fast path locks might not have been entered in the primary lock table.
-	 * If the lock we're dealing with could conflict with such a lock, we must
-	 * examine each backend's fast-path array for conflicts.
+	 * If getMask could match such a lock, we must examine each backend's
+	 * fast-path array.
 	 */
-	if (ConflictsWithRelationFastPath(locktag, lockmode))
+	for (i = 1; i <= numLockModes; i++)
+	{
+		if (((getMask & LOCKBIT_ON(i)) != 0) &&
+			EligibleForRelationFastPath(locktag, i)) {
+			checkFast = true;
+			break;
+		}
+	}
+	if (checkFast)
 	{
-		int			i;
 		Oid			relid = locktag->locktag_field2;
 		VirtualTransactionId vxid;
 
@@ -2955,12 +2994,12 @@ GetLockConflicts(const LOCKTAG *locktag, LOCKMODE lockmode, int *countp)
 
 				/*
 				 * There can only be one entry per relation, so if we found it
-				 * and it doesn't conflict, we can skip the rest of the slots.
+				 * and it doesn't match, we can skip the rest of the slots.
 				 */
-				if ((lockmask & conflictMask) == 0)
+				if ((lockmask & getMask) == 0)
 					break;
 
-				/* Conflict! */
+				/* Match! */
 				GET_VXID_FROM_PGPROC(vxid, *proc);
 
 				if (VirtualTransactionIdIsValid(vxid))
@@ -2975,7 +3014,7 @@ GetLockConflicts(const LOCKTAG *locktag, LOCKMODE lockmode, int *countp)
 		}
 	}
 
-	/* Remember how many fast-path conflicts we found. */
+	/* Remember how many fast-path matches we found. */
 	fast_count = count;
 
 	/*
@@ -3009,11 +3048,11 @@ GetLockConflicts(const LOCKTAG *locktag, LOCKMODE lockmode, int *countp)
 	{
 		proclock = dlist_container(PROCLOCK, lockLink, proclock_iter.cur);
 
-		if (conflictMask & proclock->holdMask)
+		if (getMask & proclock->holdMask)
 		{
 			PGPROC	   *proc = proclock->tag.myProc;
 
-			/* A backend never blocks itself */
+			/* A backend doesn't care about its own locks */
 			if (proc != MyProc)
 			{
 				VirtualTransactionId vxid;
@@ -3022,8 +3061,6 @@ GetLockConflicts(const LOCKTAG *locktag, LOCKMODE lockmode, int *countp)
 
 				if (VirtualTransactionIdIsValid(vxid))
 				{
-					int			i;
-
 					/* Avoid duplicate entries. */
 					for (i = 0; i < fast_count; ++i)
 						if (VirtualTransactionIdEquals(vxids[i], vxid))
@@ -3039,7 +3076,7 @@ GetLockConflicts(const LOCKTAG *locktag, LOCKMODE lockmode, int *countp)
 	LWLockRelease(partitionLock);
 
 	if (count > MaxBackends + max_prepared_xacts)	/* should never happen */
-		elog(PANIC, "too many conflicting locks found");
+		elog(PANIC, "too many locks found");
 
 	vxids[count].backendId = InvalidBackendId;
 	vxids[count].localTransactionId = InvalidLocalTransactionId;
diff --git a/src/include/storage/lock.h b/src/include/storage/lock.h
index 00679624f7..62c50597a8 100644
--- a/src/include/storage/lock.h
+++ b/src/include/storage/lock.h
@@ -574,8 +574,11 @@ extern HTAB *GetLockMethodLocalHash(void);
 #endif
 extern bool LockHasWaiters(const LOCKTAG *locktag,
 						   LOCKMODE lockmode, bool sessionLock);
+extern LOCKMASK GetLockConflictMask(const LOCKTAG *locktag, LOCKMODE lockmode);
 extern VirtualTransactionId *GetLockConflicts(const LOCKTAG *locktag,
 											  LOCKMODE lockmode, int *countp);
+extern VirtualTransactionId *GetLockers(const LOCKTAG *locktag,
+										LOCKMASK getMask, int *countp);
 extern void AtPrepare_Locks(void);
 extern void PostPrepare_Locks(TransactionId xid);
 extern bool LockCheckConflicts(LockMethod lockMethodTable,
-- 
2.34.1

