Split ProcSleep function into JoinWaitQueue and ProcSleep

Split ProcSleep into two functions: JoinWaitQueue and ProcSleep.
JoinWaitQueue is called while holding the partition lock, and inserts
the current process to the wait queue, while ProcSleep() does the
actual sleeping. ProcSleep() is now called without holding the
partition lock, and it no longer re-acquires the partition lock before
returning. That makes the wakeup a little cheaper. Once upon a time,
re-acquiring the partition lock was needed to prevent a signal handler
from longjmping out at a bad time, but these days our signal handlers
just set flags, and longjmping can only happen at points where we
explicitly run CHECK_FOR_INTERRUPTS().

If JoinWaitQueue detects an "early deadlock" before even joining the
wait queue, it returns without changing the shared lock entry, leaving
the cleanup of the shared lock entry to the caller. This makes the
handling of an early deadlock the same as the dontWait=true case.

One small user-visible side-effect of this refactoring is that we now
only set the 'ps' title to say "waiting" when we actually enter the
sleep, not when the lock is skipped because dontWait=true, or when a
deadlock is detected early before entering the sleep.

This eliminates the 'lockAwaited' global variable in proc.c, which was
largely redundant with 'awaitedLock' in lock.c

Note: Updating the local lock table is now the caller's responsibility.
JoinWaitQueue and ProcSleep are now only responsible for modifying the
shared state. Seems a little nicer that way.

Based on Thomas Munro's earlier patch and observation that ProcSleep
doesn't really need to re-acquire the partition lock.

Reviewed-by: Maxim Orlov
Discussion: https://www.postgresql.org/message-id/7c2090cd-a72a-4e34-afaa-6dd2ef31440e@iki.fi

Author: hlinnaka

src/backend/storage/lmgr/lock.c

