Multi tenant monitoring

src/backend/distributed/executor/citus_custom_scan.c

@@ -66,7 +66,6 @@ static DistributedPlan * CopyDistributedPlanWithoutCache(
 	DistributedPlan *originalDistributedPlan);
 static void CitusEndScan(CustomScanState *node);
 static void CitusReScan(CustomScanState *node);
-static void SetJobColocationId(Job *job);
 static void EnsureForceDelegationDistributionKey(Job *job);
 static void EnsureAnchorShardsInJobExist(Job *job);
 static bool AnchorShardsInTaskListExist(List *taskList);
@@ -892,7 +891,7 @@ IsCitusCustomScan(Plan *plan)
  * colocation group, the Job's colocation ID is set to the group ID, else,
  * it will be set to INVALID_COLOCATION_ID.
  */
-static void
+void
 SetJobColocationId(Job *job)
 {
 	uint32 jobColocationId = INVALID_COLOCATION_ID;

src/backend/distributed/executor/local_executor.c

@@ -97,6 +97,7 @@
 #include "distributed/relation_access_tracking.h"
 #include "distributed/remote_commands.h" /* to access LogRemoteCommands */
 #include "distributed/transaction_management.h"
+#include "distributed/utils/citus_stat_tenants.h"
 #include "distributed/version_compat.h"
 #include "distributed/worker_protocol.h"
 #include "executor/tstoreReceiver.h"
@@ -647,6 +648,16 @@ LocallyExecuteTaskPlan(PlannedStmt *taskPlan, char *queryString,
 		LocalExecutorShardId = task->anchorShardId;
 	}
 
+
+	char *partitionKeyValueString = NULL;
+	if (task->partitionKeyValue != NULL)
+	{
+		partitionKeyValueString = DatumToString(task->partitionKeyValue->constvalue,
+												task->partitionKeyValue->consttype);
+	}
+
+	AttributeTask(partitionKeyValueString, task->colocationId, taskPlan->commandType);
+
 	PG_TRY();
 	{
 		processedRows = ExecuteTaskPlan(taskPlan, queryString, tupleDest, task,

src/backend/distributed/planner/deparse_shard_query.c

@@ -26,6 +26,7 @@
 #include "distributed/multi_physical_planner.h"
 #include "distributed/multi_router_planner.h"
 #include "distributed/shard_utils.h"
+#include "distributed/utils/citus_stat_tenants.h"
 #include "distributed/version_compat.h"
 #include "lib/stringinfo.h"
 #include "nodes/makefuncs.h"
@@ -141,6 +142,10 @@ RebuildQueryStrings(Job *workerJob)
 								? "(null)"
 								: TaskQueryString(task))));
 
+		task->partitionKeyValue = workerJob->partitionKeyValue;
+		SetJobColocationId(workerJob);
+		task->colocationId = workerJob->colocationId;
+
 		UpdateTaskQueryString(query, task);
 
 		/*
@@ -387,7 +392,8 @@ SetTaskQueryIfShouldLazyDeparse(Task *task, Query *query)
 		return;
 	}
 
-	SetTaskQueryString(task, DeparseTaskQuery(task, query));
+	SetTaskQueryString(task, AnnotateQuery(DeparseTaskQuery(task, query),
+										   task->partitionKeyValue, task->colocationId));
 }
 
 

src/backend/distributed/planner/distributed_planner.c

@@ -49,6 +49,7 @@
 #include "distributed/recursive_planning.h"
 #include "distributed/shardinterval_utils.h"
 #include "distributed/shard_utils.h"
+#include "distributed/utils/citus_stat_tenants.h"
 #include "distributed/version_compat.h"
 #include "distributed/worker_shard_visibility.h"
 #include "executor/executor.h"
@@ -305,6 +306,11 @@ distributed_planner(Query *parse,
 						errhint("Consider using PL/pgSQL functions instead.")));
 	}
 
+	/*
+	 * We annotate the query for tenant statisisics.
+	 */
+	AttributeQueryIfAnnotated(query_string, parse->commandType);
+
 	return result;
 }
 

src/backend/distributed/planner/multi_router_planner.c

