v2-0001-Wait-for-vxids-in-CIC.patch

From 8739848c4ade4616f88c0ac93cf082554c8f1e13 Mon Sep 17 00:00:00 2001
From: Andrey Borodin <amborodin@acm.org>
Date: Mon, 19 Jul 2021 11:50:02 +0500
Subject: [PATCH v2] Wait for vxids in CIC

---
 src/backend/access/transam/twophase.c | 41 +++++++++++++--
 src/backend/storage/lmgr/lock.c       | 72 ++++++++++++++++++++-------
 src/include/access/twophase.h         |  2 +
 3 files changed, 94 insertions(+), 21 deletions(-)

diff --git a/src/backend/access/transam/twophase.c b/src/backend/access/transam/twophase.c
index f67d813c56..dba302d322 100644
--- a/src/backend/access/transam/twophase.c
+++ b/src/backend/access/transam/twophase.c
@@ -373,8 +373,10 @@ MarkAsPreparing(TransactionId xid, const char *gid,
 				TimestampTz prepared_at, Oid owner, Oid databaseid)
 {
 	GlobalTransaction gxact;
+
 	int			i;
 
+
 	if (strlen(gid) >= GIDSIZE)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
@@ -459,14 +461,16 @@ MarkAsPreparingGuts(GlobalTransaction gxact, TransactionId xid, const char *gid,
 	proc->pgprocno = gxact->pgprocno;
 	SHMQueueElemInit(&(proc->links));
 	proc->waitStatus = PROC_WAIT_STATUS_OK;
-	/* We set up the gxact's VXID as InvalidBackendId/XID */
-	proc->lxid = (LocalTransactionId) xid;
+
+	/* We set up the gxact's VXID as real for CIC purposes */
+	proc->lxid = MyProc->lxid;
 	proc->xid = xid;
 	Assert(proc->xmin == InvalidTransactionId);
 	proc->delayChkpt = false;
 	proc->statusFlags = 0;
 	proc->pid = 0;
-	proc->backendId = InvalidBackendId;
+	/* May be backendId of startup process */
+	proc->backendId = MyBackendId;
 	proc->databaseId = databaseid;
 	proc->roleId = owner;
 	proc->tempNamespaceId = InvalidOid;
@@ -846,6 +850,37 @@ TwoPhaseGetGXact(TransactionId xid, bool lock_held)
 	return result;
 }
 
+/*
+ * TwoPhaseGetXidByVXid
+ *		Try to lookup for vxid among prepared xacts
+ */
+TransactionId
+TwoPhaseGetXidByVXid(VirtualTransactionId vxid)
+{
+	TransactionId result = InvalidTransactionId;
+	int			i;
+
+	LWLockAcquire(TwoPhaseStateLock, LW_SHARED);
+
+	for (i = 0; i < TwoPhaseState->numPrepXacts; i++)
+	{
+		GlobalTransaction gxact = TwoPhaseState->prepXacts[i];
+		PGPROC	   *proc;
+		proc = &ProcGlobal->allProcs[gxact->pgprocno];
+
+		if (proc->backendId == vxid.backendId &&
+				proc->lxid == vxid.localTransactionId)
+		{
+			result = gxact->xid;
+			break;
+		}
+	}
+
+	LWLockRelease(TwoPhaseStateLock);
+
+	return result;
+}
+
 /*
  * TwoPhaseGetDummyBackendId
  *		Get the dummy backend ID for prepared transaction specified by XID
diff --git a/src/backend/storage/lmgr/lock.c b/src/backend/storage/lmgr/lock.c
index 108b4d9023..25e3bb6fda 100644
--- a/src/backend/storage/lmgr/lock.c
+++ b/src/backend/storage/lmgr/lock.c
@@ -386,6 +386,7 @@ static void LockRefindAndRelease(LockMethod lockMethodTable, PGPROC *proc,
 								 bool decrement_strong_lock_count);
 static void GetSingleProcBlockerStatusData(PGPROC *blocked_proc,
 										   BlockedProcsData *data);
+static bool PreparedXactLock(VirtualTransactionId vxid, TransactionId xid, bool wait);
 
 
 /*
@@ -4447,6 +4448,45 @@ VirtualXactLockTableCleanup(void)
 	}
 }
 
+/*
+ *		PreparedXactLock
+ *
+ *		Wait for xid completition if have xid. Otherwise try to find xid among
+ *		fake procarray entries.
+ */
+static bool PreparedXactLock(VirtualTransactionId vxid, TransactionId xid, bool wait)
+{
+	LockAcquireResult	lar;
+	LOCKTAG				tag;
+
+	/* Questionable heuristics */
+	if (max_prepared_xacts == 0)
+		return true;
+
+	if (!TransactionIdIsValid(xid))
+	{
+		/* We must search for vxid 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
+	}
+
+	if (!TransactionIdIsValid(xid))
+	{
+		/* Seems like vxid's 2pc is gone or never existed */
+		return true;
+	}
+
+	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;
+}
+
 /*
  *		VirtualXactLock
  *
@@ -4459,25 +4499,18 @@ VirtualXactLockTableCleanup(void)
 bool
 VirtualXactLock(VirtualTransactionId vxid, bool wait)
 {
-	LOCKTAG		tag;
-	PGPROC	   *proc;
+	LOCKTAG			tag;
+	PGPROC		   *proc;
+	TransactionId	xid = InvalidTransactionId;
 
 	Assert(VirtualTransactionIdIsValid(vxid));
 
+	/*
+	 * Already prepared transactions don't hold vxid locks.  The
+	 * LocalTransactionId is always a normal, locked XID.
+	 */
 	if (VirtualTransactionIdIsPreparedXact(vxid))
-	{
-		LockAcquireResult lar;
-
-		/*
-		 * Prepared transactions don't hold vxid locks.  The
-		 * LocalTransactionId is always a normal, locked XID.
-		 */
-		SET_LOCKTAG_TRANSACTION(tag, vxid.localTransactionId);
-		lar = LockAcquire(&tag, ShareLock, false, !wait);
-		if (lar != LOCKACQUIRE_NOT_AVAIL)
-			LockRelease(&tag, ShareLock, false);
-		return lar != LOCKACQUIRE_NOT_AVAIL;
-	}
+		return PreparedXactLock(vxid, vxid.localTransactionId, wait);
 
 	SET_LOCKTAG_VIRTUALTRANSACTION(tag, vxid);
 
@@ -4491,7 +4524,7 @@ VirtualXactLock(VirtualTransactionId vxid, bool wait)
 	 */
 	proc = BackendIdGetProc(vxid.backendId);
 	if (proc == NULL)
-		return true;
+		return PreparedXactLock(vxid, InvalidTransactionId, wait);
 
 	/*
 	 * We must acquire this lock before checking the backendId and lxid
@@ -4505,9 +4538,12 @@ VirtualXactLock(VirtualTransactionId vxid, bool wait)
 		|| proc->fpLocalTransactionId != vxid.localTransactionId)
 	{
 		LWLockRelease(&proc->fpInfoLock);
-		return true;
+		return PreparedXactLock(vxid, InvalidTransactionId, wait);
 	}
 
+	/* Save the xid to test if transaction coverted to 2pc later */
+	xid = proc->xid;
+
 	/*
 	 * If we aren't asked to wait, there's no need to set up a lock table
 	 * entry.  The transaction is still in progress, so just return false.
@@ -4559,7 +4595,7 @@ VirtualXactLock(VirtualTransactionId vxid, bool wait)
 	(void) LockAcquire(&tag, ShareLock, false, false);
 
 	LockRelease(&tag, ShareLock, false);
-	return true;
+	return PreparedXactLock(vxid, xid, wait);;
 }
 
 /*
diff --git a/src/include/access/twophase.h b/src/include/access/twophase.h
index 91786da784..19923555ed 100644
--- a/src/include/access/twophase.h
+++ b/src/include/access/twophase.h
@@ -58,4 +58,6 @@ extern void PrepareRedoAdd(char *buf, XLogRecPtr start_lsn,
 						   XLogRecPtr end_lsn, RepOriginId origin_id);
 extern void PrepareRedoRemove(TransactionId xid, bool giveWarning);
 extern void restoreTwoPhaseData(void);
+
+extern TransactionId TwoPhaseGetXidByVXid(VirtualTransactionId vxid);
 #endif							/* TWOPHASE_H */
-- 
2.24.3 (Apple Git-128)