From f11ba01cbfd0932d0c74f1bfdbfabb484bdfe943 Mon Sep 17 00:00:00 2001
From: Andrey Borodin <amborodin@acm.org>
Date: Wed, 21 Jul 2021 12:21:23 +0500
Subject: [PATCH v4 3/3] Wait on many xids returned from TwoPhaseGetXidByVXid()

---
 src/backend/access/transam/twophase.c | 18 +++++++++----
 src/backend/storage/lmgr/lock.c       | 38 ++++++++++++++++-----------
 src/include/access/twophase.h         | 12 ++++++++-
 3 files changed, 46 insertions(+), 22 deletions(-)

diff --git a/src/backend/access/transam/twophase.c b/src/backend/access/transam/twophase.c
index dba302d322..45e1c3259b 100644
--- a/src/backend/access/transam/twophase.c
+++ b/src/backend/access/transam/twophase.c
@@ -854,11 +854,13 @@ TwoPhaseGetGXact(TransactionId xid, bool lock_held)
  * TwoPhaseGetXidByVXid
  *		Try to lookup for vxid among prepared xacts
  */
-TransactionId
+XidListEntry
 TwoPhaseGetXidByVXid(VirtualTransactionId vxid)
 {
-	TransactionId result = InvalidTransactionId;
-	int			i;
+	int				i;
+	XidListEntry	result;
+	result.next = NULL;
+	result.xid = InvalidTransactionId;
 
 	LWLockAcquire(TwoPhaseStateLock, LW_SHARED);
 
@@ -871,8 +873,14 @@ TwoPhaseGetXidByVXid(VirtualTransactionId vxid)
 		if (proc->backendId == vxid.backendId &&
 				proc->lxid == vxid.localTransactionId)
 		{
-			result = gxact->xid;
-			break;
+			if (result.xid != InvalidTransactionId)
+			{
+				XidListEntry *copy = palloc(sizeof(XidListEntry));
+				copy->next = result.next;
+				copy->xid = result.xid;
+				result.next = copy;
+			}
+			result.xid = gxact->xid;
 		}
 	}
 
diff --git a/src/backend/storage/lmgr/lock.c b/src/backend/storage/lmgr/lock.c
index 50553e729c..55efb33a7b 100644
--- a/src/backend/storage/lmgr/lock.c
+++ b/src/backend/storage/lmgr/lock.c
@@ -4472,33 +4472,39 @@ static bool PreparedXactLock(VirtualTransactionId vxid, TransactionId xid, bool
 {
 	LockAcquireResult	lar;
 	LOCKTAG				tag;
+	XidListEntry		xidlist;
 
 	/* Questionable heuristics */
 	if (max_prepared_xacts == 0)
 		return true;
 
-	if (!TransactionIdIsValid(xid))
+	if (TransactionIdIsValid(xid))
 	{
-		/* We must search for vxid in 2pc state */
+		xidlist.xid = xid;
+		xidlist.next = NULL;
+	}
+	else
+	{
+		/* We must search for vxids in 2pc state */
 		/* XXX: O(N*N) complexity where N is number of prepared xacts */
-		xid = TwoPhaseGetXidByVXid(vxid);
-		if (TransactionIdIsValid(xid))							// Remove this line
-		{														// Remove this line
-			elog(NOTICE,"XXX: Sucessfully found xid by vxid");	// Remove this line
-		}														// Remove this line
+		xidlist = TwoPhaseGetXidByVXid(vxid);
+		elog(NOTICE,"XXX: Sucessfully found xid by vxid");	// Remove this line
 	}
 
-	if (!TransactionIdIsValid(xid))
+	while (true)
 	{
-		/* Seems like vxid's 2pc is gone or never existed */
-		return true;
+		if (TransactionIdIsValid(xidlist.xid))
+		{
+			SET_LOCKTAG_TRANSACTION(tag, xid);
+			lar = LockAcquire(&tag, ShareLock, false, !wait);
+			if (lar != LOCKACQUIRE_NOT_AVAIL)
+				LockRelease(&tag, ShareLock, false);
+			return lar != LOCKACQUIRE_NOT_AVAIL;
+		}
+		if (xidlist.next == NULL)
+			return true;
+		xidlist = *xidlist.next;
 	}
-
-	SET_LOCKTAG_TRANSACTION(tag, xid);
-	lar = LockAcquire(&tag, ShareLock, false, !wait);
-	if (lar != LOCKACQUIRE_NOT_AVAIL)
-		LockRelease(&tag, ShareLock, false);
-	return lar != LOCKACQUIRE_NOT_AVAIL;
 }
 
 /*
diff --git a/src/include/access/twophase.h b/src/include/access/twophase.h
index 19923555ed..70b3cb9391 100644
--- a/src/include/access/twophase.h
+++ b/src/include/access/twophase.h
@@ -59,5 +59,15 @@ extern void PrepareRedoAdd(char *buf, XLogRecPtr start_lsn,
 extern void PrepareRedoRemove(TransactionId xid, bool giveWarning);
 extern void restoreTwoPhaseData(void);
 
-extern TransactionId TwoPhaseGetXidByVXid(VirtualTransactionId vxid);
+/* 
+ * This struc is expected to be used as list very rarely. Under normal
+ * circumstances TwoPhaseGetXidByVXid() returns only one xid.
+ * But under certain conditions can return many xids or nothing.
+ */
+typedef struct XidListEntry
+{
+	TransactionId xid;
+	struct XidListEntry* next;
+} XidListEntry;
+extern XidListEntry TwoPhaseGetXidByVXid(VirtualTransactionId vxid);
 #endif							/* TWOPHASE_H */
-- 
2.24.3 (Apple Git-128)