v9-0002-Add-Selective-Invalidation-of-Cache.patch

From 7ce19b7dafdf659a9689856211e81c501dfe498f Mon Sep 17 00:00:00 2001
From: Shlok Kyal <shlok.kyal.oss@gmail.com>
Date: Wed, 25 Sep 2024 11:41:42 +0530
Subject: [PATCH v9 2/2] Add Selective Invalidation of Cache

When we alter a publication, add/drop namespace to/from publication,
alter a namespace all the cache for all the tables are invalidated. With
this patch for the above operationns we will invalidate the cache of
only the desired tables.
---
 src/backend/replication/pgoutput/pgoutput.c |  52 ++++----
 src/backend/utils/cache/inval.c             | 127 +++++++++++++++++++-
 src/include/storage/sinval.h                |   9 ++
 src/include/utils/inval.h                   |   4 +
 src/test/subscription/t/100_bugs.pl         |  91 ++++++++++++++
 5 files changed, 253 insertions(+), 30 deletions(-)

diff --git a/src/backend/replication/pgoutput/pgoutput.c b/src/backend/replication/pgoutput/pgoutput.c
index 00e7024563..ba480e7e48 100644
--- a/src/backend/replication/pgoutput/pgoutput.c
+++ b/src/backend/replication/pgoutput/pgoutput.c
@@ -126,6 +126,8 @@ typedef struct RelationSyncEntry
 {
 	Oid			relid;			/* relation oid */
 
+	Oid			schemaid;		/* schema oid */
+
 	bool		replicate_valid;	/* overall validity flag for entry */
 
 	bool		schema_sent;
@@ -216,6 +218,7 @@ static RelationSyncEntry *get_rel_sync_entry(PGOutputData *data,
 static void rel_sync_cache_relation_cb(Datum arg, Oid relid);
 static void rel_sync_cache_publication_cb(Datum arg, int cacheid,
 										  uint32 hashvalue);
+static void rel_sync_cache_namespacerel_cb(Datum arg, int nspid);
 static void set_schema_sent_in_streamed_txn(RelationSyncEntry *entry,
 											TransactionId xid);
 static bool get_schema_sent_in_streamed_txn(RelationSyncEntry *entry,
@@ -1739,12 +1742,6 @@ static void
 publication_invalidation_cb(Datum arg, int cacheid, uint32 hashvalue)
 {
 	publications_valid = false;
-
-	/*
-	 * Also invalidate per-relation cache so that next time the filtering info
-	 * is checked it will be updated with the new publication settings.
-	 */
-	rel_sync_cache_publication_cb(arg, cacheid, hashvalue);
 }
 
 /*
@@ -1911,26 +1908,7 @@ init_rel_sync_cache(MemoryContext cachectx)
 
 	/* We must update the cache entry for a relation after a relcache flush */
 	CacheRegisterRelcacheCallback(rel_sync_cache_relation_cb, (Datum) 0);
-
-	/*
-	 * Flush all cache entries after a pg_namespace change, in case it was a
-	 * schema rename affecting a relation being replicated.
-	 */
-	CacheRegisterSyscacheCallback(NAMESPACEOID,
-								  rel_sync_cache_publication_cb,
-								  (Datum) 0);
-
-	/*
-	 * Flush all cache entries after any publication changes.  (We need no
-	 * callback entry for pg_publication, because publication_invalidation_cb
-	 * will take care of it.)
-	 */
-	CacheRegisterSyscacheCallback(PUBLICATIONRELMAP,
-								  rel_sync_cache_publication_cb,
-								  (Datum) 0);
-	CacheRegisterSyscacheCallback(PUBLICATIONNAMESPACEMAP,
-								  rel_sync_cache_publication_cb,
-								  (Datum) 0);
+	CacheRegisterNspcacheCallback(rel_sync_cache_namespacerel_cb, (Datum) 0);
 
 	relation_callbacks_registered = true;
 }
@@ -2076,6 +2054,8 @@ get_rel_sync_entry(PGOutputData *data, Relation relation)
 		entry->estate = NULL;
 		memset(entry->exprstate, 0, sizeof(entry->exprstate));
 
+		entry->schemaid = schemaId;
+
 		/*
 		 * Build publication cache. We can't use one provided by relcache as
 		 * relcache considers all publications that the given relation is in,
@@ -2349,6 +2329,26 @@ rel_sync_cache_publication_cb(Datum arg, int cacheid, uint32 hashvalue)
 	}
 }
 
+/*
+ * Namespace invalidation callback
+ */
+static void
+rel_sync_cache_namespacerel_cb(Datum arg, int nspid)
+{
+	HASH_SEQ_STATUS status;
+	RelationSyncEntry *entry;
+
+	if (RelationSyncCache == NULL)
+		return;
+
+	hash_seq_init(&status, RelationSyncCache);
+	while ((entry = (RelationSyncEntry *) hash_seq_search(&status)) != NULL)
+	{
+		if (entry->replicate_valid && entry->schemaid == nspid)
+			entry->replicate_valid = false;
+	}
+}
+
 /* Send Replication origin */
 static void
 send_repl_origin(LogicalDecodingContext *ctx, RepOriginId origin_id,
diff --git a/src/backend/utils/cache/inval.c b/src/backend/utils/cache/inval.c
index 603aa4157b..fc0d91aec9 100644
--- a/src/backend/utils/cache/inval.c
+++ b/src/backend/utils/cache/inval.c
@@ -114,6 +114,7 @@
 #include "access/xact.h"
 #include "access/xloginsert.h"
 #include "catalog/catalog.h"
+#include "catalog/pg_namespace.h"
 #include "catalog/pg_constraint.h"
 #include "miscadmin.h"
 #include "storage/sinval.h"
@@ -160,6 +161,9 @@
  */
 #define CatCacheMsgs 0
 #define RelCacheMsgs 1
+#define NspCacheMsgs 2
+
+#define NumberofCache 3
 
 /* Pointers to main arrays in TopTransactionContext */
 typedef struct InvalMessageArray
@@ -168,13 +172,13 @@ typedef struct InvalMessageArray
 	int			maxmsgs;		/* current allocated size of array */
 } InvalMessageArray;
 
-static InvalMessageArray InvalMessageArrays[2];
+static InvalMessageArray InvalMessageArrays[NumberofCache];
 
 /* Control information for one logical group of messages */
 typedef struct InvalidationMsgsGroup
 {
-	int			firstmsg[2];	/* first index in relevant array */
-	int			nextmsg[2];		/* last+1 index */
+	int			firstmsg[NumberofCache];	/* first index in relevant array */
+	int			nextmsg[NumberofCache];		/* last+1 index */
 } InvalidationMsgsGroup;
 
 /* Macros to help preserve InvalidationMsgsGroup abstraction */
@@ -189,6 +193,7 @@ typedef struct InvalidationMsgsGroup
 	do { \
 		SetSubGroupToFollow(targetgroup, priorgroup, CatCacheMsgs); \
 		SetSubGroupToFollow(targetgroup, priorgroup, RelCacheMsgs); \
+		SetSubGroupToFollow(targetgroup, priorgroup, NspCacheMsgs);	\
 	} while (0)
 
 #define NumMessagesInSubGroup(group, subgroup) \
@@ -196,7 +201,8 @@ typedef struct InvalidationMsgsGroup
 
 #define NumMessagesInGroup(group) \
 	(NumMessagesInSubGroup(group, CatCacheMsgs) + \
-	 NumMessagesInSubGroup(group, RelCacheMsgs))
+	 NumMessagesInSubGroup(group, RelCacheMsgs) + \
+	 NumMessagesInSubGroup(group, NspCacheMsgs))
 
 
 /*----------------
@@ -251,6 +257,7 @@ int			debug_discard_caches = 0;
 
 #define MAX_SYSCACHE_CALLBACKS 64
 #define MAX_RELCACHE_CALLBACKS 10
+#define MAX_NSPCACHE_CALLBACKS 10
 
 static struct SYSCACHECALLBACK
 {
@@ -270,7 +277,14 @@ static struct RELCACHECALLBACK
 	Datum		arg;
 }			relcache_callback_list[MAX_RELCACHE_CALLBACKS];
 
+static struct NSPCACHECALLBACK
+{
+	NspcacheCallbackFunction function;
+	Datum		arg;
+}			nspcache_callback_list[MAX_NSPCACHE_CALLBACKS];
+
 static int	relcache_callback_count = 0;
+static int	nspcache_callback_count = 0;
 
 /* ----------------------------------------------------------------
  *				Invalidation subgroup support functions
@@ -464,6 +478,35 @@ AddRelcacheInvalidationMessage(InvalidationMsgsGroup *group,
 	AddInvalidationMessage(group, RelCacheMsgs, &msg);
 }
 
+static void
+AddNspcacheInvalidationMessage(InvalidationMsgsGroup *group,
+							   Oid dbId, Oid nspId)
+{
+	SharedInvalidationMessage msg;
+
+	/*
+	 * Don't add a duplicate item. We assume dbId need not be checked because
+	 * it will never change. InvalidOid for relId means all relations so we
+	 * don't need to add individual ones when it is present.
+	 */
+
+	ProcessMessageSubGroup(group, NspCacheMsgs,
+						   if (msg->nc.id == SHAREDINVALNSPCACHE_ID &&
+							   (msg->nc.nspId == nspId ||
+								msg->nc.nspId == InvalidOid))
+						   return);
+
+
+	/* OK, add the item */
+	msg.nc.id = SHAREDINVALNSPCACHE_ID;
+	msg.nc.dbId = dbId;
+	msg.nc.nspId = nspId;
+	/* check AddCatcacheInvalidationMessage() for an explanation */
+	VALGRIND_MAKE_MEM_DEFINED(&msg, sizeof(msg));
+
+	AddInvalidationMessage(group, NspCacheMsgs, &msg);
+}
+
 /*
  * Add a snapshot inval entry
  *
@@ -502,6 +545,7 @@ AppendInvalidationMessages(InvalidationMsgsGroup *dest,
 {
 	AppendInvalidationMessageSubGroup(dest, src, CatCacheMsgs);
 	AppendInvalidationMessageSubGroup(dest, src, RelCacheMsgs);
+	AppendInvalidationMessageSubGroup(dest, src, NspCacheMsgs);
 }
 
 /*
@@ -516,6 +560,7 @@ ProcessInvalidationMessages(InvalidationMsgsGroup *group,
 {
 	ProcessMessageSubGroup(group, CatCacheMsgs, func(msg));
 	ProcessMessageSubGroup(group, RelCacheMsgs, func(msg));
+	ProcessMessageSubGroup(group, NspCacheMsgs, func(msg));
 }
 
 /*
@@ -528,6 +573,7 @@ ProcessInvalidationMessagesMulti(InvalidationMsgsGroup *group,
 {
 	ProcessMessageSubGroupMulti(group, CatCacheMsgs, func(msgs, n));
 	ProcessMessageSubGroupMulti(group, RelCacheMsgs, func(msgs, n));
+	ProcessMessageSubGroupMulti(group, NspCacheMsgs, func(msgs, n));
 }
 
 /* ----------------------------------------------------------------
@@ -590,6 +636,18 @@ RegisterRelcacheInvalidation(Oid dbId, Oid relId)
 		transInvalInfo->RelcacheInitFileInval = true;
 }
 
+/*
+ * RegisterNspcacheInvalidation
+ *
+ * As above, but register a namespace invalidation event.
+ */
+static void
+RegisterNspcacheInvalidation(Oid dbId, Oid nspId)
+{
+	AddNspcacheInvalidationMessage(&transInvalInfo->CurrentCmdInvalidMsgs,
+								   dbId, nspId);
+}
+
 /*
  * RegisterSnapshotInvalidation
  *
@@ -660,6 +718,8 @@ PrepareInvalidationState(void)
 		InvalMessageArrays[CatCacheMsgs].maxmsgs = 0;
 		InvalMessageArrays[RelCacheMsgs].msgs = NULL;
 		InvalMessageArrays[RelCacheMsgs].maxmsgs = 0;
+		InvalMessageArrays[NspCacheMsgs].msgs = NULL;
+		InvalMessageArrays[NspCacheMsgs].maxmsgs = 0;
 	}
 
 	transInvalInfo = myInfo;
@@ -773,6 +833,20 @@ LocalExecuteInvalidationMessage(SharedInvalidationMessage *msg)
 		else if (msg->sn.dbId == MyDatabaseId)
 			InvalidateCatalogSnapshot();
 	}
+	else if (msg->id == SHAREDINVALNSPCACHE_ID)
+	{
+		if (msg->nc.dbId == MyDatabaseId || msg->nc.dbId == InvalidOid)
+		{
+			int			i;
+
+			for (i = 0; i < nspcache_callback_count; i++)
+			{
+				struct NSPCACHECALLBACK *ncitem = nspcache_callback_list + i;
+
+				ncitem->function(ncitem->arg, msg->nc.nspId);
+			}
+		}
+	}
 	else
 		elog(FATAL, "unrecognized SI message ID: %d", msg->id);
 }
@@ -944,6 +1018,18 @@ xactGetCommittedInvalidationMessages(SharedInvalidationMessage **msgs,
 										msgs,
 										n * sizeof(SharedInvalidationMessage)),
 								 nmsgs += n));
+	ProcessMessageSubGroupMulti(&transInvalInfo->PriorCmdInvalidMsgs,
+								NspCacheMsgs,
+								(memcpy(msgarray + nmsgs,
+										msgs,
+										n * sizeof(SharedInvalidationMessage)),
+								 nmsgs += n));
+	ProcessMessageSubGroupMulti(&transInvalInfo->CurrentCmdInvalidMsgs,
+								NspCacheMsgs,
+								(memcpy(msgarray + nmsgs,
+										msgs,
+										n * sizeof(SharedInvalidationMessage)),
+								 nmsgs += n));
 	Assert(nmsgs == nummsgs);
 
 	return nmsgs;
@@ -1312,6 +1398,17 @@ CacheInvalidateHeapTuple(Relation relation,
 		else
 			return;
 	}
+	else if (tupleRelId == NamespaceRelationId)
+	{
+		Form_pg_namespace nsptup = (Form_pg_namespace) GETSTRUCT(tuple);
+
+		/* get namespace id */
+		relationId = nsptup->oid;
+		databaseId = MyDatabaseId;
+
+		RegisterNspcacheInvalidation(databaseId, relationId);
+		return;
+	}
 	else
 		return;
 