@@ -409,8 +409,7 @@ static PROCLOCK *SetupLockInTable(LockMethod lockMethodTable, PGPROC *proc,
 static void GrantLockLocal(LOCALLOCK *locallock, ResourceOwner owner);
 static void BeginStrongLockAcquire(LOCALLOCK *locallock, uint32 fasthashcode);
 static void FinishStrongLockAcquire(void);
-static void WaitOnLock(LOCALLOCK *locallock, ResourceOwner owner,
-					   bool dontWait);
+static ProcWaitStatus WaitOnLock(LOCALLOCK *locallock, ResourceOwner owner);
 static void ReleaseLockIfHeld(LOCALLOCK *locallock, bool sessionLock);
 static void LockReassignOwner(LOCALLOCK *locallock, ResourceOwner parent);
 static bool UnGrantLock(LOCK *lock, LOCKMODE lockmode,
@@ -843,6 +842,7 @@ LockAcquireExtended(const LOCKTAG *locktag,
 	uint32		hashcode;
 	LWLock	   *partitionLock;
 	bool		found_conflict;
+	ProcWaitStatus waitResult;
 	bool		log_lock = false;
 
 	if (lockmethodid <= 0 || lockmethodid >= lengthof(LockMethods))
@@ -1091,94 +1091,115 @@ LockAcquireExtended(const LOCKTAG *locktag,
 	{
 		/* No conflict with held or previously requested locks */
 		GrantLock(lock, proclock, lockmode);
-		GrantLockLocal(locallock, owner);
+		waitResult = PROC_WAIT_STATUS_OK;
 	}
 	else
 	{
 		/*
-		 * Sleep till someone wakes me up. We do this even in the dontWait
-		 * case, because while trying to go to sleep, we may discover that we
-		 * can acquire the lock immediately after all.
+		 * Join the lock's wait queue.  We call this even in the dontWait
+		 * case, because JoinWaitQueue() may discover that we can acquire the
+		 * lock immediately after all.
+		 */
+		waitResult = JoinWaitQueue(locallock, lockMethodTable, dontWait);
+	}
+
+	if (waitResult == PROC_WAIT_STATUS_ERROR)
+	{
+		/*
+		 * We're not getting the lock because a deadlock was detected already
+		 * while trying to join the wait queue, or because we would have to
+		 * wait but the caller requested no blocking.
+		 *
+		 * Undo the changes to shared entries before releasing the partition
+		 * lock.
 		 */
-		WaitOnLock(locallock, owner, dontWait);
+		AbortStrongLockAcquire();
+
+		if (proclock->holdMask == 0)
+		{
+			uint32		proclock_hashcode;
+
+			proclock_hashcode = ProcLockHashCode(&proclock->tag,
+												 hashcode);
+			dlist_delete(&proclock->lockLink);
+			dlist_delete(&proclock->procLink);
+			if (!hash_search_with_hash_value(LockMethodProcLockHash,
+											 &(proclock->tag),
+											 proclock_hashcode,
+											 HASH_REMOVE,
+											 NULL))
+				elog(PANIC, "proclock table corrupted");
+		}
+		else
+			PROCLOCK_PRINT("LockAcquire: did not join wait queue", proclock);
+		lock->nRequested--;
+		lock->requested[lockmode]--;
+		LOCK_PRINT("LockAcquire: did not join wait queue",
+				   lock, lockmode);
+		Assert((lock->nRequested > 0) &&
+			   (lock->requested[lockmode] >= 0));
+		Assert(lock->nGranted <= lock->nRequested);
+		LWLockRelease(partitionLock);
+		if (locallock->nLocks == 0)
+			RemoveLocalLock(locallock);
+
+		if (dontWait)
+		{
+			if (locallockp)
+				*locallockp = NULL;
+			return LOCKACQUIRE_NOT_AVAIL;
+		}
+		else
+		{
+			DeadLockReport();
+			/* DeadLockReport() will not return */
+		}
+	}
+
+	/*
+	 * We are now in the lock queue, or the lock was already granted.  If
+	 * queued, go to sleep.
+	 */
+	if (waitResult == PROC_WAIT_STATUS_WAITING)
+	{
+		Assert(!dontWait);
+		PROCLOCK_PRINT("LockAcquire: sleeping on lock", proclock);
+		LOCK_PRINT("LockAcquire: sleeping on lock", lock, lockmode);
+		LWLockRelease(partitionLock);
+
+		waitResult = WaitOnLock(locallock, owner);
 
 		/*
 		 * NOTE: do not do any material change of state between here and
 		 * return.  All required changes in locktable state must have been
 		 * done when the lock was granted to us --- see notes in WaitOnLock.
 		 */
 
-		/*
-		 * Check the proclock entry status. If dontWait = true, this is an
-		 * expected case; otherwise, it will only happen if something in the
-		 * ipc communication doesn't work correctly.
-		 */
-		if (!(proclock->holdMask & LOCKBIT_ON(lockmode)))
+		if (waitResult == PROC_WAIT_STATUS_ERROR)
 		{
-			AbortStrongLockAcquire();
-
-			if (dontWait)
-			{
-				/*
-				 * We can't acquire the lock immediately.  If caller specified
-				 * no blocking, remove useless table entries and return
-				 * LOCKACQUIRE_NOT_AVAIL without waiting.
-				 */
-				if (proclock->holdMask == 0)
-				{
-					uint32		proclock_hashcode;
-
-					proclock_hashcode = ProcLockHashCode(&proclock->tag,
-														 hashcode);
-					dlist_delete(&proclock->lockLink);
-					dlist_delete(&proclock->procLink);
-					if (!hash_search_with_hash_value(LockMethodProcLockHash,
-													 &(proclock->tag),
-													 proclock_hashcode,
-													 HASH_REMOVE,
-													 NULL))
-						elog(PANIC, "proclock table corrupted");
-				}
-				else
-					PROCLOCK_PRINT("LockAcquire: NOWAIT", proclock);
-				lock->nRequested--;
-				lock->requested[lockmode]--;
-				LOCK_PRINT("LockAcquire: conditional lock failed",
-						   lock, lockmode);
-				Assert((lock->nRequested > 0) &&
-					   (lock->requested[lockmode] >= 0));
-				Assert(lock->nGranted <= lock->nRequested);
-				LWLockRelease(partitionLock);
-				if (locallock->nLocks == 0)
-					RemoveLocalLock(locallock);
-				if (locallockp)
-					*locallockp = NULL;
-				return LOCKACQUIRE_NOT_AVAIL;
-			}
-			else
-			{
-				/*
-				 * We should have gotten the lock, but somehow that didn't
-				 * happen. If we get here, it's a bug.
-				 */
-				PROCLOCK_PRINT("LockAcquire: INCONSISTENT", proclock);
-				LOCK_PRINT("LockAcquire: INCONSISTENT", lock, lockmode);
-				LWLockRelease(partitionLock);
-				elog(ERROR, "LockAcquire failed");
-			}
+			/*
+			 * We failed as a result of a deadlock, see CheckDeadLock(). Quit
+			 * now.
+			 */
+			Assert(!dontWait);
+			DeadLockReport();
+			/* DeadLockReport() will not return */
 		}
-		PROCLOCK_PRINT("LockAcquire: granted", proclock);
-		LOCK_PRINT("LockAcquire: granted", lock, lockmode);
 	}
+	else
+		LWLockRelease(partitionLock);
+	Assert(waitResult == PROC_WAIT_STATUS_OK);
+
+	/* The lock was granted to us.  Update the local lock entry accordingly */
+	Assert((proclock->holdMask & LOCKBIT_ON(lockmode)) != 0);
+	GrantLockLocal(locallock, owner);
 
 	/*
 	 * Lock state is fully up-to-date now; if we error out after this, no
 	 * special error cleanup is required.
 	 */
 	FinishStrongLockAcquire();
 
-	LWLockRelease(partitionLock);
-
 	/*
 	 * Emit a WAL record if acquisition of this lock needs to be replayed in a
 	 * standby server.
@@ -1819,6 +1840,15 @@ GrantAwaitedLock(void)
 	GrantLockLocal(awaitedLock, awaitedOwner);
 }
 
+/*
+ * GetAwaitedLock -- Return the lock we're currently doing WaitOnLock on.
+ */
+LOCALLOCK *
+GetAwaitedLock(void)
+{
+	return awaitedLock;
+}
+
 /*
  * MarkLockClear -- mark an acquired lock as "clear"
  *
@@ -1836,14 +1866,12 @@ MarkLockClear(LOCALLOCK *locallock)
 /*
  * WaitOnLock -- wait to acquire a lock
  *
- * The appropriate partition lock must be held at entry, and will still be
- * held at exit.
+ * This is a wrapper around ProcSleep, with extra tracing and bookkeeping.
  */
-static void
-WaitOnLock(LOCALLOCK *locallock, ResourceOwner owner, bool dontWait)
+static ProcWaitStatus
+WaitOnLock(LOCALLOCK *locallock, ResourceOwner owner)
 {
-	LOCKMETHODID lockmethodid = LOCALLOCK_LOCKMETHOD(*locallock);
-	LockMethod	lockMethodTable = LockMethods[lockmethodid];
+	ProcWaitStatus result;
 
 	TRACE_POSTGRESQL_LOCK_WAIT_START(locallock->tag.lock.locktag_field1,
 									 locallock->tag.lock.locktag_field2,
@@ -1852,12 +1880,13 @@ WaitOnLock(LOCALLOCK *locallock, ResourceOwner owner, bool dontWait)
 									 locallock->tag.lock.locktag_type,
 									 locallock->tag.mode);
 
-	LOCK_PRINT("WaitOnLock: sleeping on lock",
-			   locallock->lock, locallock->tag.mode);
-
 	/* adjust the process title to indicate that it's waiting */
 	set_ps_display_suffix("waiting");
 
+	/*
+	 * Record the fact that we are waiting for a lock, so that
+	 * LockErrorCleanup will clean up if cancel/die happens.
+	 */
 	awaitedLock = locallock;
 	awaitedOwner = owner;
 
@@ -1880,30 +1909,7 @@ WaitOnLock(LOCALLOCK *locallock, ResourceOwner owner, bool dontWait)
 	 */
 	PG_TRY();
 	{
-		/*
-		 * If dontWait = true, we handle success and failure in the same way
-		 * here. The caller will be able to sort out what has happened.
-		 */
-		if (ProcSleep(locallock, lockMethodTable, dontWait) != PROC_WAIT_STATUS_OK
-			&& !dontWait)
-		{
-
-			/*
-			 * We failed as a result of a deadlock, see CheckDeadLock(). Quit
-			 * now.
-			 */
-			awaitedLock = NULL;
-			LOCK_PRINT("WaitOnLock: aborting on lock",
-					   locallock->lock, locallock->tag.mode);
-			LWLockRelease(LockHashPartitionLock(locallock->hashcode));
-
-			/*
-			 * Now that we aren't holding the partition lock, we can give an
-			 * error report including details about the detected deadlock.
-			 */
-			DeadLockReport();
-			/* not reached */
-		}
+		result = ProcSleep(locallock);
 	}
 	PG_CATCH();
 	{
@@ -1917,20 +1923,22 @@ WaitOnLock(LOCALLOCK *locallock, ResourceOwner owner, bool dontWait)
 	}
 	PG_END_TRY();
 
+	/*
+	 * We no longer want LockErrorCleanup to do anything.
+	 */
 	awaitedLock = NULL;
 
 	/* reset ps display to remove the suffix */
 	set_ps_display_remove_suffix();
 
-	LOCK_PRINT("WaitOnLock: wakeup on lock",
-			   locallock->lock, locallock->tag.mode);
-
 	TRACE_POSTGRESQL_LOCK_WAIT_DONE(locallock->tag.lock.locktag_field1,
 									locallock->tag.lock.locktag_field2,
 									locallock->tag.lock.locktag_field3,
 									locallock->tag.lock.locktag_field4,
 									locallock->tag.lock.locktag_type,
 									locallock->tag.mode);
+
+	return result;
 }
 
 /*