diff --git a/src/backend/access/heap/heapam.c b/src/backend/access/heap/heapam.c
index 31073bc..464a73c 100644
--- a/src/backend/access/heap/heapam.c
+++ b/src/backend/access/heap/heapam.c
@@ -4222,22 +4222,27 @@ l3:
 		 */
 		if (!have_tuple_lock)
 		{
-			if (wait_policy == LockWaitBlock)
+			switch (wait_policy)
 			{
-				LockTupleTuplock(relation, tid, mode);
-			}
-			else if (wait_policy == LockWaitError)
-			{
-				if (!ConditionalLockTupleTuplock(relation, tid, mode))
-					ereport(ERROR,
-							(errcode(ERRCODE_LOCK_NOT_AVAILABLE),
-					errmsg("could not obtain lock on row in relation \"%s\"",
-						   RelationGetRelationName(relation))));
-			}
-			else /* wait_policy == LockWaitSkip */
-			{
-				if (!ConditionalLockTupleTuplock(relation, tid, mode))
-					return HeapTupleWouldBlock;
+				case LockWaitBlock:
+					LockTupleTuplock(relation, tid, mode);
+					break;
+				case LockWaitSkip:
+					if (!ConditionalLockTupleTuplock(relation, tid, mode))
+					{
+						result = HeapTupleWouldBlock;
+						/* recovery code expects to have buffer lock held */
+						LockBuffer(*buffer, BUFFER_LOCK_EXCLUSIVE);
+						goto failed;
+					}
+					break;
+				case LockWaitError:
+					if (!ConditionalLockTupleTuplock(relation, tid, mode))
+						ereport(ERROR,
+								(errcode(ERRCODE_LOCK_NOT_AVAILABLE),
+								 errmsg("could not obtain lock on row in relation \"%s\"",
+										RelationGetRelationName(relation))));
+					break;
 			}
 			have_tuple_lock = true;
 		}
@@ -4441,34 +4446,36 @@ l3:
 				if (status >= MultiXactStatusNoKeyUpdate)
 					elog(ERROR, "invalid lock mode in heap_lock_tuple");
 
-				/* wait for multixact to end */
-				if (wait_policy == LockWaitBlock)
+				/* wait for multixact to end, or die trying  */
+				switch (wait_policy)
 				{
-					MultiXactIdWait((MultiXactId) xwait, status, infomask,
-									relation, &tuple->t_data->t_ctid,
-									XLTW_Lock, NULL);
-				}
-				else if (wait_policy == LockWaitError)
-				{
-					if (!ConditionalMultiXactIdWait((MultiXactId) xwait,
-												  status, infomask, relation,
-													&tuple->t_data->t_ctid,
-													XLTW_Lock, NULL))
-						ereport(ERROR,
-								(errcode(ERRCODE_LOCK_NOT_AVAILABLE),
-								 errmsg("could not obtain lock on row in relation \"%s\"",
-										RelationGetRelationName(relation))));
-				}
-				else /* wait_policy == LockWaitSkip */
-				{
-					if (!ConditionalMultiXactIdWait((MultiXactId) xwait,
-													status, infomask, relation,
-													&tuple->t_data->t_ctid,
-													XLTW_Lock, NULL))
-					{
-						UnlockTupleTuplock(relation, tid, mode);
-						return HeapTupleWouldBlock;
-					}
+					case LockWaitBlock:
+						MultiXactIdWait((MultiXactId) xwait, status, infomask,
+										relation, &tuple->t_data->t_ctid, XLTW_Lock, NULL);
+						break;
+					case LockWaitSkip:
+						if (!ConditionalMultiXactIdWait((MultiXactId) xwait,
+														status, infomask, relation,
+														&tuple->t_data->t_ctid,
+														XLTW_Lock, NULL))
+						{
+							result = HeapTupleWouldBlock;
+							/* recovery code expects to have buffer lock held */
+							LockBuffer(*buffer, BUFFER_LOCK_EXCLUSIVE);
+							goto failed;
+						}
+						break;
+					case LockWaitError:
+						if (!ConditionalMultiXactIdWait((MultiXactId) xwait,
+														status, infomask, relation,
+														&tuple->t_data->t_ctid,
+														XLTW_Lock, NULL))
+							ereport(ERROR,
+									(errcode(ERRCODE_LOCK_NOT_AVAILABLE),
+									 errmsg("could not obtain lock on row in relation \"%s\"",
+											RelationGetRelationName(relation))));
+
+						break;
 				}
 
 				/* if there are updates, follow the update chain */
