rebalance: adds -allowed-dsts to support backfilling nodes in imbalanced replicas

Sometimes we use rebalance command to backfill a storage node that replaces a broken node:

    bucky rebalance -h node-in-replica-1:4242 node-in-replica2:4242

This only works if replica 1 and replica 2 has the same number of nodes (assume using
jump hash).

With -allowed-dsts, we can backfill the node from a replica has higher or lower storage
nodes.

    bucky rebalance -h node-in-replica-1:4242 -allowed-dsts node-in-replica-1:4242 node1-in-replica2:4242,node2-in-replica2:4242,...,nodeN-in-replica2:4242

Author: bom-d-van

Makefile

@@ -44,9 +44,10 @@ test_rebalance_health_check_setup:
 	sudo ip addr add 10.0.1.9 dev lo
 
 test_rebalance_health_check: clean bucky buckyd
-# 	go run -mod vendor testing/rebalance_health_check.go $(REBALANCE_FLAGS)
+	go run -mod vendor testing/rebalance_health_check.go $(REBALANCE_FLAGS)
 
 clean_test:
+	rm -rf bucky buckyd
 	rm -rf testdata_rebalance_*
 	rm -rf testdata_copy_*
 

cmd/bucky/rebalance.go

@@ -4,8 +4,13 @@ import (
 	"fmt"
 	"log"
 	"sort"
+	"strings"
 )
 
+var rebalanceConfig struct {
+	allowedDsts string
+}
+
 func init() {
 	usage := "[options] [additional buckyd servers...]"
 	short := "Rebalance a server or the entire cluster."
@@ -43,8 +48,8 @@ Set -offload=true to speed up rebalance.`
 	SetupSingle(c)
 
 	msFlags.registerFlags(c.Flag)
-	c.Flag.BoolVar(&listForce, "f", false,
-		"Force the remote daemons to rebuild their cache.")
+	c.Flag.BoolVar(&listForce, "f", false, "Force the remote daemons to rebuild their cache.")
+	c.Flag.StringVar(&rebalanceConfig.allowedDsts, "allowed-dsts", "", "Only copy/rebanace metrics to the allowed destinations (ip1:port,ip2:port). By default (i.e. empty), all dsts are allowed.")
 }
 
 // countMap returns the number of metrics in a server -> metrics mapping
@@ -109,6 +114,23 @@ func RebalanceMetrics(extraHostPorts []string) error {
 		}
 	}
 
+	if rebalanceConfig.allowedDsts != "" {
+		allowm := map[string]bool{}
+		for _, hostport := range strings.Split(rebalanceConfig.allowedDsts, ";") {
+			allowm[strings.TrimSpace(hostport)] = true
+		}
+
+		newJobs := map[string]map[string][]*syncJob{}
+		for dst, srcm := range jobs {
+			if allowm[dst] {
+				newJobs[dst] = srcm
+			}
+
+		}
+
+		jobs = newJobs
+	}
+
 	sort.Strings(servers)
 	for _, server := range servers {
 		log.Printf("%d metrics on %s must be relocated", moves[server], server)

testing/rebalance.go

@@ -152,7 +152,13 @@ func main() {
 
 	time.Sleep(time.Second * 3)
 	rebalanceStart := time.Now()
-	rebalanceCmd := exec.Command("./bucky", "rebalance", "-f", "-h", nodeStr(server0), "-offload", "-w", "3", "-ignore404")
+	rebalanceCmd := exec.Command(
+		"./bucky", "rebalance", "-f",
+		"-h", nodeStr(server0), "-offload",
+		"-w", "3", "-ignore404",
+		// "-allowed-dsts", "localhost:40002",
+		// "-allowed-dsts", "xxx:xxx",
+	)
 
 	rebalanceCmd.Stdout = rebalanceLog
 	rebalanceCmd.Stderr = rebalanceLog