From 6e4b9181f5742b5a71c49ffb8f8e22ed0754ef5f Mon Sep 17 00:00:00 2001
From: skycope <wings.wyang@gmail.com>
Date: Tue, 23 Apr 2024 21:33:50 +0800
Subject: [PATCH] fix "volume.fix.replication" move many replications only to
 one volumeServer (#5522)

---
 weed/shell/command_volume_balance.go         |  6 ++----
 weed/shell/command_volume_fix_replication.go | 21 ++++++++++++++------
 weed/shell/command_volume_tier_move.go       | 11 +++++-----
 3 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/weed/shell/command_volume_balance.go b/weed/shell/command_volume_balance.go
index cb201e064..e9a483f41 100644
--- a/weed/shell/command_volume_balance.go
+++ b/weed/shell/command_volume_balance.go
@@ -402,14 +402,12 @@ func adjustAfterMove(v *master_pb.VolumeInformationMessage, volumeReplicas map[u
 			replica.location = &loc
 			for diskType, diskInfo := range fullNode.info.DiskInfos {
 				if diskType == v.DiskType {
-					diskInfo.VolumeCount--
-					diskInfo.FreeVolumeCount++
+					addVolumeCount(diskInfo, -1)
 				}
 			}
 			for diskType, diskInfo := range emptyNode.info.DiskInfos {
 				if diskType == v.DiskType {
-					diskInfo.VolumeCount++
-					diskInfo.FreeVolumeCount--
+					addVolumeCount(diskInfo, 1)
 				}
 			}
 			return
diff --git a/weed/shell/command_volume_fix_replication.go b/weed/shell/command_volume_fix_replication.go
index b724f16f9..074931f40 100644
--- a/weed/shell/command_volume_fix_replication.go
+++ b/weed/shell/command_volume_fix_replication.go
@@ -4,16 +4,17 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"io"
+	"path/filepath"
+	"strconv"
+	"time"
+
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/storage/needle"
 	"github.com/seaweedfs/seaweedfs/weed/storage/needle_map"
 	"github.com/seaweedfs/seaweedfs/weed/storage/types"
 	"golang.org/x/exp/slices"
 	"google.golang.org/grpc"
-	"io"
-	"path/filepath"
-	"strconv"
-	"time"
 
 	"github.com/seaweedfs/seaweedfs/weed/operation"
 	"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
@@ -316,7 +317,7 @@ func (c *commandVolumeFixReplication) fixOneUnderReplicatedVolume(commandEnv *Co
 
 			if !takeAction {
 				// adjust volume count
-				dst.dataNode.DiskInfos[replica.info.DiskType].VolumeCount++
+				addVolumeCount(dst.dataNode.DiskInfos[replica.info.DiskType], 1)
 				break
 			}
 
@@ -350,7 +351,7 @@ func (c *commandVolumeFixReplication) fixOneUnderReplicatedVolume(commandEnv *Co
 			}
 
 			// adjust volume count
-			dst.dataNode.DiskInfos[replica.info.DiskType].VolumeCount++
+			addVolumeCount(dst.dataNode.DiskInfos[replica.info.DiskType], 1)
 			break
 		}
 	}
@@ -361,6 +362,14 @@ func (c *commandVolumeFixReplication) fixOneUnderReplicatedVolume(commandEnv *Co
 	return nil
 }
 
+func addVolumeCount(info *master_pb.DiskInfo, count int) {
+	if info == nil {
+		return
+	}
+	info.VolumeCount += int64(count)
+	info.FreeVolumeCount -= int64(count)
+}
+
 func keepDataNodesSorted(dataNodes []location, diskType types.DiskType) {
 	fn := capacityByFreeVolumeCount(diskType)
 	slices.SortFunc(dataNodes, func(a, b location) int {
diff --git a/weed/shell/command_volume_tier_move.go b/weed/shell/command_volume_tier_move.go
index e6cf4ee02..c6364757f 100644
--- a/weed/shell/command_volume_tier_move.go
+++ b/weed/shell/command_volume_tier_move.go
@@ -5,15 +5,16 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"io"
+	"path/filepath"
+	"sync"
+	"time"
+
 	"github.com/seaweedfs/seaweedfs/weed/glog"
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
 	"github.com/seaweedfs/seaweedfs/weed/storage/types"
 	"github.com/seaweedfs/seaweedfs/weed/wdclient"
-	"io"
-	"path/filepath"
-	"sync"
-	"time"
 
 	"github.com/seaweedfs/seaweedfs/weed/operation"
 	"github.com/seaweedfs/seaweedfs/weed/pb/volume_server_pb"
@@ -212,7 +213,7 @@ func (c *commandVolumeTierMove) doVolumeTierMove(commandEnv *CommandEnv, writer
 			hasFoundTarget = true
 
 			// adjust volume count
-			dst.dataNode.DiskInfos[string(toDiskType)].VolumeCount++
+			addVolumeCount(dst.dataNode.DiskInfos[string(toDiskType)], 1)
 
 			destServerAddress := pb.NewServerAddressFromDataNode(dst.dataNode)
 			c.queues[destServerAddress] <- volumeTierMoveJob{sourceVolumeServer, vid}