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.

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"

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"

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';

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)

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

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)
 

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;