From 31503be386bd5bea52bdc7248b76c7a4355fd142 Mon Sep 17 00:00:00 2001
From: Hayato Kuroda <kuroda.hayato@fujitsu.com>
Date: Wed, 21 May 2025 18:02:36 +0900
Subject: [PATCH v3-PG17] Avoid distributing invalidation messages several
 times

---
 .../replication/logical/reorderbuffer.c       | 134 ++++++++++++++----
 src/backend/replication/logical/snapbuild.c   |   5 +-
 src/include/replication/reorderbuffer.h       |   3 +
 3 files changed, 114 insertions(+), 28 deletions(-)

diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index 03eb005c39d..47e5cdc4de6 100644
--- a/src/backend/replication/logical/reorderbuffer.c
+++ b/src/backend/replication/logical/reorderbuffer.c
@@ -2093,6 +2093,49 @@ ReorderBufferResetTXN(ReorderBuffer *rb, ReorderBufferTXN *txn,
 	Assert(txn->size == 0);
 }
 
+/*
+ * Lookup changes in the transaction and apply only invalidation messages
+ */
+static void
+ReorderBufferProcessInvalidations(ReorderBuffer *rb, ReorderBufferTXN *txn)
+{
+	ReorderBufferChange *change;
+	ReorderBufferIterTXNState *volatile iterstate = NULL;
+
+	ReorderBufferIterTXNInit(rb, txn, &iterstate);
+	while ((change = ReorderBufferIterTXNNext(rb, iterstate)) != NULL)
+	{
+		CHECK_FOR_INTERRUPTS();
+
+		switch (change->action)
+		{
+				/* Only interested in invalidation messages */
+			case REORDER_BUFFER_CHANGE_INVALIDATION:
+				/* Execute the invalidation messages locally */
+				ReorderBufferExecuteInvalidations(change->data.inval.ninvalidations,
+												  change->data.inval.invalidations);
+				break;
+
+				/* Skip other changes because the transaction was aborted */
+			case REORDER_BUFFER_CHANGE_INSERT:
+			case REORDER_BUFFER_CHANGE_UPDATE:
+			case REORDER_BUFFER_CHANGE_DELETE:
+			case REORDER_BUFFER_CHANGE_MESSAGE:
+			case REORDER_BUFFER_CHANGE_INTERNAL_SNAPSHOT:
+			case REORDER_BUFFER_CHANGE_INTERNAL_COMMAND_ID:
+			case REORDER_BUFFER_CHANGE_INTERNAL_TUPLECID:
+			case REORDER_BUFFER_CHANGE_INTERNAL_SPEC_INSERT:
+			case REORDER_BUFFER_CHANGE_INTERNAL_SPEC_CONFIRM:
+			case REORDER_BUFFER_CHANGE_INTERNAL_SPEC_ABORT:
+			case REORDER_BUFFER_CHANGE_TRUNCATE:
+				break;
+		}
+	}
+
+	/* clean up the iterator */
+	ReorderBufferIterTXNFinish(rb, iterstate);
+}
+
 /*
  * Helper function for ReorderBufferReplay and ReorderBufferStreamTXN.
  *
@@ -2119,6 +2162,7 @@ ReorderBufferProcessTXN(ReorderBuffer *rb, ReorderBufferTXN *txn,
 	ReorderBufferChange *volatile specinsert = NULL;
 	volatile bool stream_started = false;
 	ReorderBufferTXN *volatile curtxn = NULL;
+	bool		contains_inval = false;
 
 	/* build data to be able to lookup the CommandIds of catalog tuples */
 	ReorderBufferBuildTupleCidHash(rb, txn);