@@ -1567,6 +1664,25 @@ CacheRegisterRelcacheCallback(RelcacheCallbackFunction func,
 	++relcache_callback_count;
 }
 
+/*
+ * CacheRegisterNspcacheCallback
+ *		Register the specified function to be called for all future
+ *		namespace invalidation events.  The OID of the namespace being
+ *		invalidated will be passed to the function.
+ */
+void
+CacheRegisterNspcacheCallback(NspcacheCallbackFunction func,
+							  Datum arg)
+{
+	if (nspcache_callback_count >= MAX_NSPCACHE_CALLBACKS)
+		elog(FATAL, "out of nspcache_callback_list slots");
+
+	nspcache_callback_list[nspcache_callback_count].function = func;
+	nspcache_callback_list[nspcache_callback_count].arg = arg;
+
+	++nspcache_callback_count;
+}
+
 /*
  * CallSyscacheCallbacks
  *
@@ -1629,6 +1745,9 @@ LogLogicalInvalidations(void)
 		ProcessMessageSubGroupMulti(group, RelCacheMsgs,
 									XLogRegisterData((char *) msgs,
 													 n * sizeof(SharedInvalidationMessage)));
+		ProcessMessageSubGroupMulti(group, NspCacheMsgs,
+									XLogRegisterData((char *) msgs,
+													 n * sizeof(SharedInvalidationMessage)));
 		XLogInsert(RM_XACT_ID, XLOG_XACT_INVALIDATIONS);
 	}
 }
diff --git a/src/include/storage/sinval.h b/src/include/storage/sinval.h
index 8f5744b21b..4c53012528 100644
--- a/src/include/storage/sinval.h
+++ b/src/include/storage/sinval.h
@@ -110,6 +110,14 @@ typedef struct
 	Oid			relId;			/* relation ID */
 } SharedInvalSnapshotMsg;
 
