Add CPU usage to citus_stat_tenants (#6844)

This PR adds CPU usage to `citus_stat_tenants` monitor.
CPU usage is tracked in periods, similar to query counts.

src/backend/distributed/sql/udfs/citus_stat_tenants/latest.sql

@@ -8,6 +8,8 @@ CREATE OR REPLACE FUNCTION pg_catalog.citus_stat_tenants (
     OUT read_count_in_last_period INT,
     OUT query_count_in_this_period INT,
     OUT query_count_in_last_period INT,
+    OUT cpu_usage_in_this_period DOUBLE PRECISION,
+    OUT cpu_usage_in_last_period DOUBLE PRECISION,
     OUT score BIGINT
 )
     RETURNS SETOF record
@@ -51,6 +53,8 @@ AS (
     read_count_in_last_period INT,
     query_count_in_this_period INT,
     query_count_in_last_period INT,
+    cpu_usage_in_this_period DOUBLE PRECISION,
+    cpu_usage_in_last_period DOUBLE PRECISION,
     score BIGINT
 )
     ORDER BY score DESC
@@ -66,7 +70,9 @@ SELECT
     read_count_in_this_period,
     read_count_in_last_period,
     query_count_in_this_period,
-    query_count_in_last_period
+    query_count_in_last_period,
+    cpu_usage_in_this_period,
+    cpu_usage_in_last_period
 FROM pg_catalog.citus_stat_tenants(FALSE);
 
 ALTER VIEW citus.citus_stat_tenants SET SCHEMA pg_catalog;

src/backend/distributed/sql/udfs/citus_stat_tenants_local/latest.sql

@@ -6,6 +6,8 @@ CREATE OR REPLACE FUNCTION pg_catalog.citus_stat_tenants_local(
     OUT read_count_in_last_period INT,
     OUT query_count_in_this_period INT,
     OUT query_count_in_last_period INT,
+    OUT cpu_usage_in_this_period DOUBLE PRECISION,
+    OUT cpu_usage_in_last_period DOUBLE PRECISION,
     OUT score BIGINT)
 RETURNS SETOF RECORD
 LANGUAGE C
@@ -19,7 +21,9 @@ SELECT
     read_count_in_this_period,
     read_count_in_last_period,
     query_count_in_this_period,
-    query_count_in_last_period
+    query_count_in_last_period,
+    cpu_usage_in_this_period,
+    cpu_usage_in_last_period
 FROM pg_catalog.citus_stat_tenants_local()
 ORDER BY score DESC;
 

src/backend/distributed/utils/citus_stat_tenants.c

@@ -12,13 +12,14 @@
 #include "unistd.h"
 
 #include "distributed/citus_safe_lib.h"
+#include "distributed/colocation_utils.h"
+#include "distributed/distributed_planner.h"
+#include "distributed/jsonbutils.h"
 #include "distributed/log_utils.h"
 #include "distributed/listutils.h"
 #include "distributed/metadata_cache.h"
-#include "distributed/jsonbutils.h"
-#include "distributed/colocation_utils.h"
+#include "distributed/multi_executor.h"
 #include "distributed/tuplestore.h"
-#include "distributed/colocation_utils.h"
 #include "distributed/utils/citus_stat_tenants.h"
 #include "executor/execdesc.h"
 #include "storage/ipc.h"
@@ -38,12 +39,14 @@ ExecutorEnd_hook_type prev_ExecutorEnd = NULL;
 
 #define ATTRIBUTE_PREFIX "/*{\"tId\":"
 #define ATTRIBUTE_STRING_FORMAT "/*{\"tId\":%s,\"cId\":%d}*/"
-#define STAT_TENANTS_COLUMNS 7
+#define STAT_TENANTS_COLUMNS 9
 #define ONE_QUERY_SCORE 1000000000
 
 static char AttributeToTenant[MAX_TENANT_ATTRIBUTE_LENGTH] = "";
 static CmdType AttributeToCommandType = CMD_UNKNOWN;
 static int AttributeToColocationGroupId = INVALID_COLOCATION_ID;
+static clock_t QueryStartClock = { 0 };
+static clock_t QueryEndClock = { 0 };
 
 static const char *SharedMemoryNameForMultiTenantMonitor =
 	"Shared memory for multi tenant monitor";
@@ -142,7 +145,9 @@ citus_stat_tenants_local(PG_FUNCTION_ARGS)
 								  tenantStats->writesInThisPeriod);
 		values[5] = Int32GetDatum(tenantStats->readsInLastPeriod +
 								  tenantStats->writesInLastPeriod);
