Avoid re-assigning the global pid for client backends and bg workers when the application_name changes

src/backend/distributed/metadata/metadata_cache.c

@@ -4545,6 +4545,16 @@ GetLocalNodeId(void)
 }
 
 
+/*
+ * CachedLocalNodeIdIsValid return true if the cached local node id is valid.
+ */
+bool
+CachedLocalNodeIdIsValid(void)
+{
+	return LocalNodeId != -1;
+}
+
+
 /*
  * RegisterLocalGroupIdCacheCallbacks registers the callbacks required to
  * maintain LocalGroupId at a consistent value. It's separate from

src/backend/distributed/shared_library_init.c

@@ -2899,6 +2899,18 @@
 	 * So we set the FinishedStartupCitusBackend flag in StartupCitusBackend to
 	 * indicate when this responsibility handoff has happened.
 	 *
+	 * On the other hand, even if now it's this hook's responsibility to update
+	 * the global pid, we cannot do so if we might need to read from catalog
+	 * but it's unsafe to do so. For Citus internal backends, this cannot be the
+	 * case because in that case AssignGlobalPID() just extracts its global pid
+	 * from the application_name and extracting doesn't require catalog access.
+	 * But for external client backends, we either need to guarantee that we won't
+	 * read from catalog tables or that it's safe to do so. The only case where
+	 * AssignGlobalPID() could read from catalog tables is when the cached local
+	 * node id is invalidated. So for this reason, for external client backends,
+	 * we require that either the cached local node id is valid or that we are in
+	 * a transaction block -because in that case it's safe to read from catalog.
+	 *
 	 * Another solution to the catalog table acccess problem would be to update
 	 * global pid lazily, like we do for HideShards. But that's not possible
 	 * for the global pid, since it is stored in shared memory instead of in a
@@ -2907,7 +2919,9 @@
 	 * as reasonably possible, which is also why we extract global pids in the
 	 * AuthHook already (extracting doesn't require catalog access).
 	 */
-	if (FinishedStartupCitusBackend)
+	if (FinishedStartupCitusBackend &&
+		(!IsExternalClientBackend() || CachedLocalNodeIdIsValid() ||
+		 IsTransactionState()))
 	{
 		AssignGlobalPID(newval);
 	}

src/include/distributed/metadata_cache.h

@@ -181,6 +181,7 @@ extern CitusTableCacheEntry * LookupCitusTableCacheEntry(Oid relationId);
 extern DistObjectCacheEntry * LookupDistObjectCacheEntry(Oid classid, Oid objid, int32
 														 objsubid);
 extern int32 GetLocalGroupId(void);
+extern bool CachedLocalNodeIdIsValid(void);
 extern int32 GetLocalNodeId(void);
 extern void CitusTableCacheFlushInvalidatedEntries(void);
 extern Oid LookupShardRelationFromCatalog(int64 shardId, bool missing_ok);

src/test/regress/expected/remove_coordinator.out

@@ -5,6 +5,35 @@ SELECT master_remove_node('localhost', :master_port);
 
 (1 row)
 
+-- to silence -potentially flaky- "could not establish connection after" warnings in below test
+SET client_min_messages TO ERROR;
+-- to fail fast when the hostname is not resolvable, as it will be the case below
+SET citus.node_connection_timeout to '1s';
+BEGIN;
+  SET application_name TO 'new_app_name';
+  -- that should fail because of bad hostname & port
+  SELECT citus_add_node('200.200.200.200', 1, 200);
+ERROR:  connection to the remote node [email protected]:1 failed
+  -- Since above command failed, now Postgres will need to revert the
+  -- application_name change made in this transaction and this will
+  -- happen within abort-transaction callback, so we won't be in a
+  -- transaction block while Postgres does that.
+  --
+  -- And when the application_name changes, Citus tries to re-assign
+  -- the global pid and doing so for Citus internal backends doesn't
+  -- require being in a transaction block and is safe.
+  --
+  -- However, for the client external backends (like us here), Citus
+  -- doesn't try to re-assign the global pid if doing so requires catalog
+  -- access and we're outside of a transaction block. Note that in that
+  -- case the catalog access may only be needed to retrive the local node
+  -- id when the cached local node is invalidated like what just happened
+  -- here because of the failed citus_add_node() call made above.
+  --
+  -- So by failing here (rather than crashing), we ensure this behavior.
+ROLLBACK;
+RESET client_min_messages;
+RESET citus.node_connection_timeout;
 -- restore coordinator for the rest of the tests
 SELECT citus_set_coordinator_host('localhost', :master_port);
  citus_set_coordinator_host

src/test/regress/sql/remove_coordinator.sql

@@ -1,5 +1,39 @@
 -- removing coordinator from pg_dist_node should update pg_dist_colocation
 SELECT master_remove_node('localhost', :master_port);
 
+-- to silence -potentially flaky- "could not establish connection after" warnings in below test
+SET client_min_messages TO ERROR;
+
+-- to fail fast when the hostname is not resolvable, as it will be the case below
+SET citus.node_connection_timeout to '1s';
+
+BEGIN;
+  SET application_name TO 'new_app_name';
+
+  -- that should fail because of bad hostname & port
+  SELECT citus_add_node('200.200.200.200', 1, 200);
+
+  -- Since above command failed, now Postgres will need to revert the
+  -- application_name change made in this transaction and this will
+  -- happen within abort-transaction callback, so we won't be in a
+  -- transaction block while Postgres does that.
+  --
+  -- And when the application_name changes, Citus tries to re-assign
+  -- the global pid and doing so for Citus internal backends doesn't
+  -- require being in a transaction block and is safe.
+  --
+  -- However, for the client external backends (like us here), Citus
+  -- doesn't try to re-assign the global pid if doing so requires catalog
+  -- access and we're outside of a transaction block. Note that in that
+  -- case the catalog access may only be needed to retrive the local node
+  -- id when the cached local node is invalidated like what just happened
+  -- here because of the failed citus_add_node() call made above.
+  --
+  -- So by failing here (rather than crashing), we ensure this behavior.
+ROLLBACK;
+
+RESET client_min_messages;
+RESET citus.node_connection_timeout;
+
 -- restore coordinator for the rest of the tests
 SELECT citus_set_coordinator_host('localhost', :master_port);