+#define SHAREDINVALNSPCACHE_ID	(-6)
+typedef struct
+{
+	int8		id;				/* type field --- must be first */
+	Oid			dbId;			/* database ID, or 0 if a shared relation */
+	Oid			nspId;			/* namespace ID */
+} SharedInvalNspcacheMsg;
+
 typedef union
 {
 	int8		id;				/* type field --- must be first */
@@ -119,6 +127,7 @@ typedef union
 	SharedInvalSmgrMsg sm;
 	SharedInvalRelmapMsg rm;
 	SharedInvalSnapshotMsg sn;
+	SharedInvalNspcacheMsg nc;
 } SharedInvalidationMessage;
 
 
diff --git a/src/include/utils/inval.h b/src/include/utils/inval.h
index 24695facf2..99a0c90b6d 100644
--- a/src/include/utils/inval.h
+++ b/src/include/utils/inval.h
@@ -22,6 +22,7 @@ extern PGDLLIMPORT int debug_discard_caches;
 
 typedef void (*SyscacheCallbackFunction) (Datum arg, int cacheid, uint32 hashvalue);
 typedef void (*RelcacheCallbackFunction) (Datum arg, Oid relid);
+typedef void (*NspcacheCallbackFunction) (Datum arg, Oid nspid);
 
 
 extern void AcceptInvalidationMessages(void);
