From 4a41002eea7eaa18bd51521798c19d191d9c6d8c Mon Sep 17 00:00:00 2001
From: Kyotaro Horiguchi <horikyota.ntt@gmail.com>
Date: Tue, 19 Jul 2022 11:50:29 +0900
Subject: [PATCH v7] Lift initial snapshot limit for logical replication

The replication command CREATE_REPLICATION_SLOT tends to fail with
"snapshot too large" when many subtransactions are active. This
problem stems from SnapBuildInitialSnapshot attempting to generate a
snapshot in which all active XIDs, including subxids, are stored
within the xip array. This patch mitigates the constraint on the
acceptable number of transaction IDs by utilizing the subxid array.

Author: Dilip Kumar and Kyotaro Horiguchi
---
 src/backend/replication/logical/snapbuild.c | 50 +++++++++++++++++----
 src/backend/storage/ipc/procarray.c         | 18 ++------
 src/include/access/transam.h                | 33 ++++++++++++++
 3 files changed, 77 insertions(+), 24 deletions(-)

diff --git a/src/backend/replication/logical/snapbuild.c b/src/backend/replication/logical/snapbuild.c
index 62542827e4..516e9ddfc8 100644
--- a/src/backend/replication/logical/snapbuild.c
+++ b/src/backend/replication/logical/snapbuild.c
@@ -565,10 +565,14 @@ Snapshot
 SnapBuildInitialSnapshot(SnapBuild *builder)
 {
 	Snapshot	snap;
+	int			ntxn;
+	int			xiplen;
 	TransactionId xid;
 	TransactionId safeXid;
 	TransactionId *newxip;
+	TransactionId *newsubxip;
 	int			newxcnt = 0;
+	int			newsubxcnt = 0;
 
 	Assert(XactIsoLevel == XACT_REPEATABLE_READ);
 	Assert(builder->building_full_snapshot);
@@ -610,9 +614,35 @@ SnapBuildInitialSnapshot(SnapBuild *builder)
 
 	MyProc->xmin = snap->xmin;
 
-	/* allocate in transaction context */
-	newxip = (TransactionId *)
-		palloc(sizeof(TransactionId) * GetMaxSnapshotXidCount());
+	/*
+	 * Aallocate in transaction context
+	 *
+	 * Since we know all the active XIDs, this snapshot won't be
+	 * suboverflowed. However, if the number of XIDs surpasses the XID arrays'
+	 * capacity, we cannot establish this snapshot because we don't know which
+	 * XIDs are top-level. Since the original snapshot is historical, we can't
+	 * reliably idetify the top-level XIDs. Thus, we have no choice but fail
+	 * when there are too many active XIDs.
+	 */
+	ntxn = TransactionIdDistance(snap->xmin, snap->xmax) - snap->xcnt;
+	xiplen = GetMaxSnapshotXidCount();
+	newxip = (TransactionId *) palloc(sizeof(TransactionId) * xiplen);
+
+	if (ntxn > GetMaxSnapshotXidCount())
+	{
+		/*
+		 * Since we can't identify the top-level XIDs, we can't carete a
+		 * suboverflowed snapshot.
+		 */
+		if (ntxn > xiplen + GetMaxSnapshotSubxidCount())
+			ereport(ERROR,
+					(errcode(ERRCODE_T_R_SERIALIZATION_FAILURE),
+					 errmsg("initial slot snapshot too large")));
+			
+		newsubxip = (TransactionId *)
+			palloc(sizeof(TransactionId) * GetMaxSnapshotSubxidCount());
+	}
+		
 
 	/*
 	 * snapbuild.c builds transactions in an "inverted" manner, which means it
@@ -633,21 +663,23 @@ SnapBuildInitialSnapshot(SnapBuild *builder)
 
 		if (test == NULL)
 		{
-			if (newxcnt >= GetMaxSnapshotXidCount())
-				ereport(ERROR,
-						(errcode(ERRCODE_T_R_SERIALIZATION_FAILURE),
-						 errmsg("initial slot snapshot too large")));
-
-			newxip[newxcnt++] = xid;
+			if (newxcnt < xiplen)
+				newxip[newxcnt++] = xid;
+			else
+				newsubxip[newsubxcnt++] = xid;
 		}
 
 		TransactionIdAdvance(xid);
 	}
 
+	Assert(newxcnt + newsubxcnt == ntxn);
+
 	/* adjust remaining snapshot fields as needed */
 	snap->snapshot_type = SNAPSHOT_MVCC;
 	snap->xcnt = newxcnt;
 	snap->xip = newxip;
+	snap->subxcnt = newsubxcnt;
+	snap->subxip = newsubxip;
 
 	return snap;
 }
diff --git a/src/backend/storage/ipc/procarray.c b/src/backend/storage/ipc/procarray.c
index ea91ce355f..a0f5ba5a36 100644
--- a/src/backend/storage/ipc/procarray.c
+++ b/src/backend/storage/ipc/procarray.c
@@ -4823,22 +4823,10 @@ KnownAssignedXidsAdd(TransactionId from_xid, TransactionId to_xid,
 	Assert(TransactionIdPrecedesOrEquals(from_xid, to_xid));
 
 	/*
-	 * Calculate how many array slots we'll need.  Normally this is cheap; in
-	 * the unusual case where the XIDs cross the wrap point, we do it the hard
-	 * way.
+	 * Calculate how many array slots we'll need. Both ends of the region are
+	 * inclusive.
 	 */
-	if (to_xid >= from_xid)
-		nxids = to_xid - from_xid + 1;
-	else
-	{
-		nxids = 1;
-		next_xid = from_xid;
-		while (TransactionIdPrecedes(next_xid, to_xid))
-		{
-			nxids++;
-			TransactionIdAdvance(next_xid);
-		}
-	}
+ 	nxids = TransactionIdDistance(from_xid, to_xid) + 1;
 
 	/*
 	 * Since only the startup process modifies the head/tail pointers, we
diff --git a/src/include/access/transam.h b/src/include/access/transam.h
index f5af6d3055..b6a3c8abee 100644
--- a/src/include/access/transam.h
+++ b/src/include/access/transam.h
@@ -370,6 +370,39 @@ FullTransactionIdNewer(FullTransactionId a, FullTransactionId b)
 	return b;
 }
 
+
+/* return the distance between the two IDs */
+static inline int
+TransactionIdDistance(TransactionId from_xid, TransactionId to_xid)
+{
+	int			nxids;
+	TransactionId next_xid;
+
+	StaticAssertDecl(sizeof(int) >= sizeof(TransactionId),
+					 "TransactionId exceeds the width of int");
+	Assert(TransactionIdPrecedesOrEquals(from_xid, to_xid));
+
+	/*
+	 * Calculate the distance between the two XIDs. Normally this is cheap; in
+	 * the unusual case where the XIDs cross the wrap point, we do it the hard
+	 * way.
+	 */
+	if (to_xid >= from_xid)
+		nxids = to_xid - from_xid;
+	else
+	{
+		nxids = 0;
+		next_xid = from_xid;
+		while (TransactionIdPrecedes(next_xid, to_xid))
+		{
+			nxids++;
+			TransactionIdAdvance(next_xid);
+		}
+	}
+
+	return nxids;
+}
+
 #endif							/* FRONTEND */
 
 #endif							/* TRANSAM_H */
-- 
2.31.1