@@ -4514,27 +4521,29 @@ l3:
 			}
 			else
 			{
-				/* wait for regular transaction to end */
-				if (wait_policy == LockWaitBlock)
-				{
-					XactLockTableWait(xwait, relation, &tuple->t_data->t_ctid,
-									  XLTW_Lock);
-				}
-				else if (wait_policy == LockWaitError)
+				/* wait for regular transaction to end, or die trying */
+				switch (wait_policy)
 				{
-					if (!ConditionalXactLockTableWait(xwait))
-						ereport(ERROR,
-								(errcode(ERRCODE_LOCK_NOT_AVAILABLE),
-								 errmsg("could not obtain lock on row in relation \"%s\"",
-										RelationGetRelationName(relation))));
-				}
-				else /* wait_policy == LockWaitSkip */
-				{
-					if (!ConditionalXactLockTableWait(xwait))
-					{
-						UnlockTupleTuplock(relation, tid, mode);
-						return HeapTupleWouldBlock;
-					}
+					case LockWaitBlock:
+						XactLockTableWait(xwait, relation, &tuple->t_data->t_ctid,
+										  XLTW_Lock);
+						break;
+					case LockWaitSkip:
+						if (!ConditionalXactLockTableWait(xwait))
+						{
+							result = HeapTupleWouldBlock;
+							/* recovery code expects to have buffer lock held */
+							LockBuffer(*buffer, BUFFER_LOCK_EXCLUSIVE);
+							goto failed;
+						}
+						break;
+					case LockWaitError:
+						if (!ConditionalXactLockTableWait(xwait))
+							ereport(ERROR,
+									(errcode(ERRCODE_LOCK_NOT_AVAILABLE),
+									 errmsg("could not obtain lock on row in relation \"%s\"",
+											RelationGetRelationName(relation))));
+						break;
 				}
 
 				/* if there are updates, follow the update chain */
@@ -4597,7 +4606,8 @@ l3:
 failed:
 	if (result != HeapTupleMayBeUpdated)
 	{
-		Assert(result == HeapTupleSelfUpdated || result == HeapTupleUpdated);
+		Assert(result == HeapTupleSelfUpdated || result == HeapTupleUpdated ||
+			   result == HeapTupleWouldBlock);
 		Assert(!(tuple->t_data->t_infomask & HEAP_XMAX_INVALID));
 		hufd->ctid = tuple->t_data->t_ctid;
 		hufd->xmax = HeapTupleHeaderGetUpdateXid(tuple->t_data);
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 71df0c2..0f6e9b0 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -2372,13 +2372,13 @@ transformLockingClause(ParseState *pstate, Query *qry, LockingClause *lc,
 			switch (rte->rtekind)
 			{
 				case RTE_RELATION:
-					applyLockingClause(qry, i,
-									   lc->strength, lc->waitPolicy, pushedDown);
+					applyLockingClause(qry, i, lc->strength, lc->waitPolicy,
+									   pushedDown);
 					rte->requiredPerms |= ACL_SELECT_FOR_UPDATE;
 					break;
 				case RTE_SUBQUERY:
-					applyLockingClause(qry, i,
-									   lc->strength, lc->waitPolicy, pushedDown);
+					applyLockingClause(qry, i, lc->strength, lc->waitPolicy,
+									   pushedDown);
 
 					/*
 					 * FOR UPDATE/SHARE of subquery is propagated to all of
@@ -2424,15 +2424,13 @@ transformLockingClause(ParseState *pstate, Query *qry, LockingClause *lc,
 					switch (rte->rtekind)
 					{
 						case RTE_RELATION:
-							applyLockingClause(qry, i,
-											   lc->strength, lc->waitPolicy,
-											   pushedDown);
+							applyLockingClause(qry, i, lc->strength,
+											   lc->waitPolicy, pushedDown);
 							rte->requiredPerms |= ACL_SELECT_FOR_UPDATE;
 							break;
 						case RTE_SUBQUERY:
-							applyLockingClause(qry, i,
-											   lc->strength, lc->waitPolicy,
-											   pushedDown);
+							applyLockingClause(qry, i, lc->strength,
+											   lc->waitPolicy, pushedDown);
 							/* see comment above */
 							transformLockingClause(pstate, rte->subquery,
 												   allrels, true);
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 9dbfb40..6d4edcd 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -62,8 +62,8 @@ static void rewriteValuesRTE(RangeTblEntry *rte, Relation target_relation,
 static void rewriteTargetListUD(Query *parsetree, RangeTblEntry *target_rte,
 					Relation target_relation);
 static void markQueryForLocking(Query *qry, Node *jtnode,
-								LockClauseStrength strength, LockWaitPolicy waitPolicy,
-								bool pushedDown);
+					LockClauseStrength strength, LockWaitPolicy waitPolicy,
+					bool pushedDown);
 static List *matchLocks(CmdType event, RuleLock *rulelocks,
 		   int varno, Query *parsetree);
 static Query *fireRIRrules(Query *parsetree, List *activeRIRs,
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
index 383df42..f5da6bf 100644
--- a/src/include/parser/analyze.h
+++ b/src/include/parser/analyze.h
@@ -39,7 +39,7 @@ extern bool analyze_requires_snapshot(Node *parseTree);
 extern char *LCS_asString(LockClauseStrength strength);
 extern void CheckSelectLocking(Query *qry, LockClauseStrength strength);
 extern void applyLockingClause(Query *qry, Index rtindex,
-							   LockClauseStrength strength,
-							   LockWaitPolicy waitPolicy, bool pushedDown);
+				   LockClauseStrength strength,
+				   LockWaitPolicy waitPolicy, bool pushedDown);
 
 #endif   /* ANALYZE_H */