v8-PG17-0001-Fix-exponential-memory-allocation-issue-in-l.patch

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

Commit cadaf0ac46 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       | 175 ++++++++++++++----
 src/backend/replication/logical/snapbuild.c   |   5 +-
 src/include/replication/reorderbuffer.h       |  17 ++
 3 files changed, 158 insertions(+), 39 deletions(-)

diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index 03eb005c39d..01a2cb87cb1 100644
--- a/src/backend/replication/logical/reorderbuffer.c
+++ b/src/backend/replication/logical/reorderbuffer.c
@@ -108,6 +108,7 @@
 #include "storage/fd.h"
 #include "storage/sinval.h"
 #include "utils/builtins.h"
+#include "utils/inval.h"
 #include "utils/memutils.h"
 #include "utils/rel.h"
 #include "utils/relfilenumbermap.h"
@@ -469,6 +470,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);
 
@@ -2574,7 +2581,14 @@ ReorderBufferProcessTXN(ReorderBuffer *rb, ReorderBufferTXN *txn,
 		AbortCurrentTransaction();
 
 		/* make sure there's no cache pollution */
-		ReorderBufferExecuteInvalidations(txn->ninvalidations, txn->invalidations);
+		if (rbtxn_inval_all_cache(txn))
+			InvalidateSystemCaches();
+		else
+		{
+			ReorderBufferExecuteInvalidations(txn->ninvalidations, txn->invalidations);
+			ReorderBufferExecuteInvalidations(txn->ninvalidations_distr,
+											  txn->distributed_invalidations);
+		}
 
 		if (using_subtxn)
 			RollbackAndReleaseCurrentSubTransaction();
@@ -2650,6 +2664,20 @@ ReorderBufferProcessTXN(ReorderBuffer *rb, ReorderBufferTXN *txn,
 			errdata = NULL;
 			curtxn->concurrent_abort = true;
 
+			/*
+			 * Make sure there's no cache pollution. Unlike the PG_TRY
+			 * section, this must be done unconditionally, since processing
+			 * may fail before we reach the invalidation messages. This isn't
+			 * needed in non-rollback cases, as the system cache is cleared
+			 * either by the API's error handling or when the walsender exits.
+			 */
+			if (rbtxn_inval_all_cache(txn))
+				InvalidateSystemCaches();
+			else
+				ReorderBufferExecuteInvalidations(txn->ninvalidations_distr,
+												  txn->distributed_invalidations);
+
+
 			/* Reset the TXN so that it is allowed to stream remaining data. */
 			ReorderBufferResetTXN(rb, txn, snapshot_now,
 								  command_id, prev_lsn,
@@ -3313,58 +3341,57 @@ ReorderBufferAddNewTupleCids(ReorderBuffer *rb, TransactionId xid,
 }
 
 /*
- * Accumulate the invalidations for executing them later.
+ * Helper function to store invalidation messages and queue them as a
+ * change in the reorder buffer. Handles both regular and distributed
+ * invalidations.
  *
- * This needs to be called for each XLOG_XACT_INVALIDATIONS message and
- * accumulates all the invalidation messages in the toplevel transaction, if
- * available, otherwise in the current transaction, as well as in the form of
- * change in reorder buffer.  We require to record it in form of the change
+ * This accumulates all the invalidation messages in the toplevel transaction,
+ * if available, otherwise in the current transaction, as well as in the form
+ * of change in reorder buffer.  We require to record it in form of the change
  * so that we can execute only the required invalidations instead of executing
  * all the invalidations on each CommandId increment.  We also need to
  * accumulate these in the txn buffer because in some cases where we skip
  * processing the transaction (see ReorderBufferForget), we need to execute
  * all the invalidations together.
  */
-void
-ReorderBufferAddInvalidations(ReorderBuffer *rb, TransactionId xid,
-							  XLogRecPtr lsn, Size nmsgs,
-							  SharedInvalidationMessage *msgs)
+static void
+ReorderBufferAddInvalidationsCommon(ReorderBuffer *rb, TransactionId xid,
+									XLogRecPtr lsn, Size nmsgs,
+									SharedInvalidationMessage *msgs,
+									ReorderBufferTXN *txn,
+									bool for_inval)
 {
-	ReorderBufferTXN *txn;
 	MemoryContext oldcontext;
 	ReorderBufferChange *change;
 
-	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.  See comments atop this
-	 * function.
-	 */
-	txn = rbtxn_get_toptxn(txn);
+	if (!for_inval || (for_inval && !rbtxn_inval_all_cache(txn)))
+	{
+		uint32	   *ninvalidations;
+		SharedInvalidationMessage **invalidations;
 
-	Assert(nmsgs > 0);
+		ninvalidations = for_inval ? &txn->ninvalidations_distr : &txn->ninvalidations;
+		invalidations = for_inval ? &txn->distributed_invalidations : &txn->invalidations;
 
-	/* Accumulate invalidations. */
-	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));
+		/* 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));
 
-		memcpy(txn->invalidations + txn->ninvalidations, msgs,
-			   nmsgs * sizeof(SharedInvalidationMessage));
-		txn->ninvalidations += nmsgs;
+			memcpy(*invalidations + *ninvalidations, msgs,
+				   nmsgs * sizeof(SharedInvalidationMessage));
+			*ninvalidations += nmsgs;
+		}
 	}
 
 	change = ReorderBufferGetChange(rb);
@@ -3380,6 +3407,80 @@ ReorderBufferAddInvalidations(ReorderBuffer *rb, TransactionId xid,
 	MemoryContextSwitchTo(oldcontext);
 }
 
+/*
+ * Accumulate the invalidations for executing them later.
+ *
+ * This needs to be called for each XLOG_XACT_INVALIDATIONS message.
+ */
+void
+ReorderBufferAddInvalidations(ReorderBuffer *rb, TransactionId xid,
+							  XLogRecPtr lsn, Size nmsgs,
+							  SharedInvalidationMessage *msgs)
+{
+	ReorderBufferTXN *txn;
+
+	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.  See comments atop
+	 * ReorderBufferAddInvalidationsCommon function.
+	 */
+	txn = rbtxn_get_toptxn(txn);
+
+	Assert(nmsgs > 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.  See comments
+	 * ReorderBufferAddInvalidationsCommon ReorderBufferAddInvalidations.
+	 */
+	txn = rbtxn_get_toptxn(txn);
+
+	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->txn_flags |= RBTXN_INVAL_ALL_CACHE;
+
+		if (txn->distributed_invalidations)
+		{
+			pfree(txn->distributed_invalidations);
+			txn->distributed_invalidations = NULL;
+			txn->ninvalidations_distr = 0;
+		}
+	}
+
+	ReorderBufferAddInvalidationsCommon(rb, xid, lsn, nmsgs, msgs, txn, true);
+}
+
 /*
  * Apply all invalidations we know. Possibly we only need parts at this point
  * in the changestream but we don't know which those are.
diff --git a/src/backend/replication/logical/snapbuild.c b/src/backend/replication/logical/snapbuild.c
index 110e0b0a044..7ca5ce71ebf 100644
--- a/src/backend/replication/logical/snapbuild.c
+++ b/src/backend/replication/logical/snapbuild.c
@@ -945,8 +945,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 4c56f219fd8..ab4a0c63151 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_ALL_CACHE		0x0200
 
 /* Does the transaction have catalog changes? */
 #define rbtxn_has_catalog_changes(txn) \
@@ -205,6 +206,12 @@ typedef struct ReorderBufferChange
 	((txn)->txn_flags & RBTXN_HAS_STREAMABLE_CHANGE) != 0 \
 )
 
+/* Should the complete cache be invalidated? */
+#define rbtxn_inval_all_cache(txn) \
+( \
+	((txn)->txn_flags & RBTXN_INVAL_ALL_CACHE) != 0 \
+)
+
 /*
  * Has this transaction been streamed to downstream?
  *
@@ -422,6 +429,13 @@ typedef struct ReorderBufferTXN
 	 * Private data pointer of the output plugin.
 	 */
 	void	   *output_plugin_private;
+
+	/*
+	 * 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 */
@@ -709,6 +723,9 @@ extern void ReorderBufferAddNewTupleCids(ReorderBuffer *rb, TransactionId xid,
 										 CommandId cmin, CommandId cmax, CommandId combocid);
 extern void ReorderBufferAddInvalidations(ReorderBuffer *rb, TransactionId xid, XLogRecPtr lsn,
 										  Size nmsgs, SharedInvalidationMessage *msgs);
+extern void ReorderBufferAddDistributedInvalidtions(ReorderBuffer *rb, TransactionId xid,
+													XLogRecPtr lsn, Size nmsgs,
+													SharedInvalidationMessage *msgs);
 extern void ReorderBufferImmediateInvalidation(ReorderBuffer *rb, uint32 ninvalidations,
 											   SharedInvalidationMessage *invalidations);
 extern void ReorderBufferProcessXid(ReorderBuffer *rb, TransactionId xid, XLogRecPtr lsn);
-- 
2.43.0