-		values[6] = Int64GetDatum(tenantStats->score);
+		values[6] = Float8GetDatum(tenantStats->cpuUsageInThisPeriod);
+		values[7] = Float8GetDatum(tenantStats->cpuUsageInLastPeriod);
+		values[8] = Int64GetDatum(tenantStats->score);
 
 		tuplestore_putvalues(tupleStore, tupleDescriptor, values, isNulls);
 	}
@@ -225,6 +230,7 @@ AttributeTask(char *tenantId, int colocationId, CmdType commandType)
 	strncpy_s(AttributeToTenant, MAX_TENANT_ATTRIBUTE_LENGTH, tenantId,
 			  MAX_TENANT_ATTRIBUTE_LENGTH - 1);
 	AttributeToCommandType = commandType;
+	QueryStartClock = clock();
 }
 
 
@@ -316,6 +322,17 @@ AttributeMetricsIfApplicable()
 		return;
 	}
 
+	/*
+	 * return if we are not in the top level to make sure we are not
+	 * stopping counting time for a sub-level execution
+	 */
+	if (ExecutorLevel != 0 || PlannerLevel != 0)
+	{
+		return;
+	}
+
+	QueryEndClock = clock();
+
 	TimestampTz queryTime = GetCurrentTimestamp();
 
 	MultiTenantMonitor *monitor = GetMultiTenantMonitor();
@@ -411,6 +428,9 @@ UpdatePeriodsIfNecessary(TenantStats *tenantStats, TimestampTz queryTime)
 
 		tenantStats->readsInLastPeriod = tenantStats->readsInThisPeriod;
 		tenantStats->readsInThisPeriod = 0;
+
+		tenantStats->cpuUsageInLastPeriod = tenantStats->cpuUsageInThisPeriod;
+		tenantStats->cpuUsageInThisPeriod = 0;
 	}
 
 	/*
@@ -422,6 +442,8 @@ UpdatePeriodsIfNecessary(TenantStats *tenantStats, TimestampTz queryTime)
 		tenantStats->writesInLastPeriod = 0;
 
 		tenantStats->readsInLastPeriod = 0;
+
+		tenantStats->cpuUsageInLastPeriod = 0;
 	}
 }
 
@@ -524,6 +546,9 @@ RecordTenantStats(TenantStats *tenantStats, TimestampTz queryTime)
 		tenantStats->writesInThisPeriod++;
 	}
 
+	double queryCpuTime = ((double) (QueryEndClock - QueryStartClock)) / CLOCKS_PER_SEC;
+	tenantStats->cpuUsageInThisPeriod += queryCpuTime;
+
 	tenantStats->lastQueryTime = queryTime;
 }
 

src/include/distributed/utils/citus_stat_tenants.h

@@ -42,6 +42,13 @@ typedef struct TenantStats
 	int writesInLastPeriod;
 	int writesInThisPeriod;
 
+
+	/*
+	 * CPU time usage of this tenant in this and last periods.
+	 */
+	double cpuUsageInLastPeriod;
+	double cpuUsageInThisPeriod;
+
 	/*
 	 * The latest time this tenant ran a query. This value is used to update the score later.
 	 */

src/test/regress/expected/citus_stat_tenants.out

