Citus stats tenants collector view (#6761)

Add a view that collects statistics from all nodes

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

@@ -10,4 +10,5 @@ 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_stats_tenants_local/11.3-1.sql"
 #include "udfs/citus_stats_tenants/11.3-1.sql"

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

@@ -21,5 +21,8 @@ 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_stats_tenants_local;
+DROP FUNCTION pg_catalog.citus_stats_tenants_local(boolean);
+
 DROP VIEW pg_catalog.citus_stats_tenants;
 DROP FUNCTION pg_catalog.citus_stats_tenants(boolean);

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

@@ -1,27 +1,66 @@
-CREATE OR REPLACE FUNCTION pg_catalog.citus_stats_tenants(
+-- cts in the query is an abbreviation for citus_stats_tenants
+CREATE OR REPLACE FUNCTION pg_catalog.citus_stats_tenants (
     return_all_tenants BOOLEAN DEFAULT FALSE,
+    OUT nodeid INT,
     OUT colocation_id INT,
     OUT tenant_attribute TEXT,
     OUT read_count_in_this_period INT,
     OUT read_count_in_last_period INT,
     OUT query_count_in_this_period INT,
     OUT query_count_in_last_period INT,
-    OUT score BIGINT)
-RETURNS SETOF RECORD
-LANGUAGE C
-AS 'citus', $$citus_stats_tenants$$;
-
+    OUT score BIGINT
+)
+    RETURNS SETOF record
+    LANGUAGE plpgsql
+    AS $function$
+BEGIN
+    RETURN QUERY
+    SELECT *
+    FROM jsonb_to_recordset((
+        SELECT
+            jsonb_agg(all_cst_rows_as_jsonb.cst_row_as_jsonb)::jsonb
+        FROM (
+            SELECT
+                jsonb_array_elements(run_command_on_all_nodes.result::jsonb)::jsonb ||
+                    ('{"nodeid":' || run_command_on_all_nodes.nodeid || '}')::jsonb AS cst_row_as_jsonb
+            FROM
+                run_command_on_all_nodes (
+                    $$
+                        SELECT
+                            coalesce(to_jsonb (array_agg(cstl.*)), '[]'::jsonb)
+                        FROM citus_stats_tenants_local($$||return_all_tenants||$$) cstl;
+                    $$,
+                    parallel:= TRUE,
+                    give_warning_for_connection_errors:= TRUE)
+            WHERE
+                success = 't')
+        AS all_cst_rows_as_jsonb))
+AS (
+    nodeid INT,
+    colocation_id INT,
+    tenant_attribute TEXT,
+    read_count_in_this_period INT,
+    read_count_in_last_period INT,
+    query_count_in_this_period INT,
+    query_count_in_last_period INT,
+    score BIGINT
+)
+    ORDER BY score DESC
+    LIMIT CASE WHEN NOT return_all_tenants THEN current_setting('citus.stats_tenants_limit')::BIGINT END;
+END;
+$function$;
 
 CREATE OR REPLACE VIEW citus.citus_stats_tenants AS
 SELECT
+    nodeid,
     colocation_id,
     tenant_attribute,
     read_count_in_this_period,
     read_count_in_last_period,
     query_count_in_this_period,
     query_count_in_last_period
-FROM pg_catalog.citus_stats_tenants()
-ORDER BY score DESC;
+FROM pg_catalog.citus_stats_tenants(FALSE);
 
 ALTER VIEW citus.citus_stats_tenants SET SCHEMA pg_catalog;
+
 GRANT SELECT ON pg_catalog.citus_stats_tenants TO PUBLIC;

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

@@ -0,0 +1,27 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_stats_tenants_local(
+    return_all_tenants BOOLEAN DEFAULT FALSE,
+    OUT colocation_id INT,
+    OUT tenant_attribute TEXT,
+    OUT read_count_in_this_period INT,
+    OUT read_count_in_last_period INT,
+    OUT query_count_in_this_period INT,
+    OUT query_count_in_last_period INT,
+    OUT score BIGINT)
+RETURNS SETOF RECORD
+LANGUAGE C
+AS 'citus', $$citus_stats_tenants_local$$;
+
+
+CREATE OR REPLACE VIEW citus.citus_stats_tenants_local AS
+SELECT
+    colocation_id,
+    tenant_attribute,
+    read_count_in_this_period,
+    read_count_in_last_period,
+    query_count_in_this_period,
+    query_count_in_last_period
+FROM pg_catalog.citus_stats_tenants_local()
+ORDER BY score DESC;
+
+ALTER VIEW citus.citus_stats_tenants_local SET SCHEMA pg_catalog;
+GRANT SELECT ON pg_catalog.citus_stats_tenants_local TO PUBLIC;

src/backend/distributed/utils/attribute.c

@@ -68,16 +68,16 @@ int CitusStatsTenantsPeriod = (time_t) 60;
 int CitusStatsTenantsLimit = 10;
 
 
-PG_FUNCTION_INFO_V1(citus_stats_tenants);
+PG_FUNCTION_INFO_V1(citus_stats_tenants_local);
 PG_FUNCTION_INFO_V1(clean_citus_stats_tenants);
 PG_FUNCTION_INFO_V1(sleep_until_next_period);
 
 
 /*
- * citus_stats_tenants finds, updates and returns the statistics for tenants.
+ * citus_stats_tenants_local finds, updates and returns the statistics for tenants.
  */
 Datum
-citus_stats_tenants(PG_FUNCTION_ARGS)
+citus_stats_tenants_local(PG_FUNCTION_ARGS)
 {
 	CheckCitusVersion(ERROR);
 

src/test/regress/expected/citus_stats_tenants.out

@@ -68,24 +68,16 @@ 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;
-\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_stats_tenants 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 FROM citus_stats_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
- 5                |                         0 |                         0 |                          1 |                          0
-(2 rows)
-
-\c - - - :worker_2_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_stats_tenants 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
----------------------------------------------------------------------
  2                |                         0 |                         0 |                          1 |                          0
  3                |                         0 |                         0 |                          1 |                          0
-(2 rows)
+ 4                |                         0 |                         0 |                          1 |                          0
+ 5                |                         0 |                         0 |                          1 |                          0
+(5 rows)
 
-\c - - - :master_port
-SET search_path TO citus_stats_tenants;
 SELECT result FROM run_command_on_all_nodes('SELECT clean_citus_stats_tenants()');
  result
 ---------------------------------------------------------------------
@@ -108,14 +100,11 @@ SELECT count(*)>=0 FROM ref_tbl WHERE a = 1;
  t
 (1 row)
 
-\c - - - :worker_1_port
-SELECT tenant_attribute, query_count_in_this_period FROM citus_stats_tenants ORDER BY tenant_attribute;
+SELECT tenant_attribute, query_count_in_this_period FROM citus_stats_tenants(true) ORDER BY tenant_attribute;
  tenant_attribute | query_count_in_this_period
 ---------------------------------------------------------------------
 (0 rows)
 
-\c - - - :master_port
-SET search_path TO citus_stats_tenants;
 -- queries with multiple tables but one tenant should be counted
 SELECT count(*)>=0 FROM dist_tbl, dist_tbl_2 WHERE dist_tbl.a = 1 AND dist_tbl_2.a = 1;
  ?column?
@@ -129,17 +118,15 @@ SELECT count(*)>=0 FROM dist_tbl JOIN dist_tbl_2 ON dist_tbl.a = dist_tbl_2.a WH
  t
 (1 row)
 
-\c - - - :worker_1_port
-SELECT tenant_attribute, query_count_in_this_period FROM citus_stats_tenants WHERE tenant_attribute = '1';
+SELECT tenant_attribute, query_count_in_this_period FROM citus_stats_tenants(true) WHERE tenant_attribute = '1';
  tenant_attribute | query_count_in_this_period
 ---------------------------------------------------------------------
  1                |                          2
 (1 row)
 
-\c - - - :master_port
-SET search_path TO citus_stats_tenants;
 -- test scoring
 -- all of these distribution column values are from second worker
+SELECT nodeid AS worker_2_nodeid FROM pg_dist_node WHERE nodeport = :worker_2_port \gset
 SELECT count(*)>=0 FROM dist_tbl WHERE a = 2;
  ?column?
 ---------------------------------------------------------------------
@@ -164,8 +151,7 @@ SELECT count(*)>=0 FROM dist_tbl_text WHERE a = 'abcd';
  t
 (1 row)
 
-\c - - - :worker_2_port
-SELECT tenant_attribute, query_count_in_this_period, score FROM citus_stats_tenants(true) ORDER BY score DESC;
+SELECT tenant_attribute, query_count_in_this_period, score FROM citus_stats_tenants(true) WHERE nodeid = :worker_2_nodeid ORDER BY score DESC, tenant_attribute;
  tenant_attribute | query_count_in_this_period |   score
 ---------------------------------------------------------------------
  2                |                          1 | 1000000000
@@ -174,8 +160,6 @@ SELECT tenant_attribute, query_count_in_this_period, score FROM citus_stats_tena
  abcd             |                          1 | 1000000000
 (4 rows)
 
-\c - - - :master_port
-SET search_path TO citus_stats_tenants;
 SELECT count(*)>=0 FROM dist_tbl_text WHERE a = 'abcd';
  ?column?
 ---------------------------------------------------------------------
@@ -200,8 +184,7 @@ SELECT count(*)>=0 FROM dist_tbl_text WHERE a = 'cdef';
  t
 (1 row)
 
-\c - - - :worker_2_port
-SELECT tenant_attribute, query_count_in_this_period, score FROM citus_stats_tenants(true) ORDER BY score DESC;
+SELECT tenant_attribute, query_count_in_this_period, score FROM citus_stats_tenants(true) WHERE nodeid = :worker_2_nodeid ORDER BY score DESC, tenant_attribute;
  tenant_attribute | query_count_in_this_period |   score
 ---------------------------------------------------------------------
  abcd             |                          3 | 3000000000
@@ -212,8 +195,6 @@ SELECT tenant_attribute, query_count_in_this_period, score FROM citus_stats_tena
  cdef             |                          1 | 1000000000
 (6 rows)
 
-\c - - - :master_port
-SET search_path TO citus_stats_tenants;
 SELECT count(*)>=0 FROM dist_tbl_text WHERE a = 'bcde';
  ?column?
 ---------------------------------------------------------------------
@@ -232,8 +213,7 @@ SELECT count(*)>=0 FROM dist_tbl_text WHERE a = 'defg';
  t
 (1 row)
 
-\c - - - :worker_2_port
-SELECT tenant_attribute, query_count_in_this_period, score FROM citus_stats_tenants(true) ORDER BY score DESC;
+SELECT tenant_attribute, query_count_in_this_period, score FROM citus_stats_tenants(true) WHERE nodeid = :worker_2_nodeid ORDER BY score DESC, tenant_attribute;
  tenant_attribute | query_count_in_this_period |   score
 ---------------------------------------------------------------------
  abcd             |                          3 | 3000000000
@@ -243,8 +223,6 @@ SELECT tenant_attribute, query_count_in_this_period, score FROM citus_stats_tena
  defg             |                          1 | 1000000000
 (5 rows)
 
-\c - - - :master_port
-SET search_path TO citus_stats_tenants;
 -- test period passing
 SELECT result FROM run_command_on_all_nodes('SELECT clean_citus_stats_tenants()');
  result
@@ -262,7 +240,7 @@ 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_stats_tenants 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 FROM citus_stats_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
@@ -277,7 +255,7 @@ 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_stats_tenants 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 FROM citus_stats_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                |                         0 |                         1 |                          0 |                          1

src/test/regress/expected/multi_extension.out

@@ -1367,10 +1367,12 @@ SELECT * FROM multi_extension.print_extension_changes();
                  | function citus_internal_start_replication_origin_tracking() void
                  | function citus_internal_stop_replication_origin_tracking() void
                  | function citus_stats_tenants(boolean) SETOF record
+                 | function citus_stats_tenants_local(boolean) SETOF record
                  | function worker_adjust_identity_column_seq_ranges(regclass) void
                  | function worker_drop_all_shell_tables(boolean)
                  | view citus_stats_tenants
-(8 rows)
+                 | view citus_stats_tenants_local
+(10 rows)
 
 DROP TABLE multi_extension.prev_objects, multi_extension.extension_diff;
 -- show running version