From 858d99be33e0057e04dc74f687f145dad89d6b59 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 9 Oct 2023 13:13:08 +0300 Subject: [PATCH] Take improvement_threshold into the account in citus_add_rebalance_strategy() (#7247) DESCRIPTION: Makes sure to take improvement_threshold into the account in `citus_add_rebalance_strategy()`. Fixes https://github.com/citusdata/citus/issues/7188. --- .../distributed/sql/citus--12.1-1--12.2-1.sql | 2 ++ .../sql/downgrades/citus--12.2-1--12.1-1.sql | 3 +- .../citus_add_rebalance_strategy/12.2-1.sql | 32 +++++++++++++++++++ .../citus_add_rebalance_strategy/latest.sql | 6 ++-- .../regress/expected/shard_rebalancer.out | 21 ++++++++++++ .../upgrade_rebalance_strategy_after.out | 2 +- src/test/regress/sql/shard_rebalancer.sql | 14 ++++++++ 7 files changed, 76 insertions(+), 4 deletions(-) create mode 100644 src/backend/distributed/sql/udfs/citus_add_rebalance_strategy/12.2-1.sql diff --git a/src/backend/distributed/sql/citus--12.1-1--12.2-1.sql b/src/backend/distributed/sql/citus--12.1-1--12.2-1.sql index bb9d22969..ec4cc7134 100644 --- a/src/backend/distributed/sql/citus--12.1-1--12.2-1.sql +++ b/src/backend/distributed/sql/citus--12.1-1--12.2-1.sql @@ -1,3 +1,5 @@ -- citus--12.1-1--12.2-1 -- bump version to 12.2-1 + +#include "udfs/citus_add_rebalance_strategy/12.2-1.sql" diff --git a/src/backend/distributed/sql/downgrades/citus--12.2-1--12.1-1.sql b/src/backend/distributed/sql/downgrades/citus--12.2-1--12.1-1.sql index b26fc16bc..93d121a12 100644 --- a/src/backend/distributed/sql/downgrades/citus--12.2-1--12.1-1.sql +++ b/src/backend/distributed/sql/downgrades/citus--12.2-1--12.1-1.sql @@ -1,2 +1,3 @@ -- citus--12.2-1--12.1-1 --- this is an empty downgrade path since citus--12.2-1--12.1-1.sql is empty for now + +#include "../udfs/citus_add_rebalance_strategy/10.1-1.sql" diff --git a/src/backend/distributed/sql/udfs/citus_add_rebalance_strategy/12.2-1.sql b/src/backend/distributed/sql/udfs/citus_add_rebalance_strategy/12.2-1.sql new file mode 100644 index 000000000..c4f157c2e --- /dev/null +++ b/src/backend/distributed/sql/udfs/citus_add_rebalance_strategy/12.2-1.sql @@ -0,0 +1,32 @@ +DROP FUNCTION pg_catalog.citus_add_rebalance_strategy; +CREATE OR REPLACE FUNCTION pg_catalog.citus_add_rebalance_strategy( + name name, + shard_cost_function regproc, + node_capacity_function regproc, + shard_allowed_on_node_function regproc, + default_threshold float4, + minimum_threshold float4 DEFAULT 0, + improvement_threshold float4 DEFAULT 0 +) + RETURNS VOID AS $$ + INSERT INTO + pg_catalog.pg_dist_rebalance_strategy( + name, + shard_cost_function, + node_capacity_function, + shard_allowed_on_node_function, + default_threshold, + minimum_threshold, + improvement_threshold + ) VALUES ( + name, + shard_cost_function, + node_capacity_function, + shard_allowed_on_node_function, + default_threshold, + minimum_threshold, + improvement_threshold + ); + $$ LANGUAGE sql; +COMMENT ON FUNCTION pg_catalog.citus_add_rebalance_strategy(name,regproc,regproc,regproc,float4, float4, float4) + IS 'adds a new rebalance strategy which can be used when rebalancing shards or draining nodes'; diff --git a/src/backend/distributed/sql/udfs/citus_add_rebalance_strategy/latest.sql b/src/backend/distributed/sql/udfs/citus_add_rebalance_strategy/latest.sql index 4c5f8ba79..c4f157c2e 100644 --- a/src/backend/distributed/sql/udfs/citus_add_rebalance_strategy/latest.sql +++ b/src/backend/distributed/sql/udfs/citus_add_rebalance_strategy/latest.sql @@ -16,14 +16,16 @@ CREATE OR REPLACE FUNCTION pg_catalog.citus_add_rebalance_strategy( node_capacity_function, shard_allowed_on_node_function, default_threshold, - minimum_threshold + minimum_threshold, + improvement_threshold ) VALUES ( name, shard_cost_function, node_capacity_function, shard_allowed_on_node_function, default_threshold, - minimum_threshold + minimum_threshold, + improvement_threshold ); $$ LANGUAGE sql; COMMENT ON FUNCTION pg_catalog.citus_add_rebalance_strategy(name,regproc,regproc,regproc,float4, float4, float4) diff --git a/src/test/regress/expected/shard_rebalancer.out b/src/test/regress/expected/shard_rebalancer.out index 7997b5e28..f5b76c14c 100644 --- a/src/test/regress/expected/shard_rebalancer.out +++ b/src/test/regress/expected/shard_rebalancer.out @@ -2184,6 +2184,27 @@ SELECT citus_add_rebalance_strategy( 0.1 ); ERROR: default_threshold cannot be smaller than minimum_threshold +SELECT citus_add_rebalance_strategy( + 'test_improvement_threshold', + 'citus_shard_cost_1', + 'capacity_high_worker_2', + 'citus_shard_allowed_on_node_true', + 0.2, + 0.1, + 0.3 + ); + citus_add_rebalance_strategy +--------------------------------------------------------------------- + +(1 row) + +SELECT * FROM pg_dist_rebalance_strategy WHERE name='test_improvement_threshold'; + name | default_strategy | shard_cost_function | node_capacity_function | shard_allowed_on_node_function | default_threshold | minimum_threshold | improvement_threshold +--------------------------------------------------------------------- + test_improvement_threshold | f | citus_shard_cost_1 | capacity_high_worker_2 | citus_shard_allowed_on_node_true | 0.2 | 0.1 | 0.3 +(1 row) + +DELETE FROM pg_catalog.pg_dist_rebalance_strategy WHERE name='test_improvement_threshold'; -- Make it a data node again SELECT * from master_set_node_property('localhost', :worker_2_port, 'shouldhaveshards', true); master_set_node_property diff --git a/src/test/regress/expected/upgrade_rebalance_strategy_after.out b/src/test/regress/expected/upgrade_rebalance_strategy_after.out index da822fffd..4036af539 100644 --- a/src/test/regress/expected/upgrade_rebalance_strategy_after.out +++ b/src/test/regress/expected/upgrade_rebalance_strategy_after.out @@ -3,6 +3,6 @@ SELECT * FROM pg_catalog.pg_dist_rebalance_strategy ORDER BY name; --------------------------------------------------------------------- by_disk_size | f | citus_shard_cost_by_disk_size | citus_node_capacity_1 | citus_shard_allowed_on_node_true | 0.1 | 0.01 | 0.5 by_shard_count | f | citus_shard_cost_1 | citus_node_capacity_1 | citus_shard_allowed_on_node_true | 0 | 0 | 0 - custom_strategy | t | upgrade_rebalance_strategy.shard_cost_2 | upgrade_rebalance_strategy.capacity_high_worker_1 | upgrade_rebalance_strategy.only_worker_2 | 0.5 | 0.2 | 0 + custom_strategy | t | upgrade_rebalance_strategy.shard_cost_2 | upgrade_rebalance_strategy.capacity_high_worker_1 | upgrade_rebalance_strategy.only_worker_2 | 0.5 | 0.2 | 0.3 (3 rows) diff --git a/src/test/regress/sql/shard_rebalancer.sql b/src/test/regress/sql/shard_rebalancer.sql index 07efa8617..5d8e89b36 100644 --- a/src/test/regress/sql/shard_rebalancer.sql +++ b/src/test/regress/sql/shard_rebalancer.sql @@ -1229,6 +1229,20 @@ SELECT citus_add_rebalance_strategy( 0.1 ); +SELECT citus_add_rebalance_strategy( + 'test_improvement_threshold', + 'citus_shard_cost_1', + 'capacity_high_worker_2', + 'citus_shard_allowed_on_node_true', + 0.2, + 0.1, + 0.3 + ); + +SELECT * FROM pg_dist_rebalance_strategy WHERE name='test_improvement_threshold'; + +DELETE FROM pg_catalog.pg_dist_rebalance_strategy WHERE name='test_improvement_threshold'; + -- Make it a data node again SELECT * from master_set_node_property('localhost', :worker_2_port, 'shouldhaveshards', true); DROP TABLE tab;