Citus stats tenants guc (#6769)

Adds a GUC for turning tenant monitoring on and off

src/backend/distributed/shared_library_init.c

@@ -229,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 },
@@ -2378,6 +2384,17 @@ RegisterCitusConfigVariables(void)
 		GUC_STANDARD,
 		NULL, NULL, NULL);
 
+	DefineCustomEnumVariable(
+		"citus.stat_tenants_track",
+		gettext_noop("enable disable"),
+		NULL,
+		&StatTenantsTrack,
+		STAT_TENANTS_TRACK_ALL,
+		stat_tenants_track_options,
+		PGC_SUSET,
+		GUC_STANDARD,
+		NULL, NULL, NULL);
+
 	DefineCustomIntVariable(
 		"citus.stats_tenants_limit",
 		gettext_noop("monitor limit"),

src/backend/distributed/utils/attribute.c

@@ -72,6 +72,7 @@ static char * UnescapeCommentChars(const char *str);
 int MultiTenantMonitoringLogLevel = CITUS_LOG_LEVEL_OFF;
 int CitusStatsTenantsPeriod = (time_t) 60;
 int CitusStatsTenantsLimit = 10;
+int StatTenantsTrack = STAT_TENANTS_TRACK_ALL;
 
 
 PG_FUNCTION_INFO_V1(citus_stats_tenants_local);
@@ -195,6 +196,11 @@ sleep_until_next_period(PG_FUNCTION_ARGS)
 void
 AttributeQueryIfAnnotated(const char *query_string, CmdType commandType)
 {
+	if (StatTenantsTrack == STAT_TENANTS_TRACK_NONE)
+	{
+		return;
+	}
+
 	strcpy_s(attributeToTenant, sizeof(attributeToTenant), "");
 
 	if (query_string == NULL)
@@ -231,7 +237,8 @@ AttributeQueryIfAnnotated(const char *query_string, CmdType commandType)
 void
 AttributeTask(char *tenantId, int colocationId, CmdType commandType)
 {
-	if (tenantId == NULL || colocationId == INVALID_COLOCATION_ID)
+	if (StatTenantsTrack == STAT_TENANTS_TRACK_NONE ||
+		tenantId == NULL || colocationId == INVALID_COLOCATION_ID)
 	{
 		return;
 	}
@@ -249,7 +256,7 @@ AttributeTask(char *tenantId, int colocationId, CmdType commandType)
 char *
 AnnotateQuery(char *queryString, char *partitionColumn, int colocationId)
 {
-	if (partitionColumn == NULL)
+	if (StatTenantsTrack == STAT_TENANTS_TRACK_NONE || partitionColumn == NULL)
 	{
 		return queryString;
 	}
@@ -322,6 +329,11 @@ CompareTenantScore(const void *leftElement, const void *rightElement)
 static void
 AttributeMetricsIfApplicable()
 {
+	if (StatTenantsTrack == STAT_TENANTS_TRACK_NONE)
+	{
+		return;
+	}
+
 	if (strcmp(attributeToTenant, "") != 0)
 	{
 		time_t queryTime = time(0);

src/include/distributed/utils/attribute.h

@@ -87,6 +87,11 @@ typedef struct MultiTenantMonitor
 	TenantStats tenants[FLEXIBLE_ARRAY_MEMBER];
 } MultiTenantMonitor;
 
+typedef enum
+{
+	STAT_TENANTS_TRACK_NONE = 0,
+	STAT_TENANTS_TRACK_ALL = 1
+} StatTenantsTrackType;
 
 extern void CitusAttributeToEnd(QueryDesc *queryDesc);
 extern void AttributeQueryIfAnnotated(const char *queryString, CmdType commandType);
@@ -99,5 +104,6 @@ extern ExecutorEnd_hook_type prev_ExecutorEnd;
 extern int MultiTenantMonitoringLogLevel;
 extern int CitusStatsTenantsPeriod;
 extern int CitusStatsTenantsLimit;
+extern int StatTenantsTrack;
 
 #endif /*CITUS_ATTRIBUTE_H */

src/test/regress/expected/citus_stats_tenants.out

@@ -314,6 +314,48 @@ CONTEXT:  PL/pgSQL function citus_stats_tenants(boolean) line XX at RAISE
  t
 (1 row)
 
+RESET client_min_messages;
+SELECT result FROM run_command_on_all_nodes('SELECT clean_citus_stats_tenants()');
+ result
+---------------------------------------------------------------------
+
+
+
+(3 rows)
+
+-- test turning monitoring on/off
+SET citus.stat_tenants_track TO "NONE";
+SELECT count(*)>=0 FROM dist_tbl WHERE a = 1;
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+INSERT INTO dist_tbl VALUES (1, 1);
+SELECT tenant_attribute, query_count_in_this_period FROM citus_stats_tenants;
+ tenant_attribute | query_count_in_this_period
+---------------------------------------------------------------------
+(0 rows)
+
+SET citus.stat_tenants_track TO "ALL";
+SELECT tenant_attribute, query_count_in_this_period FROM citus_stats_tenants;
+ tenant_attribute | query_count_in_this_period
+---------------------------------------------------------------------
+(0 rows)
+
+SELECT count(*)>=0 FROM dist_tbl WHERE a = 1;
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+INSERT INTO dist_tbl VALUES (1, 1);
+SELECT tenant_attribute, query_count_in_this_period FROM citus_stats_tenants;
+ tenant_attribute | query_count_in_this_period
+---------------------------------------------------------------------
+ 1                |                          2
+(1 row)
+
 -- test special and multibyte characters in tenant attribute
 SELECT result FROM run_command_on_all_nodes('SELECT clean_citus_stats_tenants()');
  result

src/test/regress/sql/citus_stats_tenants.sql

@@ -110,6 +110,25 @@ SET citus.multi_tenant_monitoring_log_level TO LOG;
 SELECT count(*)>=0 FROM citus_stats_tenants;
 SET citus.multi_tenant_monitoring_log_level TO DEBUG;
 SELECT count(*)>=0 FROM citus_stats_tenants;
+RESET client_min_messages;
+
+SELECT result FROM run_command_on_all_nodes('SELECT clean_citus_stats_tenants()');
+
+-- test turning monitoring on/off
+SET citus.stat_tenants_track TO "NONE";
+SELECT count(*)>=0 FROM dist_tbl WHERE a = 1;
+INSERT INTO dist_tbl VALUES (1, 1);
+
+SELECT tenant_attribute, query_count_in_this_period FROM citus_stats_tenants;
+
+SET citus.stat_tenants_track TO "ALL";
+
+SELECT tenant_attribute, query_count_in_this_period FROM citus_stats_tenants;
+
+SELECT count(*)>=0 FROM dist_tbl WHERE a = 1;
+INSERT INTO dist_tbl VALUES (1, 1);
+
+SELECT tenant_attribute, query_count_in_this_period FROM citus_stats_tenants;
 
 -- test special and multibyte characters in tenant attribute
 SELECT result FROM run_command_on_all_nodes('SELECT clean_citus_stats_tenants()');