commit effcebc965982755a0ff8822c31abf4ce806179c
Author: Simon Riggs <simon.riggs@enterprisedb.com>
Date:   Wed Jul 27 13:51:17 2022 +0100

    Optimize XidInMVCCSnapshot() for overflowed snapshots
    
    Simon Riggs
    
    Performance testing and a bug fix from Julien Tachoires

diff --git a/src/backend/access/transam/subtrans.c b/src/backend/access/transam/subtrans.c
index 66d3548155..922621b4fc 100644
--- a/src/backend/access/transam/subtrans.c
+++ b/src/backend/access/transam/subtrans.c
@@ -177,6 +177,61 @@ SubTransGetTopmostTransaction(TransactionId xid)
 	return previousXid;
 }
 
+/*
+ * SubTransGetTopmostTransactionPrecedes
+ *
+ * Different form of SubTransGetTopmostTransaction() that minimizes the number
+ * of iterations required to get the parent or stop if it is earlier than cutoff.
+ */
+bool
+SubTransGetTopmostTransactionPrecedes(TransactionId *xid, TransactionId cutoff_xid, bool standby)
+{
+	TransactionId parentXid = *xid,
+				previousXid = *xid;
+
+	/* Can't ask about stuff that might not be around anymore */
+	Assert(TransactionIdFollowsOrEquals(cutoff_xid, TransactionXmin));
+	Assert(TransactionIdFollowsOrEquals(*xid, cutoff_xid));
+
+	while (TransactionIdIsValid(parentXid))
+	{
+		previousXid = parentXid;
+
+		/*
+		 * Stop as soon as we are earlier than the cutoff. This saves multiple
+		 * lookups against subtrans when we have a deeply nested subxid with
+		 * a later snapshot with an xmin much higher than TransactionXmin.
+		 */
+		if (TransactionIdPrecedes(parentXid, cutoff_xid))
+		{
+			*xid = previousXid;
+			return true;
+		}
+		parentXid = SubTransGetParent(parentXid);
+
+		/*
+		 * By convention the parent xid gets allocated first, so should always
+		 * precede the child xid. Anything else points to a corrupted data
+		 * structure that could lead to an infinite loop, so exit.
+		 */
+		if (!TransactionIdPrecedes(parentXid, previousXid))
+			elog(ERROR, "pg_subtrans contains invalid entry: xid %u points to parent xid %u",
+				 previousXid, parentXid);
+
+		/*
+		 * subxids always point directly at parent on standby, so we can avoid
+		 * multiple loops in that case.
+		 */
+		if (standby)
+			break;
+	}
+
+	Assert(TransactionIdIsValid(previousXid));
+
+	*xid = previousXid;
+
+	return false;
+}
 
 /*
  * Initialization of shared memory for SUBTRANS
diff --git a/src/backend/storage/ipc/procarray.c b/src/backend/storage/ipc/procarray.c
index dadaa958a8..3e26632636 100644
--- a/src/backend/storage/ipc/procarray.c
+++ b/src/backend/storage/ipc/procarray.c
@@ -2370,27 +2370,34 @@ GetSnapshotData(Snapshot snapshot)
 			 *
 			 * Again, our own XIDs are not included in the snapshot.
 			 */
-			if (!suboverflowed)
+			if (subxidStates[pgxactoff].overflowed)
+				suboverflowed = true;
+
+			/*
+			 * Include all subxids, whether or not we overflowed. This is
+			 * important because it can reduce the number of accesses to SUBTRANS
+			 * when we search snapshots, which is important for scalability,
+			 * especially in the presence of both overflowed and long running xacts.
+			 * This is true when fraction of subxids held in subxip is a large
+			 * fraction of the total subxids in use. In the case where one or more
+			 * transactions had more subxids in progress than the total size of
+			 * the cache this might not be true, but we consider that case to be
+			 * unlikely, even if it is possible.
+			 */
 			{
+				int			nsubxids = subxidStates[pgxactoff].count;
 
-				if (subxidStates[pgxactoff].overflowed)
-					suboverflowed = true;
-				else
+				if (nsubxids > 0)
 				{
-					int			nsubxids = subxidStates[pgxactoff].count;
-
-					if (nsubxids > 0)
-					{
-						int			pgprocno = pgprocnos[pgxactoff];
-						PGPROC	   *proc = &allProcs[pgprocno];
+					int			pgprocno = pgprocnos[pgxactoff];
+					PGPROC	   *proc = &allProcs[pgprocno];
 
-						pg_read_barrier();	/* pairs with GetNewTransactionId */
+					pg_read_barrier();	/* pairs with GetNewTransactionId */
 
-						memcpy(snapshot->subxip + subcount,
-							   (void *) proc->subxids.xids,
-							   nsubxids * sizeof(TransactionId));
-						subcount += nsubxids;
-					}
+					memcpy(snapshot->subxip + subcount,
+						   (void *) proc->subxids.xids,
+						   nsubxids * sizeof(TransactionId));
+					subcount += nsubxids;
 				}
 			}
 		}
