Limit EC re-balancing for `ec.encode` to relevant collections when a volume ID argument is provided. (#6347)

Limit EC re-balancing for `ec.encode` to relevant collections when a volume ID is provided.

3 files changed, 76 insertions, 5 deletions

diff --git a/weed/shell/command_ec_common.go b/weed/shell/command_ec_common.go
index 4b672b2ef..30667896a 100644
--- a/weed/shell/command_ec_common.go
+++ b/weed/shell/command_ec_common.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"math/rand/v2"
+	"sort"
 	"time"
 
 	"github.com/seaweedfs/seaweedfs/weed/glog"
@@ -182,6 +183,46 @@ func collectEcNodes(commandEnv *CommandEnv) (ecNodes []*EcNode, totalFreeEcSlots
 	return collectEcNodesForDC(commandEnv, "")
 }
 
+func collectCollectionsForVolumeIds(t *master_pb.TopologyInfo, vids []needle.VolumeId) []string {
+	if len(vids) == 0 {
+		return nil
+	}
+
+	found := map[string]bool{}
+	for _, dc := range t.DataCenterInfos {
+		for _, r := range dc.RackInfos {
+			for _, dn := range r.DataNodeInfos {
+				for _, diskInfo := range dn.DiskInfos {
+					for _, vi := range diskInfo.VolumeInfos {
+						for _, vid := range vids {
+							if needle.VolumeId(vi.Id) == vid && vi.Collection != "" {
+								found[vi.Collection] = true
+							}
+						}
+					}
+					for _, ecs := range diskInfo.EcShardInfos {
+						for _, vid := range vids {
+							if needle.VolumeId(ecs.Id) == vid && ecs.Collection != "" {
+								found[ecs.Collection] = true
+							}
+						}
+					}
+				}
+			}
+		}
+	}
+	if len(found) == 0 {
+		return nil
+	}
+
+	collections := []string{}
+	for k, _ := range found {
+		collections = append(collections, k)
+	}
+	sort.Strings(collections)
+	return collections
+}
+
 func moveMountedShardToEcNode(commandEnv *CommandEnv, existingLocation *EcNode, collection string, vid needle.VolumeId, shardId erasure_coding.ShardId, destinationEcNode *EcNode, applyBalancing bool) (err error) {
 
 	if !commandEnv.isLocked() {
diff --git a/weed/shell/command_ec_common_test.go b/weed/shell/command_ec_common_test.go
index e9c9b1d99..f76840f3c 100644
--- a/weed/shell/command_ec_common_test.go
+++ b/weed/shell/command_ec_common_test.go
@@ -2,6 +2,7 @@ package shell
 
 import (
 	"fmt"
+	"reflect"
 	"strings"
 	"testing"
 
@@ -33,6 +34,39 @@ func errorCheck(got error, want string) error {
 	return nil
 }
 
+func TestCollectCollectionsForVolumeIds(t *testing.T) {
+	testCases := []struct {
+		topology *master_pb.TopologyInfo
+		vids     []needle.VolumeId
+		want     []string
+	}{
+		// normal volumes
+		{topology1, nil, nil},
+		{topology1, []needle.VolumeId{}, nil},
+		{topology1, []needle.VolumeId{needle.VolumeId(9999)}, nil},
+		{topology1, []needle.VolumeId{needle.VolumeId(2)}, nil},
+		{topology1, []needle.VolumeId{needle.VolumeId(2), needle.VolumeId(272)}, []string{"collection2"}},
+		{topology1, []needle.VolumeId{needle.VolumeId(2), needle.VolumeId(272), needle.VolumeId(299)}, []string{"collection2"}},
+		{topology1, []needle.VolumeId{needle.VolumeId(272), needle.VolumeId(299), needle.VolumeId(95)}, []string{"collection1", "collection2"}},
+		{topology1, []needle.VolumeId{needle.VolumeId(272), needle.VolumeId(299), needle.VolumeId(95), needle.VolumeId(51)}, []string{"collection1", "collection2"}},
+		{topology1, []needle.VolumeId{needle.VolumeId(272), needle.VolumeId(299), needle.VolumeId(95), needle.VolumeId(51), needle.VolumeId(15)}, []string{"collection0", "collection1", "collection2"}},
+		// EC volumes
+		{topology2, []needle.VolumeId{needle.VolumeId(9577)}, []string{"s3qldata"}},
+		{topology2, []needle.VolumeId{needle.VolumeId(9577), needle.VolumeId(12549)}, []string{"s3qldata"}},
+		// normal + EC volumes
+		{topology2, []needle.VolumeId{needle.VolumeId(18111)}, []string{"s3qldata"}},
+		{topology2, []needle.VolumeId{needle.VolumeId(8677)}, []string{"s3qldata"}},
+		{topology2, []needle.VolumeId{needle.VolumeId(18111), needle.VolumeId(8677)}, []string{"s3qldata"}},
+	}
+
+	for _, tc := range testCases {
+		got := collectCollectionsForVolumeIds(tc.topology, tc.vids)
+		if !reflect.DeepEqual(got, tc.want) {
+			t.Errorf("for %v: got %v, want %v", tc.vids, got, tc.want)
+		}
+	}
+}
+
 func TestParseReplicaPlacementArg(t *testing.T) {
 	getDefaultReplicaPlacementOrig := getDefaultReplicaPlacement
 	getDefaultReplicaPlacement = func(commandEnv *CommandEnv) (*super_block.ReplicaPlacement, error) {
diff --git a/weed/shell/command_ec_encode.go b/weed/shell/command_ec_encode.go
index 6ee530256..f8b881d7c 100644
--- a/weed/shell/command_ec_encode.go
+++ b/weed/shell/command_ec_encode.go
@@ -115,11 +115,7 @@ func (c *commandEcEncode) Do(args []string, commandEnv *CommandEnv, writer io.Wr
 	if *collection != "" {
 		collections = []string{*collection}
 	} else {
-		// TODO: should we limit this to collections associated with the provided volume ID?
-		collections, err = ListCollectionNames(commandEnv, false, true)
-		if err != nil {
-			return err
-		}
+		collections = collectCollectionsForVolumeIds(topologyInfo, volumeIds)
 	}
 
 	// encode all requested volumes...