-- -- BACKGROUND_REBALANCE_PARALLEL -- -- Test to check if the background tasks scheduled by the background rebalancer -- have the correct dependencies -- -- Test to verify that we do not allow parallel rebalancer moves involving a -- particular node (either as source or target) more than -- citus.max_background_task_executors_per_node, and that we can change the GUC on -- the fly, and that will affect the ongoing balance as it should -- -- Test to verify that there's a hard dependency when a specific node is first being -- used as a source for a move, and then later as a target. -- CREATE SCHEMA background_rebalance_parallel; SET search_path TO background_rebalance_parallel; SET citus.next_shard_id TO 85674000; SET citus.shard_replication_factor TO 1; SET client_min_messages TO ERROR; ALTER SEQUENCE pg_dist_background_job_job_id_seq RESTART 17777; ALTER SEQUENCE pg_dist_background_task_task_id_seq RESTART 1000; ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 50050; SELECT nextval('pg_catalog.pg_dist_groupid_seq') AS last_group_id_cls \gset SELECT nextval('pg_catalog.pg_dist_node_nodeid_seq') AS last_node_id_cls \gset ALTER SEQUENCE pg_catalog.pg_dist_groupid_seq RESTART 50; ALTER SEQUENCE pg_catalog.pg_dist_node_nodeid_seq RESTART 50; SELECT 1 FROM master_remove_node('localhost', :worker_1_port); ?column? --------------------------------------------------------------------- 1 (1 row) SELECT 1 FROM master_remove_node('localhost', :worker_2_port); ?column? --------------------------------------------------------------------- 1 (1 row) SELECT 1 FROM master_add_node('localhost', :worker_1_port); ?column? --------------------------------------------------------------------- 1 (1 row) SELECT 1 FROM master_add_node('localhost', :worker_2_port); ?column? --------------------------------------------------------------------- 1 (1 row) ALTER SYSTEM SET citus.background_task_queue_interval TO '1s'; SELECT pg_reload_conf(); pg_reload_conf --------------------------------------------------------------------- t (1 row) -- Colocation group 1: create two tables table1_colg1, table2_colg1 and in a colocation group CREATE TABLE table1_colg1 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg1', 'a', shard_count => 4, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) CREATE TABLE table2_colg1 (b int PRIMARY KEY); SELECT create_distributed_table('table2_colg1', 'b', colocate_with => 'table1_colg1'); create_distributed_table --------------------------------------------------------------------- (1 row) -- Colocation group 2: create two tables table1_colg2, table2_colg2 and in a colocation group CREATE TABLE table1_colg2 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg2', 'a', shard_count => 4, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) CREATE TABLE table2_colg2 (b int primary key); SELECT create_distributed_table('table2_colg2', 'b', colocate_with => 'table1_colg2'); create_distributed_table --------------------------------------------------------------------- (1 row) -- Colocation group 3: create two tables table1_colg3, table2_colg3 and in a colocation group CREATE TABLE table1_colg3 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg3', 'a', shard_count => 4, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) CREATE TABLE table2_colg3 (b int primary key); SELECT create_distributed_table('table2_colg3', 'b', colocate_with => 'table1_colg3'); create_distributed_table --------------------------------------------------------------------- (1 row) -- Add two new nodes so that we can rebalance SELECT 1 FROM citus_add_node('localhost', :worker_3_port); ?column? --------------------------------------------------------------------- 1 (1 row) SELECT 1 FROM citus_add_node('localhost', :worker_4_port); ?column? --------------------------------------------------------------------- 1 (1 row) SELECT * FROM get_rebalance_table_shards_plan() ORDER BY shardid; table_name | shardid | shard_size | sourcename | sourceport | targetname | targetport --------------------------------------------------------------------- table1_colg1 | 85674000 | 0 | localhost | 57637 | localhost | 57640 table1_colg1 | 85674001 | 0 | localhost | 57638 | localhost | 57639 table2_colg1 | 85674004 | 0 | localhost | 57637 | localhost | 57640 table2_colg1 | 85674005 | 0 | localhost | 57638 | localhost | 57639 table1_colg2 | 85674008 | 0 | localhost | 57637 | localhost | 57640 table1_colg2 | 85674009 | 0 | localhost | 57638 | localhost | 57639 table2_colg2 | 85674012 | 0 | localhost | 57637 | localhost | 57640 table2_colg2 | 85674013 | 0 | localhost | 57638 | localhost | 57639 table1_colg3 | 85674016 | 0 | localhost | 57637 | localhost | 57640 table1_colg3 | 85674017 | 0 | localhost | 57638 | localhost | 57639 table2_colg3 | 85674020 | 0 | localhost | 57637 | localhost | 57640 table2_colg3 | 85674021 | 0 | localhost | 57638 | localhost | 57639 (12 rows) SELECT * FROM citus_rebalance_start(); citus_rebalance_start --------------------------------------------------------------------- 17777 (1 row) SELECT citus_rebalance_wait(); citus_rebalance_wait --------------------------------------------------------------------- (1 row) -- PART 1 -- Test to check if the background tasks scheduled by the background rebalancer -- have the correct dependencies -- Check that a move is dependent on -- any other move scheduled earlier in its colocation group. SELECT S.shardid, P.colocationid FROM pg_dist_shard S, pg_dist_partition P WHERE S.logicalrelid = P.logicalrelid ORDER BY S.shardid ASC; shardid | colocationid --------------------------------------------------------------------- 85674000 | 50050 85674001 | 50050 85674002 | 50050 85674003 | 50050 85674004 | 50050 85674005 | 50050 85674006 | 50050 85674007 | 50050 85674008 | 50051 85674009 | 50051 85674010 | 50051 85674011 | 50051 85674012 | 50051 85674013 | 50051 85674014 | 50051 85674015 | 50051 85674016 | 50052 85674017 | 50052 85674018 | 50052 85674019 | 50052 85674020 | 50052 85674021 | 50052 85674022 | 50052 85674023 | 50052 (24 rows) SELECT D.task_id, (SELECT T.command FROM pg_dist_background_task T WHERE T.task_id = D.task_id), D.depends_on, (SELECT T.command FROM pg_dist_background_task T WHERE T.task_id = D.depends_on) FROM pg_dist_background_task_depend D WHERE job_id = 17777 ORDER BY D.task_id, D.depends_on ASC; task_id | command | depends_on | command --------------------------------------------------------------------- 1001 | SELECT pg_catalog.citus_move_shard_placement(85674000,50,53,'auto') | 1000 | SELECT pg_catalog.citus_move_shard_placement(85674001,51,52,'auto') 1003 | SELECT pg_catalog.citus_move_shard_placement(85674008,50,53,'auto') | 1002 | SELECT pg_catalog.citus_move_shard_placement(85674009,51,52,'auto') 1005 | SELECT pg_catalog.citus_move_shard_placement(85674016,50,53,'auto') | 1004 | SELECT pg_catalog.citus_move_shard_placement(85674017,51,52,'auto') (3 rows) -- Check that if there is a reference table that needs to be synched to a node, -- any move without a dependency must depend on the move task for reference table. SELECT 1 FROM citus_drain_node('localhost',:worker_4_port); ?column? --------------------------------------------------------------------- 1 (1 row) SELECT public.wait_for_resource_cleanup(); wait_for_resource_cleanup --------------------------------------------------------------------- (1 row) SELECT 1 FROM citus_disable_node('localhost', :worker_4_port, synchronous:=true); ?column? --------------------------------------------------------------------- 1 (1 row) -- Drain worker_3 so that we can move only one colocation group to worker_3 -- to create an unbalance that would cause parallel rebalancing. SELECT 1 FROM citus_drain_node('localhost',:worker_3_port); ?column? --------------------------------------------------------------------- 1 (1 row) SELECT citus_set_node_property('localhost', :worker_3_port, 'shouldhaveshards', true); citus_set_node_property --------------------------------------------------------------------- (1 row) CALL citus_cleanup_orphaned_resources(); CREATE TABLE ref_table(a int PRIMARY KEY); SELECT create_reference_table('ref_table'); create_reference_table --------------------------------------------------------------------- (1 row) -- Move all the shards of Colocation group 3 to worker_3. SELECT master_move_shard_placement(shardid, 'localhost', nodeport, 'localhost', :worker_3_port, 'block_writes') FROM pg_dist_shard NATURAL JOIN pg_dist_shard_placement WHERE logicalrelid = 'table1_colg3'::regclass AND nodeport <> :worker_3_port ORDER BY shardid; master_move_shard_placement --------------------------------------------------------------------- (4 rows) CALL citus_cleanup_orphaned_resources(); -- Activate and new nodes so that we can rebalance. SELECT 1 FROM citus_activate_node('localhost', :worker_4_port); ?column? --------------------------------------------------------------------- 1 (1 row) SELECT citus_set_node_property('localhost', :worker_4_port, 'shouldhaveshards', true); citus_set_node_property --------------------------------------------------------------------- (1 row) SELECT 1 FROM citus_add_node('localhost', :worker_5_port); ?column? --------------------------------------------------------------------- 1 (1 row) SELECT 1 FROM citus_add_node('localhost', :worker_6_port); ?column? --------------------------------------------------------------------- 1 (1 row) SELECT * FROM citus_rebalance_start(); citus_rebalance_start --------------------------------------------------------------------- 17778 (1 row) SELECT citus_rebalance_wait(); citus_rebalance_wait --------------------------------------------------------------------- (1 row) SELECT S.shardid, P.colocationid FROM pg_dist_shard S, pg_dist_partition P WHERE S.logicalrelid = P.logicalrelid ORDER BY S.shardid ASC; shardid | colocationid --------------------------------------------------------------------- 85674000 | 50050 85674001 | 50050 85674002 | 50050 85674003 | 50050 85674004 | 50050 85674005 | 50050 85674006 | 50050 85674007 | 50050 85674008 | 50051 85674009 | 50051 85674010 | 50051 85674011 | 50051 85674012 | 50051 85674013 | 50051 85674014 | 50051 85674015 | 50051 85674016 | 50052 85674017 | 50052 85674018 | 50052 85674019 | 50052 85674020 | 50052 85674021 | 50052 85674022 | 50052 85674023 | 50052 85674024 | 50053 (25 rows) SELECT D.task_id, (SELECT T.command FROM pg_dist_background_task T WHERE T.task_id = D.task_id), D.depends_on, (SELECT T.command FROM pg_dist_background_task T WHERE T.task_id = D.depends_on) FROM pg_dist_background_task_depend D WHERE job_id = 17778 ORDER BY D.task_id, D.depends_on ASC; task_id | command | depends_on | command --------------------------------------------------------------------- 1007 | SELECT pg_catalog.citus_move_shard_placement(85674016,52,53,'auto') | 1006 | SELECT pg_catalog.replicate_reference_tables('auto') 1008 | SELECT pg_catalog.citus_move_shard_placement(85674003,51,54,'auto') | 1006 | SELECT pg_catalog.replicate_reference_tables('auto') 1009 | SELECT pg_catalog.citus_move_shard_placement(85674000,50,55,'auto') | 1008 | SELECT pg_catalog.citus_move_shard_placement(85674003,51,54,'auto') 1010 | SELECT pg_catalog.citus_move_shard_placement(85674017,52,53,'auto') | 1007 | SELECT pg_catalog.citus_move_shard_placement(85674016,52,53,'auto') 1011 | SELECT pg_catalog.citus_move_shard_placement(85674008,51,54,'auto') | 1006 | SELECT pg_catalog.replicate_reference_tables('auto') 1012 | SELECT pg_catalog.citus_move_shard_placement(85674001,50,55,'auto') | 1009 | SELECT pg_catalog.citus_move_shard_placement(85674000,50,55,'auto') (6 rows) -- PART 2 -- Test to verify that we do not allow parallel rebalancer moves involving a -- particular node (either as source or target) -- more than citus.max_background_task_executors_per_node -- and that we can change the GUC on the fly -- citus_task_wait calls are used to ensure consistent pg_dist_background_task query -- output i.e. to avoid flakiness -- First let's restart the scenario DROP SCHEMA background_rebalance_parallel CASCADE; TRUNCATE pg_dist_background_job CASCADE; TRUNCATE pg_dist_background_task CASCADE; TRUNCATE pg_dist_background_task_depend; SELECT public.wait_for_resource_cleanup(); wait_for_resource_cleanup --------------------------------------------------------------------- (1 row) select citus_remove_node('localhost', :worker_2_port); citus_remove_node --------------------------------------------------------------------- (1 row) select citus_remove_node('localhost', :worker_3_port); citus_remove_node --------------------------------------------------------------------- (1 row) select citus_remove_node('localhost', :worker_4_port); citus_remove_node --------------------------------------------------------------------- (1 row) select citus_remove_node('localhost', :worker_5_port); citus_remove_node --------------------------------------------------------------------- (1 row) select citus_remove_node('localhost', :worker_6_port); citus_remove_node --------------------------------------------------------------------- (1 row) CREATE SCHEMA background_rebalance_parallel; SET search_path TO background_rebalance_parallel; -- Create 8 tables in 4 colocation groups, and populate them CREATE TABLE table1_colg1 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg1', 'a', shard_count => 3, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table1_colg1 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table2_colg1 (b int PRIMARY KEY); SELECT create_distributed_table('table2_colg1', 'b', colocate_with => 'table1_colg1'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table2_colg1 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table1_colg2 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg2', 'a', shard_count => 3, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table1_colg2 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table2_colg2 (b int PRIMARY KEY); SELECT create_distributed_table('table2_colg2', 'b', colocate_with => 'table1_colg2'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table2_colg2 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table1_colg3 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg3', 'a', shard_count => 3, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table1_colg3 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table2_colg3 (b int primary key); SELECT create_distributed_table('table2_colg3', 'b', colocate_with => 'table1_colg3'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table2_colg3 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table1_colg4 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg4', 'a', shard_count => 3, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table1_colg4 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table2_colg4 (b int PRIMARY KEY); SELECT create_distributed_table('table2_colg4', 'b', colocate_with => 'table1_colg4'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table2_colg4 SELECT i FROM generate_series(0, 100)i; -- Add nodes so that we can rebalance SELECT citus_add_node('localhost', :worker_2_port); citus_add_node --------------------------------------------------------------------- 56 (1 row) SELECT citus_add_node('localhost', :worker_3_port); citus_add_node --------------------------------------------------------------------- 57 (1 row) SELECT citus_rebalance_start AS job_id from citus_rebalance_start() \gset -- see dependent tasks to understand which tasks remain runnable because of -- citus.max_background_task_executors_per_node -- and which tasks are actually blocked from colocation group dependencies SELECT task_id, depends_on FROM pg_dist_background_task_depend WHERE job_id in (:job_id) ORDER BY 1, 2 ASC; task_id | depends_on --------------------------------------------------------------------- 1014 | 1013 1016 | 1015 1018 | 1017 1020 | 1019 (4 rows) -- default citus.max_background_task_executors_per_node is 1 -- show that first exactly one task per node is running -- among the tasks that are not blocked SELECT citus_task_wait(1013, desired_status => 'running'); citus_task_wait --------------------------------------------------------------------- (1 row) SELECT job_id, task_id, status, nodes_involved FROM pg_dist_background_task WHERE job_id in (:job_id) ORDER BY task_id; job_id | task_id | status | nodes_involved --------------------------------------------------------------------- 17779 | 1013 | running | {50,56} 17779 | 1014 | blocked | {50,57} 17779 | 1015 | runnable | {50,56} 17779 | 1016 | blocked | {50,57} 17779 | 1017 | runnable | {50,56} 17779 | 1018 | blocked | {50,57} 17779 | 1019 | runnable | {50,56} 17779 | 1020 | blocked | {50,57} (8 rows) -- increase citus.max_background_task_executors_per_node SELECT citus_task_wait(1013, desired_status => 'done'); citus_task_wait --------------------------------------------------------------------- (1 row) ALTER SYSTEM SET citus.max_background_task_executors_per_node = 2; SELECT pg_reload_conf(); pg_reload_conf --------------------------------------------------------------------- t (1 row) SELECT citus_task_wait(1014, desired_status => 'running'); citus_task_wait --------------------------------------------------------------------- (1 row) SELECT citus_task_wait(1015, desired_status => 'running'); citus_task_wait --------------------------------------------------------------------- (1 row) -- show that at most 2 tasks per node are running -- among the tasks that are not blocked SELECT job_id, task_id, status, nodes_involved FROM pg_dist_background_task WHERE job_id in (:job_id) ORDER BY task_id; job_id | task_id | status | nodes_involved --------------------------------------------------------------------- 17779 | 1013 | done | {50,56} 17779 | 1014 | running | {50,57} 17779 | 1015 | running | {50,56} 17779 | 1016 | blocked | {50,57} 17779 | 1017 | runnable | {50,56} 17779 | 1018 | blocked | {50,57} 17779 | 1019 | runnable | {50,56} 17779 | 1020 | blocked | {50,57} (8 rows) -- decrease to default (1) ALTER SYSTEM RESET citus.max_background_task_executors_per_node; SELECT pg_reload_conf(); pg_reload_conf --------------------------------------------------------------------- t (1 row) SELECT citus_task_wait(1015, desired_status => 'done'); citus_task_wait --------------------------------------------------------------------- (1 row) SELECT citus_task_wait(1014, desired_status => 'done'); citus_task_wait --------------------------------------------------------------------- (1 row) SELECT citus_task_wait(1016, desired_status => 'running'); citus_task_wait --------------------------------------------------------------------- (1 row) -- show that exactly one task per node is running -- among the tasks that are not blocked SELECT job_id, task_id, status, nodes_involved FROM pg_dist_background_task WHERE job_id in (:job_id) ORDER BY task_id; job_id | task_id | status | nodes_involved --------------------------------------------------------------------- 17779 | 1013 | done | {50,56} 17779 | 1014 | done | {50,57} 17779 | 1015 | done | {50,56} 17779 | 1016 | running | {50,57} 17779 | 1017 | runnable | {50,56} 17779 | 1018 | blocked | {50,57} 17779 | 1019 | runnable | {50,56} 17779 | 1020 | blocked | {50,57} (8 rows) SELECT citus_rebalance_stop(); citus_rebalance_stop --------------------------------------------------------------------- (1 row) -- PART 3 -- Test to verify that there's a hard dependency when A specific node is first being used as a -- source for a move, and then later as a target. -- First let's restart the scenario DROP SCHEMA background_rebalance_parallel CASCADE; TRUNCATE pg_dist_background_job CASCADE; TRUNCATE pg_dist_background_task CASCADE; TRUNCATE pg_dist_background_task_depend; SELECT public.wait_for_resource_cleanup(); wait_for_resource_cleanup --------------------------------------------------------------------- (1 row) select citus_remove_node('localhost', :worker_1_port); citus_remove_node --------------------------------------------------------------------- (1 row) select citus_remove_node('localhost', :worker_2_port); citus_remove_node --------------------------------------------------------------------- (1 row) select citus_remove_node('localhost', :worker_3_port); citus_remove_node --------------------------------------------------------------------- (1 row) CREATE SCHEMA background_rebalance_parallel; SET search_path TO background_rebalance_parallel; SET citus.next_shard_id TO 85674051; ALTER SEQUENCE pg_catalog.pg_dist_node_nodeid_seq RESTART 61; -- add the first node -- nodeid here is 61 select citus_add_node('localhost', :worker_1_port); citus_add_node --------------------------------------------------------------------- 61 (1 row) -- create, populate and distribute 6 tables, each with 1 shard, none colocated with each other CREATE TABLE table1_colg1 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg1', 'a', shard_count => 1, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table1_colg1 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table1_colg2 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg2', 'a', shard_count => 1, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table1_colg2 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table1_colg3 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg3', 'a', shard_count => 1, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table1_colg3 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table1_colg4 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg4', 'a', shard_count => 1, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table1_colg4 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table1_colg5 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg5', 'a', shard_count => 1, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table1_colg5 SELECT i FROM generate_series(0, 100)i; CREATE TABLE table1_colg6 (a int PRIMARY KEY); SELECT create_distributed_table('table1_colg6', 'a', shard_count => 1, colocate_with => 'none'); create_distributed_table --------------------------------------------------------------------- (1 row) INSERT INTO table1_colg6 SELECT i FROM generate_series(0, 100)i; -- add two other nodes -- nodeid here is 62 select citus_add_node('localhost', :worker_2_port); citus_add_node --------------------------------------------------------------------- 62 (1 row) -- nodeid here is 63 select citus_add_node('localhost', :worker_3_port); citus_add_node --------------------------------------------------------------------- 63 (1 row) CREATE OR REPLACE FUNCTION shard_placement_rebalance_array( worker_node_list json[], shard_placement_list json[], threshold float4 DEFAULT 0, max_shard_moves int DEFAULT 1000000, drain_only bool DEFAULT false, improvement_threshold float4 DEFAULT 0.5 ) RETURNS json[] AS 'citus' LANGUAGE C STRICT VOLATILE; -- we are simulating the following from shard_rebalancer_unit.sql -- the following steps are all according to this scenario -- where the third move should be dependent of the first two -- because the third move's target is the source of the first two SELECT unnest(shard_placement_rebalance_array( ARRAY['{"node_name": "hostname1", "disallowed_shards": "1,2,3,5,6"}', '{"node_name": "hostname2", "disallowed_shards": "4"}', '{"node_name": "hostname3", "disallowed_shards": "4"}' ]::json[], ARRAY['{"shardid":1, "nodename":"hostname1"}', '{"shardid":2, "nodename":"hostname1"}', '{"shardid":3, "nodename":"hostname2"}', '{"shardid":4, "nodename":"hostname2"}', '{"shardid":5, "nodename":"hostname3"}', '{"shardid":6, "nodename":"hostname3"}' ]::json[] )); unnest --------------------------------------------------------------------- {"updatetype":1,"shardid":1,"sourcename":"hostname1","sourceport":5432,"targetname":"hostname2","targetport":5432} {"updatetype":1,"shardid":2,"sourcename":"hostname1","sourceport":5432,"targetname":"hostname3","targetport":5432} {"updatetype":1,"shardid":4,"sourcename":"hostname2","sourceport":5432,"targetname":"hostname1","targetport":5432} (3 rows) -- manually balance the cluster such that we have -- a balanced cluster like above with 1,2,3,4,5,6 and hostname1/2/3 -- shardid 85674051 (1) nodeid 61 (hostname1) -- shardid 85674052 (2) nodeid 61 (hostname1) -- shardid 85674053 (3) nodeid 62 (hostname2) -- shardid 85674054 (4) nodeid 62 (hostname2) -- shardid 85674055 (5) nodeid 63 (hostname3) -- shardid 85674056 (6) nodeid 63 (hostname3) SELECT pg_catalog.citus_move_shard_placement(85674053,61,62,'auto'); citus_move_shard_placement --------------------------------------------------------------------- (1 row) SELECT pg_catalog.citus_move_shard_placement(85674054,61,62,'auto'); citus_move_shard_placement --------------------------------------------------------------------- (1 row) SELECT pg_catalog.citus_move_shard_placement(85674055,61,63,'auto'); citus_move_shard_placement --------------------------------------------------------------------- (1 row) SELECT pg_catalog.citus_move_shard_placement(85674056,61,63,'auto'); citus_move_shard_placement --------------------------------------------------------------------- (1 row) -- now create another rebalance strategy in order to simulate moves -- which use as target a node that has been previously used as source CREATE OR REPLACE FUNCTION test_shard_allowed_on_node(shardid bigint, nodeid int) RETURNS boolean AS $$ -- analogous to '{"node_name": "hostname1", "disallowed_shards": "1,2,3,5,6"}' select case when (shardid != 85674054 and nodeid = 61) then false -- analogous to '{"node_name": "hostname2", "disallowed_shards": "4"}' -- AND '{"node_name": "hostname2", "disallowed_shards": "4"}' when (shardid = 85674054 and nodeid != 61) then false else true end; $$ LANGUAGE sql; -- insert the new test rebalance strategy INSERT INTO pg_catalog.pg_dist_rebalance_strategy( name, default_strategy, shard_cost_function, node_capacity_function, shard_allowed_on_node_function, default_threshold, minimum_threshold, improvement_threshold ) VALUES ( 'test_source_then_target', false, 'citus_shard_cost_1', 'citus_node_capacity_1', 'background_rebalance_parallel.test_shard_allowed_on_node', 0, 0, 0 ); SELECT * FROM get_rebalance_table_shards_plan(rebalance_strategy := 'test_source_then_target'); table_name | shardid | shard_size | sourcename | sourceport | targetname | targetport --------------------------------------------------------------------- table1_colg1 | 85674051 | 0 | localhost | 57637 | localhost | 57638 table1_colg2 | 85674052 | 0 | localhost | 57637 | localhost | 57639 table1_colg4 | 85674054 | 0 | localhost | 57638 | localhost | 57637 (3 rows) SELECT citus_rebalance_start AS job_id from citus_rebalance_start(rebalance_strategy := 'test_source_then_target') \gset -- check that the third move is blocked and depends on the first two SELECT job_id, task_id, status, nodes_involved FROM pg_dist_background_task WHERE job_id in (:job_id) ORDER BY task_id; job_id | task_id | status | nodes_involved --------------------------------------------------------------------- 17780 | 1021 | runnable | {61,62} 17780 | 1022 | runnable | {61,63} 17780 | 1023 | blocked | {62,61} (3 rows) SELECT D.task_id, (SELECT T.command FROM pg_dist_background_task T WHERE T.task_id = D.task_id), D.depends_on, (SELECT T.command FROM pg_dist_background_task T WHERE T.task_id = D.depends_on) FROM pg_dist_background_task_depend D WHERE job_id in (:job_id) ORDER BY D.task_id, D.depends_on ASC; task_id | command | depends_on | command --------------------------------------------------------------------- 1023 | SELECT pg_catalog.citus_move_shard_placement(85674054,62,61,'auto') | 1021 | SELECT pg_catalog.citus_move_shard_placement(85674051,61,62,'auto') 1023 | SELECT pg_catalog.citus_move_shard_placement(85674054,62,61,'auto') | 1022 | SELECT pg_catalog.citus_move_shard_placement(85674052,61,63,'auto') (2 rows) SELECT citus_rebalance_stop(); citus_rebalance_stop --------------------------------------------------------------------- (1 row) DELETE FROM pg_catalog.pg_dist_rebalance_strategy WHERE name='test_source_then_target'; DROP SCHEMA background_rebalance_parallel CASCADE; TRUNCATE pg_dist_background_job CASCADE; TRUNCATE pg_dist_background_task CASCADE; TRUNCATE pg_dist_background_task_depend; SELECT public.wait_for_resource_cleanup(); wait_for_resource_cleanup --------------------------------------------------------------------- (1 row) select citus_remove_node('localhost', :worker_3_port); citus_remove_node --------------------------------------------------------------------- (1 row) -- keep the rest of the tests inact that depends node/group ids ALTER SEQUENCE pg_catalog.pg_dist_groupid_seq RESTART :last_group_id_cls; ALTER SEQUENCE pg_catalog.pg_dist_node_nodeid_seq RESTART :last_node_id_cls;