From 125532e98abe2a024bc022880297353d55a0c4fa Mon Sep 17 00:00:00 2001
From: Gokhan Gulbiz <ggulbiz@gmail.com>
Date: Wed, 22 Mar 2023 11:03:13 +0300
Subject: [PATCH] temp - local query monitoring

---
 .../distributed/planner/deparse_shard_query.c |  2 +
 src/backend/distributed/utils/attribute.c     | 42 ++++++++++++-------
 src/include/distributed/utils/attribute.h     |  1 +
 3 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/src/backend/distributed/planner/deparse_shard_query.c b/src/backend/distributed/planner/deparse_shard_query.c
index 34650a377..4be978c78 100644
--- a/src/backend/distributed/planner/deparse_shard_query.c
+++ b/src/backend/distributed/planner/deparse_shard_query.c
@@ -400,6 +400,8 @@ SetTaskQueryIfShouldLazyDeparse(Task *task, Query *query)
 		task->taskQuery.queryType = TASK_QUERY_OBJECT;
 		task->taskQuery.data.jobQueryReferenceForLazyDeparsing = query;
 		task->queryCount = 1;
+		AttributeTask(task->partitionColumn, task->colocationId, query->commandType);
+
 		return;
 	}
 
diff --git a/src/backend/distributed/utils/attribute.c b/src/backend/distributed/utils/attribute.c
index 30d070c0a..7b94ac769 100644
--- a/src/backend/distributed/utils/attribute.c
+++ b/src/backend/distributed/utils/attribute.c
@@ -23,7 +23,8 @@
 #include "utils/builtins.h"
 #include "utils/json.h"
 #include "distributed/utils/attribute.h"
-#include "common/base64.h"
+#include "miscadmin.h"
+
 
 #include <time.h>
 
@@ -172,8 +173,6 @@ AttributeQueryIfAnnotated(const char *query_string, CmdType commandType)
 {
 	strcpy_s(attributeToTenant, sizeof(attributeToTenant), "");
 
-	attributeCommandType = commandType;
-
 	if (query_string == NULL)
 	{
 		return;
@@ -187,28 +186,43 @@ AttributeQueryIfAnnotated(const char *query_string, CmdType commandType)
 			Datum jsonbDatum = DirectFunctionCall1(jsonb_in, PointerGetDatum(annotation));
 
 			text *tenantIdTextP = ExtractFieldTextP(jsonbDatum, "tId");
+			char *tenantId = NULL;
 			if (tenantIdTextP != NULL)
 			{
-				char *tenantId = UnescapeCommentChars(text_to_cstring(tenantIdTextP));
-				strcpy_s(attributeToTenant, sizeof(attributeToTenant), tenantId);
+				tenantId = UnescapeCommentChars(text_to_cstring(tenantIdTextP));
 			}
 
-			colocationGroupId = ExtractFieldInt32(jsonbDatum, "cId", 0);
+			int colocationId = ExtractFieldInt32(jsonbDatum, "cId",
+												  0);
 
-			if (MultiTenantMonitoringLogLevel != CITUS_LOG_LEVEL_OFF)
-			{
-				ereport(NOTICE, (errmsg(
-									 "attributing query to tenant: %s, colocationGroupId: %d",
-									 quote_literal_cstr(attributeToTenant),
-									 colocationGroupId)));
-			}
+			AttributeTask(tenantId, colocationId, commandType);
 		}
 	}
+}
+
+/*
+ * AttributeTask assigns the given attributes of a tenant and starts a timer
+ */
+void AttributeTask(char *tenantId, int colocationId, CmdType commandType)
+{
+	ereport(NOTICE, (errmsg("MyProcPid: %d", MyProcPid)));
+	sleep(10);
+
+	colocationGroupId = colocationId;
+	strcpy_s(attributeToTenant, sizeof(attributeToTenant), tenantId);
+	attributeCommandType = commandType;
+
+	if (MultiTenantMonitoringLogLevel != CITUS_LOG_LEVEL_OFF)
+	{
+		ereport(NOTICE, (errmsg(
+								"attributing query to tenant: %s, colocationGroupId: %d",
+								quote_literal_cstr(attributeToTenant),
+								colocationGroupId)));
+	}
 
 	attributeToTenantStart = clock();
 }
 
-
 /*
  * AnnotateQuery annotates the query with tenant attributes.
  */
diff --git a/src/include/distributed/utils/attribute.h b/src/include/distributed/utils/attribute.h
index c05ec950e..f2453b818 100644
--- a/src/include/distributed/utils/attribute.h
+++ b/src/include/distributed/utils/attribute.h
@@ -56,6 +56,7 @@ extern void CitusAttributeToEnd(QueryDesc *queryDesc);
 extern void AttributeQueryIfAnnotated(const char *queryString, CmdType commandType);
 extern char * AnnotateQuery(char *queryString, char *partitionColumn, int colocationId);
 extern void InitializeMultiTenantMonitorSMHandleManagement(void);
+extern void AttributeTask(char *tenantId, int colocationGroupId, CmdType commandType);
 
 extern ExecutorEnd_hook_type prev_ExecutorEnd;