From 3c5a146ae188387d26401c4b487bb6e790169b7a Mon Sep 17 00:00:00 2001
From: Onder Kalaci <onderkalaci@gmail.com>
Date: Wed, 16 Feb 2022 16:20:16 +0100
Subject: [PATCH] Properly set worker_query and use

---
 .../distributed/transaction/backend_data.c    | 19 ++++++++-----------
 .../transaction/citus_dist_stat_activity.c    |  7 +++----
 src/include/distributed/backend_data.h        |  2 +-
 3 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/src/backend/distributed/transaction/backend_data.c b/src/backend/distributed/transaction/backend_data.c
index bc54e1da9..2467f115e 100644
--- a/src/backend/distributed/transaction/backend_data.c
+++ b/src/backend/distributed/transaction/backend_data.c
@@ -152,7 +152,6 @@ assign_distributed_transaction_id(PG_FUNCTION_ARGS)
 
 	MyBackendData->citusBackend.initiatorNodeIdentifier =
 		MyBackendData->transactionId.initiatorNodeIdentifier;
-	MyBackendData->citusBackend.transactionOriginator = false;
 
 	SpinLockRelease(&MyBackendData->mutex);
 
@@ -410,15 +409,12 @@ StoreAllActiveTransactions(Tuplestorestate *tupleStore, TupleDesc tupleDescripto
 		initiatorNodeIdentifier = currentBackend->citusBackend.initiatorNodeIdentifier;
 
 		/*
-		 * We prefer to use worker_query instead of transactionOriginator in the user facing
-		 * functions since its more intuitive. Thus, we negate the result before returning.
-		 *
-		 * We prefer to use citusBackend's transactionOriginator field over transactionId's
-		 * field with the same name. The reason is that it also covers backends that are not
-		 * inside a distributed transaction.
+		 * We prefer to use worker_query instead of distributedCommandOriginator in
+		 * the user facing functions since its more intuitive. Thus,
+		 * we negate the result before returning.
 		 */
 		bool coordinatorOriginatedQuery =
-			currentBackend->citusBackend.transactionOriginator;
+			currentBackend->distributedCommandOriginator;
 
 		transactionNumber = currentBackend->transactionId.transactionNumber;
 		TimestampTz transactionIdTimestamp = currentBackend->transactionId.timestamp;
@@ -663,7 +659,6 @@ UnSetDistributedTransactionId(void)
 		MyBackendData->transactionId.timestamp = 0;
 
 		MyBackendData->citusBackend.initiatorNodeIdentifier = -1;
-		MyBackendData->citusBackend.transactionOriginator = false;
 
 		SpinLockRelease(&MyBackendData->mutex);
 	}
@@ -682,6 +677,7 @@ UnSetGlobalPID(void)
 		SpinLockAcquire(&MyBackendData->mutex);
 
 		MyBackendData->globalPID = 0;
+		MyBackendData->distributedCommandOriginator = false;
 
 		SpinLockRelease(&MyBackendData->mutex);
 	}
@@ -776,7 +772,6 @@ AssignDistributedTransactionId(void)
 	MyBackendData->transactionId.timestamp = currentTimestamp;
 
 	MyBackendData->citusBackend.initiatorNodeIdentifier = localGroupId;
-	MyBackendData->citusBackend.transactionOriginator = true;
 
 	SpinLockRelease(&MyBackendData->mutex);
 }
@@ -798,7 +793,6 @@ MarkCitusInitiatedCoordinatorBackend(void)
 	SpinLockAcquire(&MyBackendData->mutex);
 
 	MyBackendData->citusBackend.initiatorNodeIdentifier = localGroupId;
-	MyBackendData->citusBackend.transactionOriginator = true;
 
 	SpinLockRelease(&MyBackendData->mutex);
 }
@@ -814,10 +808,12 @@ void
 AssignGlobalPID(void)
 {
 	uint64 globalPID = INVALID_CITUS_INTERNAL_BACKEND_GPID;
+	bool distributedCommandOriginator = false;
 
 	if (!IsCitusInternalBackend())
 	{
 		globalPID = GenerateGlobalPID();
+		distributedCommandOriginator = true;
 	}
 	else
 	{
@@ -826,6 +822,7 @@ AssignGlobalPID(void)
 
 	SpinLockAcquire(&MyBackendData->mutex);
 	MyBackendData->globalPID = globalPID;
+	MyBackendData->distributedCommandOriginator = distributedCommandOriginator;
 	SpinLockRelease(&MyBackendData->mutex);
 }
 
diff --git a/src/backend/distributed/transaction/citus_dist_stat_activity.c b/src/backend/distributed/transaction/citus_dist_stat_activity.c
index 3fb15295f..ec4f6e8a3 100644
--- a/src/backend/distributed/transaction/citus_dist_stat_activity.c
+++ b/src/backend/distributed/transaction/citus_dist_stat_activity.c
@@ -157,10 +157,9 @@ FROM \
 WHERE \
 	backend_type = 'client backend' \
 	AND \
-	pg_stat_activity.query NOT ILIKE '%stat_activity%' \
+	worker_query = False \
 	AND \
-	pg_stat_activity.application_name NOT SIMILAR TO 'citus_internal gpid=\\d+'; \
-"
+	pg_stat_activity.query NOT ILIKE '%stat_activity%';"
 
 #define CITUS_WORKER_STAT_ACTIVITY_QUERY \
 	"\
@@ -195,7 +194,7 @@ FROM \
 	get_all_active_transactions() AS dist_txs(database_id, process_id, initiator_node_identifier, worker_query, transaction_number, transaction_stamp, global_id) \
 	ON pg_stat_activity.pid = dist_txs.process_id \
 WHERE \
-	pg_stat_activity.application_name SIMILAR TO 'citus_internal gpid=\\d+' \
+	worker_query = True \
 	AND \
 	pg_stat_activity.query NOT ILIKE '%stat_activity%';"
 
diff --git a/src/include/distributed/backend_data.h b/src/include/distributed/backend_data.h
index 7f3a81e88..a4c93538b 100644
--- a/src/include/distributed/backend_data.h
+++ b/src/include/distributed/backend_data.h
@@ -29,7 +29,6 @@
 typedef struct CitusInitiatedBackend
 {
 	int initiatorNodeIdentifier;
-	bool transactionOriginator;
 } CitusInitiatedBackend;
 
 
@@ -51,6 +50,7 @@ typedef struct BackendData
 	slock_t mutex;
 	bool cancelledDueToDeadlock;
 	uint64 globalPID;
+	bool distributedCommandOriginator;
 	CitusInitiatedBackend citusBackend;
 	DistributedTransactionId transactionId;
 } BackendData;