@@ -2423,6 +2467,9 @@ ReorderBufferProcessTXN(ReorderBuffer *rb, ReorderBufferTXN *txn,
 					/* Execute the invalidation messages locally */
 					ReorderBufferExecuteInvalidations(change->data.inval.ninvalidations,
 													  change->data.inval.invalidations);
+
+					/* Also tracks the existence to iterate changes again */
+					contains_inval = true;
 					break;
 
 				case REORDER_BUFFER_CHANGE_INTERNAL_SNAPSHOT:
@@ -2576,6 +2623,14 @@ ReorderBufferProcessTXN(ReorderBuffer *rb, ReorderBufferTXN *txn,
 		/* make sure there's no cache pollution */
 		ReorderBufferExecuteInvalidations(txn->ninvalidations, txn->invalidations);
 
+		/*
+		 * The change queue must be checked again if the transaction contains
+		 * the inval messages because some them are not stored in
+		 * txn->invalidations.
+		 */
+		if (contains_inval)
+			ReorderBufferProcessInvalidations(rb, txn);
+
 		if (using_subtxn)
 			RollbackAndReleaseCurrentSubTransaction();
 
@@ -2650,6 +2705,14 @@ ReorderBufferProcessTXN(ReorderBuffer *rb, ReorderBufferTXN *txn,
 			errdata = NULL;
 			curtxn->concurrent_abort = true;
 
+			/*
+			 * Iterate changes in the transaction to execute invalidation
+			 * messages. Unlike the PG_TRY part, this must be done
+			 * unconditionally because the processing might fail before we
+			 * reach invalidation messages.
+			 */
+			ReorderBufferProcessInvalidations(rb, txn);
+
 			/* Reset the TXN so that it is allowed to stream remaining data. */
 			ReorderBufferResetTXN(rb, txn, snapshot_now,
 								  command_id, prev_lsn,
@@ -3330,41 +3393,60 @@ ReorderBufferAddInvalidations(ReorderBuffer *rb, TransactionId xid,
 							  XLogRecPtr lsn, Size nmsgs,
 							  SharedInvalidationMessage *msgs)
 {
-	ReorderBufferTXN *txn;
+	ReorderBufferAddInvalidationsExtended(rb, xid, lsn, nmsgs, msgs, true);
+}
+
+/*
+ * An extended workhorse version of ReorderBufferAddInvalidations(). This can
+ * control whether invalidation messages can be accumulated in reorderbuffer
+ * transactions, which would be distributed to all concurrent transactions at
+ * commit.
+ */
+void
+ReorderBufferAddInvalidationsExtended(ReorderBuffer *rb, TransactionId xid,
+									  XLogRecPtr lsn, Size nmsgs,
+									  SharedInvalidationMessage *msgs,
+									  bool needs_distribute)
+{
 	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);
-
 	Assert(nmsgs > 0);
 
-	/* Accumulate invalidations. */
-	if (txn->ninvalidations == 0)
-	{
-		txn->ninvalidations = nmsgs;
-		txn->invalidations = (SharedInvalidationMessage *)
-			palloc(sizeof(SharedInvalidationMessage) * nmsgs);
-		memcpy(txn->invalidations, msgs,
-			   sizeof(SharedInvalidationMessage) * nmsgs);
-	}
-	else
+	if (needs_distribute)
 	{
-		txn->invalidations = (SharedInvalidationMessage *)
-			repalloc(txn->invalidations, sizeof(SharedInvalidationMessage) *
-					 (txn->ninvalidations + nmsgs));
+		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 ReorderBufferAddInvalidations.
+		 */
+		txn = rbtxn_get_toptxn(txn);
 
-		memcpy(txn->invalidations + txn->ninvalidations, msgs,
-			   nmsgs * sizeof(SharedInvalidationMessage));
-		txn->ninvalidations += nmsgs;
+		/* 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));
+
+			memcpy(txn->invalidations + txn->ninvalidations, msgs,
+				   nmsgs * sizeof(SharedInvalidationMessage));
+			txn->ninvalidations += nmsgs;
+		}
 	}
 
 	change = ReorderBufferGetChange(rb);
diff --git a/src/backend/replication/logical/snapbuild.c b/src/backend/replication/logical/snapbuild.c
index 110e0b0a044..a8225c1d905 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);
+				ReorderBufferAddInvalidationsExtended(builder->reorder, txn->xid, lsn,
+													  ninvalidations, msgs,
+													  false);
 			}
 		}
 	}
diff --git a/src/include/replication/reorderbuffer.h b/src/include/replication/reorderbuffer.h
index 4c56f219fd8..001d0e716aa 100644
--- a/src/include/replication/reorderbuffer.h
+++ b/src/include/replication/reorderbuffer.h
@@ -709,6 +709,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 ReorderBufferAddInvalidationsExtended(ReorderBuffer *rb, TransactionId xid, XLogRecPtr lsn,
+												  Size nmsgs, SharedInvalidationMessage *msgs,
+												  bool needs_distribute);
 extern void ReorderBufferImmediateInvalidation(ReorderBuffer *rb, uint32 ninvalidations,
 											   SharedInvalidationMessage *invalidations);
 extern void ReorderBufferProcessXid(ReorderBuffer *rb, TransactionId xid, XLogRecPtr lsn);
-- 
2.47.1