diff --git a/src/backend/utils/time/snapmgr.c b/src/backend/utils/time/snapmgr.c
index 5bc2a15160..42130f1fea 100644
--- a/src/backend/utils/time/snapmgr.c
+++ b/src/backend/utils/time/snapmgr.c
@@ -638,13 +638,9 @@ CopySnapshot(Snapshot snapshot)
 		newsnap->xip = NULL;
 
 	/*
-	 * Setup subXID array. Don't bother to copy it if it had overflowed,
-	 * though, because it's not used anywhere in that case. Except if it's a
-	 * snapshot taken during recovery; all the top-level XIDs are in subxip as
-	 * well in that case, so we mustn't lose them.
+	 * Setup subXID array.
 	 */
-	if (snapshot->subxcnt > 0 &&
-		(!snapshot->suboverflowed || snapshot->takenDuringRecovery))
+	if (snapshot->subxcnt > 0)
 	{
 		newsnap->subxip = (TransactionId *) ((char *) newsnap + subxipoff);
 		memcpy(newsnap->subxip, snapshot->subxip,
@@ -1238,8 +1234,10 @@ ExportSnapshot(Snapshot snapshot)
 		snapshot->subxcnt + nchildren > GetMaxSnapshotSubxidCount())
 		appendStringInfoString(&buf, "sof:1\n");
 	else
-	{
 		appendStringInfoString(&buf, "sof:0\n");
+
+	/* then unconditionally, since we always include all subxids */
+	{
 		appendStringInfo(&buf, "sxcnt:%d\n", snapshot->subxcnt + nchildren);
 		for (i = 0; i < snapshot->subxcnt; i++)
 			appendStringInfo(&buf, "sxp:%u\n", snapshot->subxip[i]);
@@ -1490,7 +1488,7 @@ ImportSnapshot(const char *idstr)
 
 	snapshot.suboverflowed = parseIntFromText("sof:", &filebuf, path);
 
-	if (!snapshot.suboverflowed)
+	/* then unconditionally, since we always include all subxids */
 	{
 		snapshot.subxcnt = xcnt = parseIntFromText("sxcnt:", &filebuf, path);
 
@@ -1504,11 +1502,6 @@ ImportSnapshot(const char *idstr)
 		for (i = 0; i < xcnt; i++)
 			snapshot.subxip[i] = parseXidFromText("sxp:", &filebuf, path);
 	}
-	else
-	{
-		snapshot.subxcnt = 0;
-		snapshot.subxip = NULL;
-	}
 
 	snapshot.takenDuringRecovery = parseIntFromText("rec:", &filebuf, path);
 
@@ -2285,6 +2278,7 @@ bool
 XidInMVCCSnapshot(TransactionId xid, Snapshot snapshot)
 {
 	uint32		i;
+	bool		have_parent = false;
 
 	/*
 	 * Make a quick range check to eliminate most XIDs without looking at the
@@ -2301,93 +2295,137 @@ XidInMVCCSnapshot(TransactionId xid, Snapshot snapshot)
 	if (TransactionIdFollowsOrEquals(xid, snapshot->xmax))
 		return true;
 
+	/*
+	 * This patch reorders the operations in XidMVCCSnapshot, so as to reduce
+	 * calls to SubTransGetParent to the absolute minimum needed. It also
+	 * introduces a number of optimizations based on detailed knowledge of
+	 * how KnonAssignedXids module works in procarray.c - the two are already
+	 * intimately connected, so no modularity violation here.
+	 * The previous code was neat, but not efficient for the overflow case.
+	 */
+retry_search:
+
 	/*
 	 * Snapshot information is stored slightly differently in snapshots taken
-	 * during recovery.
+	 * during recovery. xip is empty on standbys.
 	 */
 	if (!snapshot->takenDuringRecovery)
 	{
-		/*
-		 * If the snapshot contains full subxact data, the fastest way to
-		 * check things is just to compare the given XID against both subxact
-		 * XIDs and top-level XIDs.  If the snapshot overflowed, we have to
-		 * use pg_subtrans to convert a subxact XID to its parent XID, but
-		 * then we need only look at top-level XIDs not subxacts.
-		 */
-		if (!snapshot->suboverflowed)
-		{
-			/* we have full data, so search subxip */
-			int32		j;
-
-			for (j = 0; j < snapshot->subxcnt; j++)
-			{
-				if (TransactionIdEquals(xid, snapshot->subxip[j]))
-					return true;
-			}
+		int32 j;
 
-			/* not there, fall through to search xip[] */
-		}
-		else
+		for (i = 0; i < snapshot->xcnt; i++)
 		{
-			/*
-			 * Snapshot overflowed, so convert xid to top-level.  This is safe
-			 * because we eliminated too-old XIDs above.
-			 */
-			xid = SubTransGetTopmostTransaction(xid);
-
-			/*
-			 * If xid was indeed a subxact, we might now have an xid < xmin,
-			 * so recheck to avoid an array scan.  No point in rechecking
-			 * xmax.
-			 */
-			if (TransactionIdPrecedes(xid, snapshot->xmin))
-				return false;
+			if (TransactionIdEquals(xid, snapshot->xip[i]))
+				return true;
 		}
 
-		for (i = 0; i < snapshot->xcnt; i++)
+		/*
+		 * If we have the parent xid, then the xid is not in snapshot
+		 */
+		if (have_parent)
+			return false;
+
+		/*
+		 * subxcnt will be zero in some cases
+		 */
+		for (j = 0; j < snapshot->subxcnt; j++)
 		{
-			if (TransactionIdEquals(xid, snapshot->xip[i]))
+			if (TransactionIdEquals(xid, snapshot->subxip[j]))
 				return true;
 		}
 	}
 	else
 	{
-		int32		j;
+		int	first;
+		int	last;
+		int	result_index = -1;
 
-		/*
-		 * In recovery we store all xids in the subxact array because it is by
-		 * far the bigger array, and we mostly don't know which xids are
-		 * top-level and which are subxacts. The xip array is empty.
-		 *
-		 * We start by searching subtrans, if we overflowed.
-		 */
-		if (snapshot->suboverflowed)
+#ifdef USE_ASSERT_CHECKING
+		if (snapshot->subxcnt > 1)
 		{
-			/*
-			 * Snapshot overflowed, so convert xid to top-level.  This is safe
-			 * because we eliminated too-old XIDs above.
-			 */
-			xid = SubTransGetTopmostTransaction(xid);
+			int32 j;
 
 			/*
-			 * If xid was indeed a subxact, we might now have an xid < xmin,
-			 * so recheck to avoid an array scan.  No point in rechecking
-			 * xmax.
+			 * Confirm invariant that subxip array is in sorted order.
+			 * Start at the 2nd element, j=1
 			 */
-			if (TransactionIdPrecedes(xid, snapshot->xmin))
-				return false;
+			for (j = 1; j < snapshot->subxcnt; j++)
+			{
+				if (TransactionIdPrecedes(snapshot->subxip[j - 1], snapshot->subxip[j]))
+				{
+					elog(ERROR, "snapshot subxip array not in sorted order during recovery");
+					Assert(false);
+				}
+			}
 		}
+#endif
 
 		/*
-		 * We now have either a top-level xid higher than xmin or an
-		 * indeterminate xid. We don't know whether it's top level or subxact
-		 * but it doesn't matter. If it's present, the xid is visible.
+		 * KnownAssignedXids stores xids in sorted order, a property that is
+		 * relied upon for the removal of xids at commit. When we take a snapshot
+		 * the xids are also copied in sorted order, which means we can also use
+		 * binary search here to probe the xip array, rather than scan whole thing.
+		 *
+		 * Note that standby snapshots also contain full subxid information for
+		 * non-overflowed transactions, which increases the chances that we will
+		 * find our xid here, even if the snapshot as a whole has overflowed.
+		 * That sounds like standby snapshots are larger, but we filter out any
+		 * subxids greater than xmax, which we do not do in normal running.
 		 */
-		for (j = 0; j < snapshot->subxcnt; j++)
+		first = 0;
+		last = snapshot->subxcnt - 1;
+		while (first <= last)
 		{
-			if (TransactionIdEquals(xid, snapshot->subxip[j]))
-				return true;
+			int         mid_index;
+			TransactionId mid_xid;
+
+			mid_index = (first + last) / 2;
+			mid_xid = snapshot->subxip[mid_index];
+
+			if (xid == mid_xid)
+			{
+				result_index = mid_index;
+				break;
+			}
+			else if (TransactionIdPrecedes(xid, mid_xid))
+				last = mid_index - 1;
+			else
+				first = mid_index + 1;
 		}
+
+		if (result_index >= 0)
+			return true;			/* found in array */
+	}
+
+	if (!have_parent && snapshot->suboverflowed)
+	{
+		/*
+		 * If we haven't found xid yet, it might be because it is a subxid
+		 * that is not present because we overflowed, but it might also be
+		 * because the xid is not in the snapshot.
+		 *
+		 * It is important we do this step last because it is expensive,
+		 * and if everybody does this then SubTransSLRU glows white hot.
+		 *
+		 * Use SubTransGetTopmostTransactionPrecedes(), which has been
+		 * specially provided to help here. This does two things for us:
+		 *
+		 * 1. On standby, get the parent directly, since in a standby SUBTRANS
+		 * always stores the direct parent only. Doing this avoids
+		 * one lookup of subtrans, since SubTransGetTopmostTransaction()
+		 * always does at least 2 SUBTRANS lookups, the last lookup is
+		 * how the loop knows it has found the parent in normal running.
+		 *
+		 * 2. Stops the iteration to find the parent as soon as we find an
+		 * xid earlier than snapshot->xmin, so we do the minimum lookups.
+		 */
+		if (SubTransGetTopmostTransactionPrecedes(&xid,
+												  snapshot->xmin,
+												  snapshot->takenDuringRecovery))
+			return false;
+
+		have_parent = true;
+		goto retry_search; /* search arrays again, now we have parent */
 	}
 
 	return false;
diff --git a/src/include/access/subtrans.h b/src/include/access/subtrans.h
index f94e116640..fc4103b769 100644
--- a/src/include/access/subtrans.h
+++ b/src/include/access/subtrans.h
@@ -17,6 +17,7 @@
 extern void SubTransSetParent(TransactionId xid, TransactionId parent);
 extern TransactionId SubTransGetParent(TransactionId xid);
 extern TransactionId SubTransGetTopmostTransaction(TransactionId xid);
+extern bool SubTransGetTopmostTransactionPrecedes(TransactionId *xid, TransactionId cutoff_xid, bool standby);
 
 extern Size SUBTRANSShmemSize(void);
 extern void SUBTRANSShmemInit(void);
diff --git a/src/test/isolation/expected/subxid-overflow.out b/src/test/isolation/expected/subxid-overflow.out
new file mode 100644
index 0000000000..f6c96f9a34
--- /dev/null
+++ b/src/test/isolation/expected/subxid-overflow.out
@@ -0,0 +1,17 @@
+Parsed test spec with 2 sessions
+
+starting permutation: subxid xmax s2sel s1c
+step subxid: BEGIN; SELECT gen_subxids(100);
+gen_subxids
+-----------
+           
+(1 row)
+
+step xmax: BEGIN; INSERT INTO subxids VALUES (1); COMMIT;
+step s2sel: SELECT count(*) FROM subxids WHERE subx = 0;
+count
+-----
+    0
+(1 row)
+
+step s1c: COMMIT;
diff --git a/src/test/isolation/isolation_schedule b/src/test/isolation/isolation_schedule
index 529a4cbd4d..7e10cfe022 100644
--- a/src/test/isolation/isolation_schedule
+++ b/src/test/isolation/isolation_schedule
@@ -1,3 +1,4 @@
+test: subxid-overflow
 test: read-only-anomaly
 test: read-only-anomaly-2
 test: read-only-anomaly-3
diff --git a/src/test/isolation/specs/subxid-overflow.spec b/src/test/isolation/specs/subxid-overflow.spec
new file mode 100644
index 0000000000..fda9f449f1
--- /dev/null
+++ b/src/test/isolation/specs/subxid-overflow.spec
@@ -0,0 +1,46 @@
+# Subtransaction overflow
+#
+# This test is designed to cover some code paths which only occur when
+# one transaction has overflowed the subtransaction cache.
+
+setup
+{
+DROP TABLE IF EXISTS subxids;
+CREATE TABLE subxids (subx integer);
+
+CREATE OR REPLACE FUNCTION gen_subxids (n integer)
+ RETURNS VOID
+ LANGUAGE plpgsql
+AS $$
+BEGIN
+  IF n <= 0 THEN
+    INSERT INTO subxids VALUES (0);
+    RETURN;
+  ELSE
+    PERFORM gen_subxids(n - 1);
+    RETURN;
+  END IF;
+EXCEPTION /* generates a subxid */
+  WHEN raise_exception THEN NULL;
+END;
+$$;
+}
+
+teardown
+{
+ DROP TABLE subxids;
+ DROP FUNCTION gen_subxids(integer);
+}
+
+session s1
+step subxid	{ BEGIN; SELECT gen_subxids(100); }
+step s1c	{ COMMIT; }
+
+session s2
+# move xmax forwards
+step xmax	{ BEGIN; INSERT INTO subxids VALUES (1); COMMIT;}
+# will see step subxid as still running
+step s2sel	{ SELECT count(*) FROM subxids WHERE subx = 0; }
+
+# s2sel will see subxid as still running
+permutation subxid xmax s2sel s1c
