From a853ad2474b01d08e6f392655b749be404c1d273 Mon Sep 17 00:00:00 2001
From: Hayato Kuroda <kuroda.hayato@fujitsu.com>
Date: Thu, 5 Jun 2025 21:00:09 +0900
Subject: [PATCH] implement limitation approach for PG17

---
 .../replication/logical/reorderbuffer.c       | 116 ++++++++++++++----
 src/include/replication/reorderbuffer.h       |   7 ++
 2 files changed, 98 insertions(+), 25 deletions(-)

diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index 03eb005c39d..a73e58e9188 100644
--- a/src/backend/replication/logical/reorderbuffer.c
+++ b/src/backend/replication/logical/reorderbuffer.c
@@ -109,9 +109,22 @@
 #include "storage/sinval.h"
 #include "utils/builtins.h"
 #include "utils/memutils.h"
+#include "utils/inval.h"
 #include "utils/rel.h"
 #include "utils/relfilenumbermap.h"
 
+/*
+ * Each transaction is limited to a maximum of 512MB of inval messages. Once
+ * the number of distributed inval messages reach this threshold, the
+ * transaction is marked as RBTXN_INVAL_OVERFLOWED, invalidating all caches
+ * instead as we have lost some inval messages and hence
+ * don't know what needs to be invalidated.
+ *
+ * XXX set to very small value for testing purpose
+ */
+#define MAX_DISTR_INVAL_MSG_PER_TXN \
+	((5 * 16) / sizeof(SharedInvalidationMessage))
+
 /* entry for a hash table we use to map from xid to our transaction state */
 typedef struct ReorderBufferTXNByIdEnt
 {
@@ -2574,7 +2587,14 @@ ReorderBufferProcessTXN(ReorderBuffer *rb, ReorderBufferTXN *txn,
 		AbortCurrentTransaction();
 
 		/* make sure there's no cache pollution */
-		ReorderBufferExecuteInvalidations(txn->ninvalidations, txn->invalidations);
+		if (rbtxn_inval_overflowed(txn))
+		{
+			Assert(txn->ninvalidations == 0);
+			InvalidateSystemCaches();
+		}
+		else
+			ReorderBufferExecuteInvalidations(txn->ninvalidations,
+											  txn->invalidations);
 
 		if (using_subtxn)
 			RollbackAndReleaseCurrentSubTransaction();
@@ -2620,8 +2640,13 @@ ReorderBufferProcessTXN(ReorderBuffer *rb, ReorderBufferTXN *txn,
 		AbortCurrentTransaction();
 
 		/* make sure there's no cache pollution */
-		ReorderBufferExecuteInvalidations(txn->ninvalidations,
-										  txn->invalidations);
+		if (rbtxn_inval_overflowed(txn))
+		{
+			Assert(txn->ninvalidations == 0);
+			InvalidateSystemCaches();
+		}
+		else
+			ReorderBufferExecuteInvalidations(txn->ninvalidations, txn->invalidations);
 
 		if (using_subtxn)
 			RollbackAndReleaseCurrentSubTransaction();
@@ -2910,8 +2935,15 @@ ReorderBufferFinishPrepared(ReorderBuffer *rb, TransactionId xid,
 		rb->rollback_prepared(rb, txn, prepare_end_lsn, prepare_time);
 
 	/* cleanup: make sure there's no cache pollution */
-	ReorderBufferExecuteInvalidations(txn->ninvalidations,
-									  txn->invalidations);
+	if (rbtxn_inval_overflowed(txn))
+	{
+		Assert(txn->ninvalidations == 0);
+		InvalidateSystemCaches();
+	}
+	else
+		ReorderBufferExecuteInvalidations(txn->ninvalidations,
+										  txn->invalidations);
+
 	ReorderBufferCleanupTXN(rb, txn);
 }
 
@@ -2953,7 +2985,12 @@ ReorderBufferAbort(ReorderBuffer *rb, TransactionId xid, XLogRecPtr lsn,
 		 * happened in this transaction). We don't want the decoding of future
 		 * transactions to use those cache entries so execute invalidations.
 		 */
-		if (txn->ninvalidations > 0)
+		if (rbtxn_inval_overflowed(txn))
+		{
+			Assert(txn->ninvalidations == 0);
+			InvalidateSystemCaches();
+		}
+		else if (txn->ninvalidations > 0)
 			ReorderBufferImmediateInvalidation(rb, txn->ninvalidations,
 											   txn->invalidations);
 	}
@@ -3042,9 +3079,14 @@ ReorderBufferForget(ReorderBuffer *rb, TransactionId xid, XLogRecPtr lsn)
 	 * interested in the transaction's contents, it could have manipulated the
 	 * catalog and we need to update the caches according to that.
 	 */
-	if (txn->base_snapshot != NULL && txn->ninvalidations > 0)
-		ReorderBufferImmediateInvalidation(rb, txn->ninvalidations,
-										   txn->invalidations);
+	if (txn->base_snapshot != NULL)
+	{
+		if (rbtxn_inval_overflowed(txn))
+			InvalidateSystemCaches();
+		else if (txn->ninvalidations > 0)
+			ReorderBufferImmediateInvalidation(rb, txn->ninvalidations,
+											   txn->invalidations);
+	}
 	else
 		Assert(txn->ninvalidations == 0);
 
@@ -3347,24 +3389,48 @@ ReorderBufferAddInvalidations(ReorderBuffer *rb, TransactionId xid,
 
 	Assert(nmsgs > 0);
 
-	/* Accumulate invalidations. */
-	if (txn->ninvalidations == 0)
+	/* Accumulate the invalidation messages unless it's overflowed */
+	if (!rbtxn_inval_overflowed(txn))
 	{
-		txn->ninvalidations = nmsgs;
-		txn->invalidations = (SharedInvalidationMessage *)
-			palloc(sizeof(SharedInvalidationMessage) * nmsgs);
-		memcpy(txn->invalidations, msgs,
-			   sizeof(SharedInvalidationMessage) * nmsgs);
-	}
-	else
-	{
-		txn->invalidations = (SharedInvalidationMessage *)
-			repalloc(txn->invalidations, sizeof(SharedInvalidationMessage) *
-					 (txn->ninvalidations + nmsgs));
+		/*
+		 * Check the transaction has enough space for storing invalidation
+		 * messages.
+		 */
+		if (txn->ninvalidations + nmsgs >= MAX_DISTR_INVAL_MSG_PER_TXN)
+		{
+			/*
+			 * Mark the invalidation message as overflowed and free up the
+			 * messages accumulated so far.
+			 */
+			txn->txn_flags |= RBTXN_INVAL_OVERFLOWED;
+
+			elog(LOG, "RBTXN_INVAL_OVERFLOWED flag is set");
+
+			if (txn->invalidations)
+			{
+				pfree(txn->invalidations);
+				txn->invalidations = NULL;
+				txn->ninvalidations = 0;
+			}
+		}
+		else if (txn->ninvalidations == 0)
+		{
+			txn->ninvalidations = nmsgs;
+			txn->invalidations = (SharedInvalidationMessage *)
+				palloc(sizeof(SharedInvalidationMessage) * nmsgs);
+			memcpy(txn->invalidations, msgs,
+				sizeof(SharedInvalidationMessage) * nmsgs);
+		}
+		else
+		{
+			txn->invalidations = (SharedInvalidationMessage *)
+				repalloc(txn->invalidations, sizeof(SharedInvalidationMessage) *
+						(txn->ninvalidations + nmsgs));
 
-		memcpy(txn->invalidations + txn->ninvalidations, msgs,
-			   nmsgs * sizeof(SharedInvalidationMessage));
-		txn->ninvalidations += nmsgs;
+			memcpy(txn->invalidations + txn->ninvalidations, msgs,
+				nmsgs * sizeof(SharedInvalidationMessage));
+			txn->ninvalidations += nmsgs;
+		}
 	}
 
 	change = ReorderBufferGetChange(rb);
diff --git a/src/include/replication/reorderbuffer.h b/src/include/replication/reorderbuffer.h
index 4c56f219fd8..a23aa7fefb0 100644
--- a/src/include/replication/reorderbuffer.h
+++ b/src/include/replication/reorderbuffer.h
@@ -168,6 +168,7 @@ typedef struct ReorderBufferChange
 #define RBTXN_PREPARE             	0x0040
 #define RBTXN_SKIPPED_PREPARE	  	0x0080
 #define RBTXN_HAS_STREAMABLE_CHANGE	0x0100
+#define RBTXN_INVAL_OVERFLOWED		0x0200
 
 /* Does the transaction have catalog changes? */
 #define rbtxn_has_catalog_changes(txn) \
@@ -231,6 +232,12 @@ typedef struct ReorderBufferChange
 	((txn)->txn_flags & RBTXN_SKIPPED_PREPARE) != 0 \
 )
 
+/* Is the array of inval message overflowed? */
+#define rbtxn_inval_overflowed(txn) \
+( \
+	((txn)->txn_flags & RBTXN_INVAL_OVERFLOWED) != 0 \
+)
+
 /* Is this a top-level transaction? */
 #define rbtxn_is_toptxn(txn) \
 ( \
-- 
2.47.1