@@ -164,7 +164,8 @@ static int CompareInsertValuesByShardId(const void *leftElement,
 static List * SingleShardTaskList(Query *query, uint64 jobId,
 								  List *relationShardList, List *placementList,
 								  uint64 shardId, bool parametersInQueryResolved,
-								  bool isLocalTableModification);
+								  bool isLocalTableModification, Const *partitionKeyValue,
+								  int colocationId);
 static bool RowLocksOnRelations(Node *node, List **rtiLockList);
 static void ReorderTaskPlacementsByTaskAssignmentPolicy(Job *job,
 														TaskAssignmentPolicyType
@@ -1939,11 +1940,14 @@ GenerateSingleShardRouterTaskList(Job *job, List *relationShardList,
 
 	if (originalQuery->commandType == CMD_SELECT)
 	{
+		SetJobColocationId(job);
+
 		job->taskList = SingleShardTaskList(originalQuery, job->jobId,
 											relationShardList, placementList,
 											shardId,
 											job->parametersInJobQueryResolved,
-											isLocalTableModification);
+											isLocalTableModification,
+											job->partitionKeyValue, job->colocationId);
 
 		/*
 		 * Queries to reference tables, or distributed tables with multiple replica's have
@@ -1967,11 +1971,14 @@ GenerateSingleShardRouterTaskList(Job *job, List *relationShardList,
 	}
 	else
 	{
+		SetJobColocationId(job);
+
 		job->taskList = SingleShardTaskList(originalQuery, job->jobId,
 											relationShardList, placementList,
 											shardId,
 											job->parametersInJobQueryResolved,
-											isLocalTableModification);
+											isLocalTableModification,
+											job->partitionKeyValue, job->colocationId);
 	}
 }
 
@@ -2065,7 +2072,8 @@ static List *
 SingleShardTaskList(Query *query, uint64 jobId, List *relationShardList,
 					List *placementList, uint64 shardId,
 					bool parametersInQueryResolved,
-					bool isLocalTableModification)
+					bool isLocalTableModification, Const *partitionKeyValue,
+					int colocationId)
 {
 	TaskType taskType = READ_TASK;
 	char replicationModel = 0;
@@ -2135,6 +2143,8 @@ SingleShardTaskList(Query *query, uint64 jobId, List *relationShardList,
 	 * that the query cannot be executed locally.
 	 */
 	task->taskPlacementList = placementList;
+	task->partitionKeyValue = partitionKeyValue;
+	task->colocationId = colocationId;
 	SetTaskQueryIfShouldLazyDeparse(task, query);
 	task->anchorShardId = shardId;
 	task->jobId = jobId;

src/backend/distributed/shared_library_init.c

@@ -91,6 +91,7 @@
 #include "distributed/resource_lock.h"
 #include "distributed/transaction_management.h"
 #include "distributed/transaction_recovery.h"
+#include "distributed/utils/citus_stat_tenants.h"
 #include "distributed/utils/directory.h"
 #include "distributed/worker_log_messages.h"
 #include "distributed/worker_manager.h"
@@ -228,6 +229,12 @@ static const struct config_enum_entry stat_statements_track_options[] = {
 	{ NULL, 0, false }
 };
 
+static const struct config_enum_entry stat_tenants_track_options[] = {
+	{ "none", STAT_TENANTS_TRACK_NONE, false },
+	{ "all", STAT_TENANTS_TRACK_ALL, false },
+	{ NULL, 0, false }
+};
+
 static const struct config_enum_entry task_assignment_policy_options[] = {
 	{ "greedy", TASK_ASSIGNMENT_GREEDY, false },
 	{ "first-replica", TASK_ASSIGNMENT_FIRST_REPLICA, false },
@@ -447,6 +454,8 @@ _PG_init(void)
 	ExecutorStart_hook = CitusExecutorStart;
 	ExecutorRun_hook = CitusExecutorRun;
 	ExplainOneQuery_hook = CitusExplainOneQuery;
+	prev_ExecutorEnd = ExecutorEnd_hook;
+	ExecutorEnd_hook = CitusAttributeToEnd;
 
 	/* register hook for error messages */
 	emit_log_hook = multi_log_hook;
@@ -491,6 +500,8 @@ _PG_init(void)
 	/* initialize shard split shared memory handle management */
 	InitializeShardSplitSMHandleManagement();
 
+	InitializeMultiTenantMonitorSMHandleManagement();
+
 	/* enable modification of pg_catalog tables during pg_upgrade */
 	if (IsBinaryUpgrade)
 	{
@@ -2363,6 +2374,50 @@ RegisterCitusConfigVariables(void)
 		GUC_STANDARD,
 		NULL, NULL, NULL);
 
+	DefineCustomIntVariable(
+		"citus.stat_tenants_limit",
+		gettext_noop("Number of tenants to be shown in citus_stat_tenants."),
+		NULL,
+		&StatTenantsLimit,
+		100, 1, 10000,
+		PGC_POSTMASTER,
+		GUC_STANDARD,
+		NULL, NULL, NULL);
+
+	DefineCustomEnumVariable(
+		"citus.stat_tenants_log_level",
+		gettext_noop("Sets the level of citus_stat_tenants log messages"),
+		NULL,
+		&StatTenantsLogLevel,
+		CITUS_LOG_LEVEL_OFF, log_level_options,
+		PGC_USERSET,
+		GUC_STANDARD,
+		NULL, NULL, NULL);
+
+	DefineCustomIntVariable(
+		"citus.stat_tenants_period",
+		gettext_noop("Period in seconds to be used for calculating the tenant "
+					 "statistics in citus_stat_tenants."),
+		NULL,
+		&StatTenantsPeriod,
+		60, 1, 60 * 60 * 24,
+		PGC_USERSET,
+		GUC_STANDARD,
+		NULL, NULL, NULL);
+
+	DefineCustomEnumVariable(
+		"citus.stat_tenants_track",
+		gettext_noop("Enables/Disables the stats collection for citus_stat_tenants."),
+		gettext_noop("Enables the stats collection when set to 'all'. "
+					 "Disables when set to 'none'. Disabling can be useful for "
+					 "avoiding extra CPU cycles needed for the calculations."),
+		&StatTenantsTrack,
+		STAT_TENANTS_TRACK_NONE,
+		stat_tenants_track_options,
+		PGC_SUSET,
+		GUC_STANDARD,
+		NULL, NULL, NULL);
+
 	DefineCustomBoolVariable(
 		"citus.subquery_pushdown",
 		gettext_noop("Usage of this GUC is highly discouraged, please read the long "

src/backend/distributed/sql/citus--11.2-1--11.3-1.sql

@@ -10,3 +10,8 @@ ALTER TABLE pg_catalog.pg_dist_transaction REPLICA IDENTITY USING INDEX pg_dist_
 
 #include "udfs/worker_drop_all_shell_tables/11.3-1.sql"
 #include "udfs/citus_internal_mark_node_not_synced/11.3-1.sql"
+#include "udfs/citus_stat_tenants_local/11.3-1.sql"
+#include "udfs/citus_stat_tenants/11.3-1.sql"
+
+#include "udfs/citus_stat_tenants_local_reset/11.3-1.sql"
+#include "udfs/citus_stat_tenants_reset/11.3-1.sql"

src/backend/distributed/sql/downgrades/citus--11.3-1--11.2-1.sql

@@ -20,3 +20,12 @@ ALTER TABLE pg_catalog.pg_dist_transaction REPLICA IDENTITY NOTHING;
 
 DROP PROCEDURE pg_catalog.worker_drop_all_shell_tables(bool);
 DROP FUNCTION pg_catalog.citus_internal_mark_node_not_synced(int, int);
+
+DROP VIEW pg_catalog.citus_stat_tenants_local;
+DROP FUNCTION pg_catalog.citus_stat_tenants_local(boolean);
+
+DROP VIEW pg_catalog.citus_stat_tenants;
+DROP FUNCTION pg_catalog.citus_stat_tenants(boolean);
+
+DROP FUNCTION pg_catalog.citus_stat_tenants_local_reset();
+DROP FUNCTION pg_catalog.citus_stat_tenants_reset();