@@ -71,14 +71,17 @@ INSERT INTO dist_tbl VALUES (2, 'abcd');
 UPDATE dist_tbl SET b = a + 1 WHERE a = 3;
 UPDATE dist_tbl SET b = a + 1 WHERE a = 4;
 DELETE FROM dist_tbl WHERE a = 5;
-SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period FROM citus_stat_tenants(true) ORDER BY tenant_attribute;
- tenant_attribute | read_count_in_this_period | read_count_in_last_period | query_count_in_this_period | query_count_in_last_period
----------------------------------------------------------------------
- 1                |                         0 |                         0 |                          1 |                          0
- 2                |                         0 |                         0 |                          1 |                          0
- 3                |                         0 |                         0 |                          1 |                          0
- 4                |                         0 |                         0 |                          1 |                          0
- 5                |                         0 |                         0 |                          1 |                          0
+SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period,
+    (cpu_usage_in_this_period>0) AS cpu_is_used_in_this_period, (cpu_usage_in_last_period>0) AS cpu_is_used_in_last_period
+FROM citus_stat_tenants(true)
+ORDER BY tenant_attribute;
+ tenant_attribute | read_count_in_this_period | read_count_in_last_period | query_count_in_this_period | query_count_in_last_period | cpu_is_used_in_this_period | cpu_is_used_in_last_period
+---------------------------------------------------------------------
+ 1                |                         0 |                         0 |                          1 |                          0 | t                          | f
+ 2                |                         0 |                         0 |                          1 |                          0 | t                          | f
+ 3                |                         0 |                         0 |                          1 |                          0 | t                          | f
+ 4                |                         0 |                         0 |                          1 |                          0 | t                          | f
+ 5                |                         0 |                         0 |                          1 |                          0 | t                          | f
 (5 rows)
 
 SELECT citus_stat_tenants_reset();
@@ -241,11 +244,14 @@ SELECT count(*)>=0 FROM dist_tbl WHERE a = 1;
 
 INSERT INTO dist_tbl VALUES (5, 'abcd');
 \c - - - :worker_1_port
-SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period FROM citus_stat_tenants_local ORDER BY tenant_attribute;
- tenant_attribute | read_count_in_this_period | read_count_in_last_period | query_count_in_this_period | query_count_in_last_period
----------------------------------------------------------------------
- 1                |                         1 |                         0 |                          1 |                          0
- 5                |                         0 |                         0 |                          1 |                          0
+SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period,
+    (cpu_usage_in_this_period>0) AS cpu_is_used_in_this_period, (cpu_usage_in_last_period>0) AS cpu_is_used_in_last_period
+FROM citus_stat_tenants_local
+ORDER BY tenant_attribute;
+ tenant_attribute | read_count_in_this_period | read_count_in_last_period | query_count_in_this_period | query_count_in_last_period | cpu_is_used_in_this_period | cpu_is_used_in_last_period
+---------------------------------------------------------------------
+ 1                |                         1 |                         0 |                          1 |                          0 | t                          | f
+ 5                |                         0 |                         0 |                          1 |                          0 | t                          | f
 (2 rows)
 
 -- simulate passing the period
@@ -256,11 +262,14 @@ SELECT sleep_until_next_period();
 
 (1 row)
 
-SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period FROM citus_stat_tenants_local ORDER BY tenant_attribute;
- tenant_attribute | read_count_in_this_period | read_count_in_last_period | query_count_in_this_period | query_count_in_last_period
+SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period,
+    (cpu_usage_in_this_period>0) AS cpu_is_used_in_this_period, (cpu_usage_in_last_period>0) AS cpu_is_used_in_last_period
+FROM citus_stat_tenants_local
+ORDER BY tenant_attribute;
+ tenant_attribute | read_count_in_this_period | read_count_in_last_period | query_count_in_this_period | query_count_in_last_period | cpu_is_used_in_this_period | cpu_is_used_in_last_period
 ---------------------------------------------------------------------
