Return early in CitusHasBeenLoaded when creating a different ex… (#3178)

Return early in CitusHasBeenLoaded when creating a different extension

src/backend/distributed/master/master_delete_protocol.c

@@ -270,6 +270,12 @@ master_drop_sequences(PG_FUNCTION_ARGS)
 	bool isNull = false;
 	StringInfo dropSeqCommand = makeStringInfo();
 
+	if (!CitusHasBeenLoaded())
+	{
+		/* ignore calls during CREATE EXTENSION citus */
+		PG_RETURN_VOID();
+	}
+
 	CheckCitusVersion(ERROR);
 
 	/*

src/backend/distributed/utils/metadata_cache.c

@@ -202,6 +202,7 @@ static char * InstalledExtensionVersion(void);
 static bool HasOverlappingShardInterval(ShardInterval **shardIntervalArray,
 										int shardIntervalArrayLength,
 										FmgrInfo *shardIntervalSortCompareFunction);
+static bool CitusHasBeenLoadedInternal(void);
 static void InitializeCaches(void);
 static void InitializeDistCache(void);
 static void InitializeDistObjectCache(void);
@@ -1528,41 +1529,27 @@ HasOverlappingShardInterval(ShardInterval **shardIntervalArray,
 bool
 CitusHasBeenLoaded(void)
 {
-	/* recheck presence until citus has been loaded */
 	if (!MetadataCache.extensionLoaded || creating_extension)
 	{
-		bool extensionPresent = false;
-		bool extensionScriptExecuted = true;
+		/*
+		 * Refresh if we have not determined whether the extension has been
+		 * loaded yet, or in case of ALTER EXTENSION since we want to treat
+		 * Citus as "not loaded" during ALTER EXTENSION citus.
+		 */
+		bool extensionLoaded = CitusHasBeenLoadedInternal();
 
-		Oid extensionOid = get_extension_oid("citus", true);
-		if (extensionOid != InvalidOid)
+		if (extensionLoaded && !MetadataCache.extensionLoaded)
 		{
-			extensionPresent = true;
-		}
-
-		if (extensionPresent)
-		{
-			/* check if Citus extension objects are still being created */
-			if (creating_extension && CurrentExtensionObject == extensionOid)
-			{
-				extensionScriptExecuted = false;
-			}
+			/*
+			 * Loaded Citus for the first time in this session, or first time after
+			 * CREATE/ALTER EXTENSION citus. Do some initialisation.
+			 */
 
 			/*
-			 * Whenever the extension exists, even when currently creating it,
-			 * we need the infrastructure to run citus in this database to be
-			 * ready.
+			 * Make sure the maintenance daemon is running if it was not already.
 			 */
 			StartupCitusBackend();
-		}
 
-		/* we disable extension features during pg_upgrade */
-		MetadataCache.extensionLoaded = extensionPresent &&
-										extensionScriptExecuted &&
-										!IsBinaryUpgrade;
-
-		if (MetadataCache.extensionLoaded)
-		{
 			/*
 			 * InvalidateDistRelationCacheCallback resets state such as extensionLoaded
 			 * when it notices changes to pg_dist_partition (which usually indicate
@@ -1576,26 +1563,58 @@ CitusHasBeenLoaded(void)
 			 */
 			DistPartitionRelationId();
 
-
 			/*
 			 * This needs to be initialized so we can receive foreign relation graph
 			 * invalidation messages in InvalidateForeignRelationGraphCacheCallback().
 			 * See the comments of InvalidateForeignKeyGraph for more context.
 			 */
 			DistColocationRelationId();
-
-			/*
-			 * We also reset citusVersionKnownCompatible, so it will be re-read in
-			 * case of extension update.
-			 */
-			citusVersionKnownCompatible = false;
 		}
+
+		MetadataCache.extensionLoaded = extensionLoaded;
 	}
 
 	return MetadataCache.extensionLoaded;
 }
 
 
+/*
+ * CitusHasBeenLoadedInternal returns true if the citus extension has been created
+ * in the current database and the extension script has been executed. Otherwise,
+ * it returns false.
+ */
+static bool
+CitusHasBeenLoadedInternal(void)
+{
+	Oid citusExtensionOid = InvalidOid;
+
+	if (IsBinaryUpgrade)
+	{
+		/* never use Citus logic during pg_upgrade */
+		return false;
+	}
+
+	citusExtensionOid = get_extension_oid("citus", true);
+	if (citusExtensionOid == InvalidOid)
+	{
+		/* Citus extension does not exist yet */
+		return false;
+	}
+
+	if (creating_extension && CurrentExtensionObject == citusExtensionOid)
+	{
+		/*
+		 * We do not use Citus hooks during CREATE/ALTER EXTENSION citus
+		 * since the objects used by the C code might be not be there yet.
+		 */
+		return false;
+	}
+
+	/* citus extension exists and has been created */
+	return true;
+}
+
+
 /*
  * CheckCitusVersion checks whether there is a version mismatch between the
  * available version and the loaded version or between the installed version

src/test/regress/expected/isolation_dump_global_wait_edges.out

@@ -29,11 +29,11 @@ step detector-dump-wait-edges:
 
 waiting_transaction_numblocking_transaction_numblocking_transaction_waiting
 
-290            289            f              
+289            288            f              
 transactionnumberwaitingtransactionnumbers
 
-289                           
-290            289            
+288                           
+289            288            
 step s1-abort: 
     ABORT;
 
@@ -77,14 +77,14 @@ step detector-dump-wait-edges:
 
 waiting_transaction_numblocking_transaction_numblocking_transaction_waiting
 
-294            293            f              
-295            293            f              
-295            294            t              
+293            292            f              
+294            292            f              
+294            293            t              
 transactionnumberwaitingtransactionnumbers
 
-293                           
-294            293            
-295            293,294        
+292                           
+293            292            
+294            292,293        
 step s1-abort: 
     ABORT;