From 7592d013fe83c67d1f9071d513fbfbd2cd17270b Mon Sep 17 00:00:00 2001
From: themarkchen <132556106+themarkchen@users.noreply.github.com>
Date: Thu, 4 May 2023 22:35:37 +0800
Subject: [PATCH] fix shell volume.balance bug (#4447)

---
 weed/shell/command_volume_balance.go      | 38 +++++------
 weed/shell/command_volume_balance_test.go | 79 ++++++++++++++++++++++-
 2 files changed, 93 insertions(+), 24 deletions(-)

diff --git a/weed/shell/command_volume_balance.go b/weed/shell/command_volume_balance.go
index 1c599b8a0..b2b65d5d2 100644
--- a/weed/shell/command_volume_balance.go
+++ b/weed/shell/command_volume_balance.go
@@ -3,14 +3,15 @@ package shell
 import (
 	"flag"
 	"fmt"
+	"io"
+	"os"
+	"time"
+
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/storage/erasure_coding"
 	"github.com/seaweedfs/seaweedfs/weed/storage/super_block"
 	"github.com/seaweedfs/seaweedfs/weed/storage/types"
 	"golang.org/x/exp/slices"
-	"io"
-	"os"
-	"time"
 
 	"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
 	"github.com/seaweedfs/seaweedfs/weed/storage/needle"
@@ -371,33 +372,24 @@ func isGoodMove(placement *super_block.ReplicaPlacement, existingReplicas []*Vol
 			return false
 		}
 	}
-	dcs, racks := make(map[string]bool), make(map[string]int)
+
+	// existing replicas except the one on sourceNode
+	existingReplicasExceptSourceNode := make([]*VolumeReplica, 0)
 	for _, replica := range existingReplicas {
 		if replica.location.dataNode.Id != sourceNode.info.Id {
-			dcs[replica.location.DataCenter()] = true
-			racks[replica.location.Rack()]++
+			existingReplicasExceptSourceNode = append(existingReplicasExceptSourceNode, replica)
 		}
 	}
 
-	dcs[targetNode.dc] = true
-	racks[fmt.Sprintf("%s %s", targetNode.dc, targetNode.rack)]++
-
-	if len(dcs) != placement.DiffDataCenterCount+1 {
-		return false
+	// target location
+	targetLocation := location{
+		dc:       targetNode.dc,
+		rack:     targetNode.rack,
+		dataNode: targetNode.info,
 	}
 
-	if len(racks) != placement.DiffRackCount+placement.DiffDataCenterCount+1 {
-		return false
-	}
-
-	for _, sameRackCount := range racks {
-		if sameRackCount != placement.SameRackCount+1 {
-			return false
-		}
-	}
-
-	return true
-
+	// check if this satisfies replication requirements
+	return satisfyReplicaPlacement(placement, existingReplicasExceptSourceNode, targetLocation)
 }
 
 func adjustAfterMove(v *master_pb.VolumeInformationMessage, volumeReplicas map[uint32][]*VolumeReplica, fullNode *Node, emptyNode *Node) {
diff --git a/weed/shell/command_volume_balance_test.go b/weed/shell/command_volume_balance_test.go
index 5bd170e71..20c5abdf8 100644
--- a/weed/shell/command_volume_balance_test.go
+++ b/weed/shell/command_volume_balance_test.go
@@ -1,9 +1,10 @@
 package shell
 
 import (
+	"testing"
+
 	"github.com/seaweedfs/seaweedfs/weed/storage/types"
 	"github.com/stretchr/testify/assert"
-	"testing"
 
 	"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
 	"github.com/seaweedfs/seaweedfs/weed/storage/super_block"
@@ -149,6 +150,82 @@ func TestIsGoodMove(t *testing.T) {
 			targetLocation: location{"dc1", "r2", &master_pb.DataNodeInfo{Id: "dn3"}},
 			expected:       true,
 		},
+
+		{
+			name:        "test 011 switch which rack has more replicas",
+			replication: "011",
+			replicas: []*VolumeReplica{
+				{
+					location: &location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn1"}},
+				},
+				{
+					location: &location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn2"}},
+				},
+				{
+					location: &location{"dc1", "r2", &master_pb.DataNodeInfo{Id: "dn3"}},
+				},
+			},
+			sourceLocation: location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn1"}},
+			targetLocation: location{"dc1", "r2", &master_pb.DataNodeInfo{Id: "dn4"}},
+			expected:       true,
+		},
+
+		{
+			name:        "test 011 move the lonely replica to another racks",
+			replication: "011",
+			replicas: []*VolumeReplica{
+				{
+					location: &location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn1"}},
+				},
+				{
+					location: &location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn2"}},
+				},
+				{
+					location: &location{"dc1", "r2", &master_pb.DataNodeInfo{Id: "dn3"}},
+				},
+			},
+			sourceLocation: location{"dc1", "r2", &master_pb.DataNodeInfo{Id: "dn3"}},
+			targetLocation: location{"dc1", "r3", &master_pb.DataNodeInfo{Id: "dn4"}},
+			expected:       true,
+		},
+
+		{
+			name:        "test 011 move to wrong racks",
+			replication: "011",
+			replicas: []*VolumeReplica{
+				{
+					location: &location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn1"}},
+				},
+				{
+					location: &location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn2"}},
+				},
+				{
+					location: &location{"dc1", "r2", &master_pb.DataNodeInfo{Id: "dn3"}},
+				},
+			},
+			sourceLocation: location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn1"}},
+			targetLocation: location{"dc1", "r3", &master_pb.DataNodeInfo{Id: "dn4"}},
+			expected:       false,
+		},
+
+		{
+			name:        "test 011 move all to the same rack",
+			replication: "011",
+			replicas: []*VolumeReplica{
+				{
+					location: &location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn1"}},
+				},
+				{
+					location: &location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn2"}},
+				},
+				{
+					location: &location{"dc1", "r2", &master_pb.DataNodeInfo{Id: "dn3"}},
+				},
+			},
+			sourceLocation: location{"dc1", "r2", &master_pb.DataNodeInfo{Id: "dn3"}},
+			targetLocation: location{"dc1", "r1", &master_pb.DataNodeInfo{Id: "dn4"}},
+			expected:       false,
+		},
 	}
 
 	for _, tt := range tests {