From 06f46c3c5db9133ee3723fd484c5759cc99d3179 Mon Sep 17 00:00:00 2001
From: Vignesh <vignesh21@gmail.com>
Date: Mon, 2 Jun 2025 13:53:46 +0530
Subject: [PATCH v8-PG13] Fix exponential memory allocation issue in logical
 decoding due to invalidation message distribution

Commit 247ee94150 added logic to distribute invalidation messages from
one transaction to all concurrent in-progress transactions. However,
since the source of each invalidation message was not tracked, transactions
could redundantly re-distribute already propagated messages. This led to an
exponential increase in invalidation messages and, ultimately, memory
allocation failures.

To fix this, a new field is introduced to track distributed invalidation
messages separately. This ensures that only transactions in the pre-existing
list receive the invalidations, preventing redundant propagation.

The accumulated distributed invalidations are now consumed at the end of
transaction processing in `ReorderBufferProcessTXN()`.
---
 .../replication/logical/reorderbuffer.c       | 136 ++++++++++++++----
 src/backend/replication/logical/snapbuild.c   |   5 +-
 src/include/replication/reorderbuffer.h       |  18 +++
 3 files changed, 131 insertions(+), 28 deletions(-)

diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index fa9413fa2a0..447ab745630 100644
--- a/src/backend/replication/logical/reorderbuffer.c
+++ b/src/backend/replication/logical/reorderbuffer.c
@@ -103,6 +103,7 @@
 #include "storage/sinval.h"
 #include "utils/builtins.h"
 #include "utils/combocid.h"
+#include "utils/inval.h"
 #include "utils/memdebug.h"
 #include "utils/memutils.h"
 #include "utils/rel.h"
@@ -220,7 +221,8 @@ static void ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 static ReorderBufferChange *ReorderBufferIterTXNNext(ReorderBuffer *rb, ReorderBufferIterTXNState *state);
 static void ReorderBufferIterTXNFinish(ReorderBuffer *rb,
 									   ReorderBufferIterTXNState *state);
-static void ReorderBufferExecuteInvalidations(ReorderBuffer *rb, ReorderBufferTXN *txn);
+static void ReorderBufferExecuteInvalidations(ReorderBuffer *rb, ReorderBufferTXN *txn,
+											  bool is_distributed);
 
 /*
  * ---------------------------------------
@@ -406,6 +408,12 @@ ReorderBufferReturnTXN(ReorderBuffer *rb, ReorderBufferTXN *txn)
 		txn->invalidations = NULL;
 	}
 
+	if (txn->distributed_invalidations)
+	{
+		pfree(txn->distributed_invalidations);
+		txn->distributed_invalidations = NULL;
+	}
+
 	/* Reset the toast hash */
 	ReorderBufferToastReset(rb, txn);
 
@@ -1883,7 +1891,7 @@ ReorderBufferCommit(ReorderBuffer *rb, TransactionId xid,
 						 * see new catalog contents, so execute all
 						 * invalidations.
 						 */
-						ReorderBufferExecuteInvalidations(rb, txn);
+						ReorderBufferExecuteInvalidations(rb, txn, false);
 					}
 
 					break;