@@ -59,6 +60,9 @@ extern void CacheRegisterSyscacheCallback(int cacheid,
 extern void CacheRegisterRelcacheCallback(RelcacheCallbackFunction func,
 										  Datum arg);
 
+extern void CacheRegisterNspcacheCallback(NspcacheCallbackFunction func,
+										  Datum arg);
+
 extern void CallSyscacheCallbacks(int cacheid, uint32 hashvalue);
 
 extern void InvalidateSystemCaches(void);
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index 85d5c0d016..e038dd8a87 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -611,6 +611,97 @@ is( $result, qq(1
 	'Verify that the incremental data for table sch3.tab_conc added after table synchronization is replicated to the subscriber'
 );
 
+$background_psql1->query_safe(
+	qq[
+	BEGIN;
+	INSERT INTO tab_conc VALUES (4);
+]);
+
+# Maintain an active transaction with a schema table that will be added to the
+# publication.
+$background_psql2->query_safe(
+	qq[
+	BEGIN;
+	INSERT INTO sch3.tab_conc VALUES (4);
+]);
+
+$background_psql3->query_safe(
+	qq[
+	ALTER PUBLICATION regress_pub1 DROP TABLE tab_conc, TABLES IN SCHEMA sch3;
+]);
+
+$background_psql1->query_safe(qq[COMMIT]);
+$background_psql2->query_safe(qq[COMMIT]);
+
+$node_publisher->safe_psql(
+	'postgres', qq(
+	INSERT INTO tab_conc VALUES (5);
+	INSERT INTO sch3.tab_conc VALUES (5);
+));
+
+$node_publisher->wait_for_catchup('regress_sub1');
+
+# Verify that the insert is replicated to the subscriber.
+$result = $node_subscriber->safe_psql('postgres', "SELECT * FROM tab_conc");
+is( $result, qq(1
+2
+3
+4),
+	'Verify that the incremental data for table tab_conc added after table synchronization is replicated to the subscriber'
+);
+
+$result =
+  $node_subscriber->safe_psql('postgres', "SELECT * FROM sch3.tab_conc");
+is( $result, qq(1
+2
+3
+4),
+	'Verify that the incremental data for table sch3.tab_conc added after table synchronization is replicated to the subscriber'
+);
+
+$background_psql3->query_safe(
+	qq[
+	ALTER PUBLICATION regress_pub1 ADD TABLE tab_conc, TABLES IN SCHEMA sch3;
+]);
+
+$background_psql1->query_safe(
+	qq[
+	BEGIN;
+	INSERT INTO tab_conc VALUES (6);
+]);
+
+$background_psql2->query_safe(
+	qq[
+	BEGIN;
+	INSERT INTO sch3.tab_conc VALUES (6);
+]);
+
+$background_psql3->query_safe(
+	qq[
+	DROP PUBLICATION regress_pub1;
+]);
+
+$background_psql1->query_safe(
+	qq[
+	INSERT INTO tab_conc VALUES (7);
+]);
+
+$background_psql2->query_safe(
+	qq[
+	INSERT INTO sch3.tab_conc VALUES (7);
+]);
+
+$background_psql1->query_safe(qq[COMMIT]);
+$background_psql2->query_safe(qq[COMMIT]);
+
+my $offset = -s $node_subscriber->logfile;
+$node_subscriber->wait_for_log(
+	qr/ERROR:  publication "regress_pub1" does not exist/,
+	$offset);
+
+$node_subscriber->safe_psql('postgres',
+	'DROP SUBSCRIPTION regress_sub1;');
+
 $background_psql1->quit;
 $background_psql2->quit;
 $background_psql3->quit;
-- 
2.34.1