- 1                |                         0 |                         1 |                          0 |                          1
- 5                |                         0 |                         0 |                          0 |                          1
+ 1                |                         0 |                         1 |                          0 |                          1 | f                          | t
+ 5                |                         0 |                         0 |                          0 |                          1 | f                          | t
 (2 rows)
 
 SELECT sleep_until_next_period();
@@ -269,11 +278,14 @@ SELECT sleep_until_next_period();
 
 (1 row)
 
-SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period FROM citus_stat_tenants_local ORDER BY tenant_attribute;
- tenant_attribute | read_count_in_this_period | read_count_in_last_period | query_count_in_this_period | query_count_in_last_period
+SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period,
+    (cpu_usage_in_this_period>0) AS cpu_is_used_in_this_period, (cpu_usage_in_last_period>0) AS cpu_is_used_in_last_period
+FROM citus_stat_tenants_local
+ORDER BY tenant_attribute;
+ tenant_attribute | read_count_in_this_period | read_count_in_last_period | query_count_in_this_period | query_count_in_last_period | cpu_is_used_in_this_period | cpu_is_used_in_last_period
 ---------------------------------------------------------------------
- 1                |                         0 |                         0 |                          0 |                          0
- 5                |                         0 |                         0 |                          0 |                          0
+ 1                |                         0 |                         0 |                          0 |                          0 | f                          | f
+ 5                |                         0 |                         0 |                          0 |                          0 | f                          | f
 (2 rows)
 
 \c - - - :master_port

src/test/regress/sql/citus_stat_tenants.sql

@@ -35,7 +35,10 @@ UPDATE dist_tbl SET b = a + 1 WHERE a = 3;
 UPDATE dist_tbl SET b = a + 1 WHERE a = 4;
 DELETE FROM dist_tbl WHERE a = 5;
 
-SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period FROM citus_stat_tenants(true) ORDER BY tenant_attribute;
+SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period,
+    (cpu_usage_in_this_period>0) AS cpu_is_used_in_this_period, (cpu_usage_in_last_period>0) AS cpu_is_used_in_last_period
+FROM citus_stat_tenants(true)
+ORDER BY tenant_attribute;
 
 SELECT citus_stat_tenants_reset();
 
@@ -84,17 +87,26 @@ SELECT count(*)>=0 FROM dist_tbl WHERE a = 1;
 INSERT INTO dist_tbl VALUES (5, 'abcd');
 
 \c - - - :worker_1_port
-SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period FROM citus_stat_tenants_local ORDER BY tenant_attribute;
+SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period,
+    (cpu_usage_in_this_period>0) AS cpu_is_used_in_this_period, (cpu_usage_in_last_period>0) AS cpu_is_used_in_last_period
+FROM citus_stat_tenants_local
+ORDER BY tenant_attribute;
 
 -- simulate passing the period
 SET citus.stat_tenants_period TO 2;
 SELECT sleep_until_next_period();
 
-SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period FROM citus_stat_tenants_local ORDER BY tenant_attribute;
+SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period,
+    (cpu_usage_in_this_period>0) AS cpu_is_used_in_this_period, (cpu_usage_in_last_period>0) AS cpu_is_used_in_last_period
+FROM citus_stat_tenants_local
+ORDER BY tenant_attribute;
 
 SELECT sleep_until_next_period();
 
-SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period FROM citus_stat_tenants_local ORDER BY tenant_attribute;
+SELECT tenant_attribute, read_count_in_this_period, read_count_in_last_period, query_count_in_this_period, query_count_in_last_period,
+    (cpu_usage_in_this_period>0) AS cpu_is_used_in_this_period, (cpu_usage_in_last_period>0) AS cpu_is_used_in_last_period
+FROM citus_stat_tenants_local
+ORDER BY tenant_attribute;
 
 \c - - - :master_port
 SET search_path TO citus_stat_tenants;