Go API: Add EnableRebalance and DisableRebalance API functions

Signed-off-by: Aaron Wilson <[email protected]>
NVIDIA · Aug 20, 2024 · f5deb20 · f5deb20
1 parent 4a4b131
commit f5deb20
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 0 deletions.
diff --git a/ais/test/config_test.go b/ais/test/config_test.go
@@ -358,3 +358,26 @@ func TestConfigOverrideAndResetCluster(t *testing.T) {
 			errWMConfigNotExpected, config.Disk.DiskUtilLowWM, daemonConfig.Disk.DiskUtilLowWM)
 	}
 }
+
+func TestConfigRebalance(t *testing.T) {
+	var (
+		oRebalance = tools.GetClusterConfig(t).Rebalance.Enabled
+	)
+	defer tools.SetClusterConfig(t, cos.StrKVs{
+		"rebalance.enabled": strconv.FormatBool(oRebalance),
+	})
+
+	tools.EnableRebalance(t)
+
+	nRebalance := tools.GetClusterConfig(t).Rebalance.Enabled
+
+	if !nRebalance {
+		t.Errorf("Rebalance was not enabled: current value %v, should be: %v", nRebalance, true)
+	}
+	tools.DisableRebalance(t)
+
+	nRebalance = tools.GetClusterConfig(t).Rebalance.Enabled
+	if nRebalance {
+		t.Errorf("Rebalance was not disabled: current value %v, should be: %v", nRebalance, false)
+	}
+}
diff --git a/api/cluster.go b/api/cluster.go
@@ -262,6 +262,23 @@ func SetClusterConfigUsingMsg(bp BaseParams, configToUpdate *cmn.ConfigToSet, tr
 	return err
 }
 
+func setRebalance(bp BaseParams, enabled bool) error {
+	configToSet := &cmn.ConfigToSet{
+		Rebalance: &cmn.RebalanceConfToSet{
+			Enabled: apc.Ptr(enabled),
+		},
+	}
+	return SetClusterConfigUsingMsg(bp, configToSet, false /*transient*/)
+}
+
+func EnableRebalance(bp BaseParams) error {
+	return setRebalance(bp, true)
+}
+
+func DisableRebalance(bp BaseParams) error {
+	return setRebalance(bp, false)
+}
+
 // all nodes: reset configuration to cluster defaults
 func ResetClusterConfig(bp BaseParams) error {
 	return _putCluster(bp, apc.ActMsg{Action: apc.ActResetConfig})

diff --git a/tools/client.go b/tools/client.go
@@ -659,6 +659,20 @@ func SetClusterConfigUsingMsg(t *testing.T, toUpdate *cmn.ConfigToSet) {
 	tassert.CheckFatal(t, err)
 }
 
+func EnableRebalance(t *testing.T) {
+	proxyURL := GetPrimaryURL()
+	bp := BaseAPIParams(proxyURL)
+	err := api.EnableRebalance(bp)
+	tassert.CheckError(t, err)
+}
+
+func DisableRebalance(t *testing.T) {
+	proxyURL := GetPrimaryURL()
+	bp := BaseAPIParams(proxyURL)
+	err := api.DisableRebalance(bp)
+	tassert.CheckError(t, err)
+}
+
 func SetRemAisConfig(t *testing.T, nvs cos.StrKVs) {
 	remoteBP := BaseAPIParams(RemoteCluster.URL)
 	err := api.SetClusterConfig(remoteBP, nvs, false /*transient*/)