@@ -1921,7 +1929,13 @@ ReorderBufferCommit(ReorderBuffer *rb, TransactionId xid,
 		AbortCurrentTransaction();
 
 		/* make sure there's no cache pollution */
-		ReorderBufferExecuteInvalidations(rb, txn);
+		if (rbtxn_inval_all_cache(txn))
+			InvalidateSystemCaches();
+		else
+		{
+			ReorderBufferExecuteInvalidations(rb, txn, false);
+			ReorderBufferExecuteInvalidations(rb, txn, true);
+		}
 
 		if (using_subtxn)
 			RollbackAndReleaseCurrentSubTransaction();
@@ -1947,7 +1961,7 @@ ReorderBufferCommit(ReorderBuffer *rb, TransactionId xid,
 		AbortCurrentTransaction();
 
 		/* make sure there's no cache pollution */
-		ReorderBufferExecuteInvalidations(rb, txn);
+		ReorderBufferExecuteInvalidations(rb, txn, false);
 
 		if (using_subtxn)
 			RollbackAndReleaseCurrentSubTransaction();
@@ -2253,6 +2267,42 @@ ReorderBufferAddNewTupleCids(ReorderBuffer *rb, TransactionId xid,
 	txn->ntuplecids++;
 }
 
+/*
+ * Helper function to store invalidation messages, handles both regular and
+ * distributed invalidations.
+ */
+static void
+ReorderBufferAddInvalidationsCommon(ReorderBuffer *rb, TransactionId xid,
+									XLogRecPtr lsn, Size nmsgs,
+									SharedInvalidationMessage *msgs,
+									ReorderBufferTXN *txn, bool for_inval)
+{
+	MemoryContext oldcontext;
+	uint32	   *ninvalidations;
+	SharedInvalidationMessage **invalidations;
+
+	oldcontext = MemoryContextSwitchTo(rb->context);
+
+	ninvalidations = for_inval ? &txn->ninvalidations_distr : &txn->ninvalidations;
+	invalidations = for_inval ? &txn->distributed_invalidations : &txn->invalidations;
+
+	/* Accumulate invalidations. */
+	if (*ninvalidations == 0)
+	{
+		*ninvalidations = nmsgs;
+		*invalidations = (SharedInvalidationMessage *)
+			palloc(sizeof(SharedInvalidationMessage) * nmsgs);
+		memcpy(*invalidations, msgs, sizeof(SharedInvalidationMessage) * nmsgs);
+	}
+	else
+	{
+		*invalidations = (SharedInvalidationMessage *)
+			repalloc(*invalidations, sizeof(SharedInvalidationMessage) *
+					 (*ninvalidations + nmsgs));
+	}
+	MemoryContextSwitchTo(oldcontext);
+}
+
 /*
  * Setup the invalidation of the toplevel transaction.
  *
@@ -2264,12 +2314,9 @@ ReorderBufferAddInvalidations(ReorderBuffer *rb, TransactionId xid,
 							  SharedInvalidationMessage *msgs)
 {
 	ReorderBufferTXN *txn;
-	MemoryContext oldcontext;
 
 	txn = ReorderBufferTXNByXid(rb, xid, true, NULL, lsn, true);
 
-	oldcontext = MemoryContextSwitchTo(rb->context);
-
 	/*
 	 * Collect all the invalidations under the top transaction, if available,
 	 * so that we can execute them all together.
@@ -2282,27 +2329,58 @@ ReorderBufferAddInvalidations(ReorderBuffer *rb, TransactionId xid,
 
 	Assert(nmsgs > 0);
 
-	/* Accumulate invalidations. */
-	if (txn->ninvalidations == 0)
+	ReorderBufferAddInvalidationsCommon(rb, xid, lsn, nmsgs, msgs, txn, false);
+}
+
+/*
+ * Accumulate the invalidations sent by committed transactions for executing
+ * them later.
+ *
+ * This needs to be called by committed transactions to distribute the
+ * invalidations to the in-progress transactions.
+ */
+void
+ReorderBufferAddDistributedInvalidtions(ReorderBuffer *rb, TransactionId xid,
+										XLogRecPtr lsn, Size nmsgs,
+										SharedInvalidationMessage *msgs)
+{
+	ReorderBufferTXN *txn;
+	Size		req_mem_size;
+
+	txn = ReorderBufferTXNByXid(rb, xid, true, NULL, lsn, true);
+
+	/*
+	 * Collect all the invalidations under the top transaction, if available,
+	 * so that we can execute them all together.
+	 */
+	if (txn->toplevel_xid)
 	{
-		txn->ninvalidations = nmsgs;
-		txn->invalidations = (SharedInvalidationMessage *)
-			palloc(sizeof(SharedInvalidationMessage) * nmsgs);
-		memcpy(txn->invalidations, msgs,
-			   sizeof(SharedInvalidationMessage) * nmsgs);
+		txn = ReorderBufferTXNByXid(rb, txn->toplevel_xid, true, NULL, lsn,
+									true);
 	}
-	else
+
+	Assert(nmsgs > 0);
+
+	req_mem_size = sizeof(SharedInvalidationMessage) * (txn->ninvalidations_distr + nmsgs);
+
+	/*
+	 * If the number of invalidation messages is larger than 8MB, it's more
+	 * efficient to invalidate the entire cache rather than processing each
+	 * message individually.
+	 */
+	if (req_mem_size > (8 * 1024 * 1024) || rbtxn_inval_all_cache(txn))
 	{
-		txn->invalidations = (SharedInvalidationMessage *)
-			repalloc(txn->invalidations, sizeof(SharedInvalidationMessage) *
-					 (txn->ninvalidations + nmsgs));
+		txn->txn_flags |= RBTXN_INVAL_ALL_CACHE;
 
-		memcpy(txn->invalidations + txn->ninvalidations, msgs,
-			   nmsgs * sizeof(SharedInvalidationMessage));
-		txn->ninvalidations += nmsgs;
+		if (txn->distributed_invalidations)
+		{
+			pfree(txn->distributed_invalidations);
+			txn->distributed_invalidations = NULL;
+			txn->ninvalidations_distr = 0;
+		}
 	}
-
-	MemoryContextSwitchTo(oldcontext);
+	else
+		ReorderBufferAddInvalidationsCommon(rb, xid, lsn, nmsgs, msgs, txn, true);
 }
 
 /*
@@ -2310,12 +2388,18 @@ ReorderBufferAddInvalidations(ReorderBuffer *rb, TransactionId xid,
  * in the changestream but we don't know which those are.
  */
 static void
-ReorderBufferExecuteInvalidations(ReorderBuffer *rb, ReorderBufferTXN *txn)
+ReorderBufferExecuteInvalidations(ReorderBuffer *rb, ReorderBufferTXN *txn,
+								  bool is_distributed)
 {
 	int			i;
+	uint32		ninvalidations;
+	SharedInvalidationMessage *invalidations;
 
-	for (i = 0; i < txn->ninvalidations; i++)
-		LocalExecuteInvalidationMessage(&txn->invalidations[i]);
+	ninvalidations = is_distributed ? txn->ninvalidations_distr : txn->ninvalidations;
+	invalidations = is_distributed ? txn->distributed_invalidations : txn->invalidations;
+
+	for (i = 0; i < ninvalidations; i++)
+		LocalExecuteInvalidationMessage(&invalidations[i]);
 }
 
 /*
diff --git a/src/backend/replication/logical/snapbuild.c b/src/backend/replication/logical/snapbuild.c
index 3bda41c5251..86e508e14be 100644
--- a/src/backend/replication/logical/snapbuild.c
+++ b/src/backend/replication/logical/snapbuild.c
@@ -940,8 +940,9 @@ SnapBuildDistributeSnapshotAndInval(SnapBuild *builder, XLogRecPtr lsn, Transact
 			{
 				Assert(msgs != NULL);
 
-				ReorderBufferAddInvalidations(builder->reorder, txn->xid, lsn,
-											  ninvalidations, msgs);
+				ReorderBufferAddDistributedInvalidtions(builder->reorder,
+														txn->xid, lsn,
+														ninvalidations, msgs);
 			}
 		}
 	}
diff --git a/src/include/replication/reorderbuffer.h b/src/include/replication/reorderbuffer.h
index 545cee891ed..6e097cac495 100644
--- a/src/include/replication/reorderbuffer.h
+++ b/src/include/replication/reorderbuffer.h
@@ -69,6 +69,12 @@ enum ReorderBufferChangeType
 /* forward declaration */
 struct ReorderBufferTXN;
 
+/* Should the complete cache be invalidated? */
+#define rbtxn_inval_all_cache(txn) \
+( \
+	((txn)->txn_flags & RBTXN_INVAL_ALL_CACHE) != 0 \
+)
+
 /*
  * a single 'change', can be an insert (with one tuple), an update (old, new),
  * or a delete (old).
@@ -163,6 +169,7 @@ typedef struct ReorderBufferChange
 #define RBTXN_HAS_CATALOG_CHANGES 0x0001
 #define RBTXN_IS_SUBXACT          0x0002
 #define RBTXN_IS_SERIALIZED       0x0004
+#define RBTXN_INVAL_ALL_CACHE     0x0008
 
 /* Does the transaction have catalog changes? */
 #define rbtxn_has_catalog_changes(txn) \
@@ -311,6 +318,13 @@ typedef struct ReorderBufferTXN
 	 * Size of this transaction (changes currently in memory, in bytes).
 	 */
 	Size		size;
+
+	/*
+	 * Stores the cache invalidation messages distributed by the committing
+	 * transaction.
+	 */
+	uint32          ninvalidations_distr;
+	SharedInvalidationMessage *distributed_invalidations;
 } ReorderBufferTXN;
 
 /* so we can define the callbacks used inside struct ReorderBuffer itself */
@@ -451,6 +465,10 @@ void		ReorderBufferAddNewTupleCids(ReorderBuffer *, TransactionId, XLogRecPtr ls
 										 CommandId cmin, CommandId cmax, CommandId combocid);
 void		ReorderBufferAddInvalidations(ReorderBuffer *, TransactionId, XLogRecPtr lsn,
 										  Size nmsgs, SharedInvalidationMessage *msgs);
+void		ReorderBufferAddDistributedInvalidtions(ReorderBuffer *rb,
+													TransactionId xid,
+													XLogRecPtr lsn, Size nmsgs,
+													SharedInvalidationMessage *msgs);
 void		ReorderBufferImmediateInvalidation(ReorderBuffer *, uint32 ninvalidations,
 											   SharedInvalidationMessage *invalidations);
 void		ReorderBufferProcessXid(ReorderBuffer *, TransactionId xid, XLogRecPtr lsn);
-- 
2.43.0

