refactor(ec_balance): make the balance planner per-volume ratio-capable (#9960)

* refactor(ec_balance): make the balance planner per-volume ratio-capable

Thread a per-volume EC ratio through the balance planner: Plan resolves each
volume's data/parity from a new Options.VolumeRatio (falling back to the
collection Ratio, then the build default, when it reports 0), and keys the
global phase's ratio maps by volume instead of collection. The shell and
worker balance paths build the per-volume lookup from each shard's heartbeat
via the new ecbalancer.VolumeShardRatio.

In OSS this is behavior-preserving: VolumeShardRatio returns 0 because the
per-volume data_shards/parity_shards heartbeat fields are an enterprise
feature, so every volume falls back to the collection ratio -- the existing
standard-scheme behavior. The refactor keeps the shared planner in sync with
the enterprise fork, which overrides VolumeShardRatio to classify and spread
a mixed-ratio collection by each volume's own data/parity split.

* perf(ec_balance): hoist the collection ratio out of the per-volume loop

The collection ratio is constant for every volume in a collection, so
resolve it once per collection instead of per volume; a custom Ratio func
may do map lookups or locking. Addresses a review comment.

weed/shell/command_ec_common.go

@@ -898,12 +898,16 @@ func shellECRatio(_ string) (int, int) {
 // balance plans EC shard moves with the shared planner and executes them. When
 // collections is empty all collections present are balanced.
 func (ecb *ecBalancer) balance(collections []string) error {
-	topo := toBalancerTopology(ecb.ecNodes, collections, ecb.diskType)
+	topo, volumeRatio := toBalancerTopology(ecb.ecNodes, collections, ecb.diskType)
 	moves := ecbalancer.Plan(topo, ecbalancer.Options{
 		DiskType:           string(ecb.diskType),
 		ImbalanceThreshold: 0, // the shell balances to an even distribution
 		ReplicaPlacement:   ecb.replicaPlacement,
 		Ratio:              shellECRatio,
+		// Prefer each volume's own heartbeat-reported ratio over the collection
+		// default so a mixed-ratio collection is spread per volume; 0 defers to
+		// shellECRatio (and is the always-0 OSS case).
+		VolumeRatio: volumeRatio,
 		// Balance the global phase by fractional fullness so heterogeneous-capacity
 		// nodes fill proportionally (matching the worker). This is identical to raw
 		// shard count when capacities are uniform.
@@ -914,12 +918,21 @@ func (ecb *ecBalancer) balance(collections []string) error {
 
 // toBalancerTopology builds an ecbalancer.Topology from the shell's EcNode model,
 // including the shards of the requested collections (all collections when empty).
-func toBalancerTopology(ecNodes []*EcNode, collections []string, diskType types.DiskType) *ecbalancer.Topology {
+// It also returns a per-volume ratio lookup built from each shard's heartbeat
+// (0,0 when unreported, e.g. always in OSS), which Plan prefers over the
+// collection ratio for mixed-ratio clusters.
+func toBalancerTopology(ecNodes []*EcNode, collections []string, diskType types.DiskType) (*ecbalancer.Topology, func(collection string, vid uint32) (int, int)) {
 	allowed := make(map[string]bool, len(collections))
 	for _, c := range collections {
 		allowed[c] = true
 	}
 
+	type volRatioKey struct {
+		collection string
+		vid        uint32
+	}
+	volRatios := make(map[volRatioKey][2]int)
+
 	topo := ecbalancer.NewTopology()
 	for _, en := range ecNodes {
 		rackKey := string(en.dc) + ":" + string(en.rack)
@@ -939,9 +952,17 @@ func toBalancerTopology(ecNodes []*EcNode, collections []string, diskType types.
 				continue
 			}
 			node.AddShards(eci.Id, eci.Collection, eci.DiskId, erasure_coding.ShardBits(eci.EcIndexBits))
+			if d, p := ecbalancer.VolumeShardRatio(eci); d > 0 || p > 0 {
+				volRatios[volRatioKey{eci.Collection, eci.Id}] = [2]int{d, p}
+			}
 		}
 	}
-	return topo
+
+	volumeRatio := func(collection string, vid uint32) (int, int) {
+		r := volRatios[volRatioKey{collection, vid}]
+		return r[0], r[1]
+	}
+	return topo, volumeRatio
 }
 
 // executeMoves carries out the planned moves. Phases run in order (a within-rack

weed/storage/erasure_coding/ecbalancer/balancer.go

@@ -84,6 +84,14 @@ type Options struct {
 	// Ratio returns a collection's (dataShards, parityShards); nil defaults to the
 	// standard scheme. This is where a caller plugs in custom-ratio resolution.
 	Ratio func(collection string) (dataShards, parityShards int)
+	// VolumeRatio returns a single volume's (dataShards, parityShards) when its
+	// heartbeat reported a per-volume ratio; either return <=0 to defer to Ratio
+	// (the collection ratio). A single collection can hold volumes of mixed ratios
+	// (a ratio change after some volumes were encoded), so placement must classify
+	// and spread each volume by its OWN data/parity split, not the collection's.
+	// nil (and the 0/OSS case) makes the planner fall back to Ratio per collection,
+	// preserving the collection-keyed behavior.
+	VolumeRatio func(collection string, vid uint32) (dataShards, parityShards int)
 	// GlobalMaxMovesPerRack caps how many shards the global (cross-volume) phase
 	// moves out of one rack in a single Plan. 0 means unlimited (drain to balance
 	// in one pass), which the shell uses; the worker sets a small value to make
@@ -185,8 +193,7 @@ func Plan(topo *Topology, opts Options) []Move {
 	racks := buildRacks(nodes)
 
 	// Group volumes by collection (deterministic order), keyed by (collection, id)
-	// so volumes that reuse a numeric id across collections stay distinct. Resolve
-	// each collection's data-shard count once for the global phase's disk scoring.
+	// so volumes that reuse a numeric id across collections stay distinct.
 	byCollection := make(map[string][]volKey)
 	seen := make(map[volKey]bool)
 	for _, n := range nodes {
@@ -198,39 +205,57 @@ func Plan(topo *Topology, opts Options) []Move {
 		}
 	}
 	collections := make([]string, 0, len(byCollection))
-	dataShardsByCollection := make(map[string]int)
-	parityShardsByCollection := make(map[string]int)
 	for c := range byCollection {
 		collections = append(collections, c)
 		sort.Slice(byCollection[c], func(i, j int) bool { return byCollection[c][i].vid < byCollection[c][j].vid })
-		d, p := ratio(c)
-		dataShardsByCollection[c] = d
-		parityShardsByCollection[c] = p
 	}
 	sort.Strings(collections)
 
+	// Resolve each volume's data/parity split: prefer the per-volume ratio the
+	// heartbeat reported (Options.VolumeRatio), fall back to the collection ratio,
+	// then the build defaults via `ratio`. Keyed by volume so a mixed-ratio
+	// collection (e.g. a 9+3 volume beside a 10+4 one) is classified and spread per
+	// volume rather than against one collection-wide split.
+	dataShardsByVolume := make(map[volKey]int)
+	parityShardsByVolume := make(map[volKey]int)
+	for _, collection := range collections {
+		defaultD, defaultP := ratio(collection)
+		for _, vk := range byCollection[collection] {
+			d, p := defaultD, defaultP
+			if opts.VolumeRatio != nil {
+				vd, vp := opts.VolumeRatio(vk.collection, vk.vid)
+				if vd > 0 {
+					d = vd
+				}
+				if vp > 0 {
+					p = vp
+				}
+			}
+			dataShardsByVolume[vk] = d
+			parityShardsByVolume[vk] = p
+		}
+	}
+
 	var all []*move
 	for _, collection := range collections {
-		dataShards, parityShards := ratio(collection)
-
 		for _, vk := range byCollection[collection] {
 			m := detectDuplicateShards(vk, nodes)
 			applyMovesToTopology(m, racks)
 			all = append(all, m...)
 		}
 		for _, vk := range byCollection[collection] {
-			m := detectCrossRackImbalance(vk, nodes, racks, opts.DiskType, opts.ImbalanceThreshold, dataShards, parityShards, opts.ReplicaPlacement)
+			m := detectCrossRackImbalance(vk, nodes, racks, opts.DiskType, opts.ImbalanceThreshold, dataShardsByVolume[vk], parityShardsByVolume[vk], opts.ReplicaPlacement)
 			applyMovesToTopology(m, racks)
 			all = append(all, m...)
 		}
 		for _, vk := range byCollection[collection] {
-			m := detectWithinRackImbalance(vk, nodes, racks, opts.DiskType, opts.ImbalanceThreshold, dataShards, parityShards, opts.ReplicaPlacement)
+			m := detectWithinRackImbalance(vk, nodes, racks, opts.DiskType, opts.ImbalanceThreshold, dataShardsByVolume[vk], parityShardsByVolume[vk], opts.ReplicaPlacement)
 			applyMovesToTopology(m, racks)
 			all = append(all, m...)
 		}
 	}
 
-	all = append(all, detectGlobalImbalance(nodes, racks, opts.DiskType, opts.ImbalanceThreshold, dataShardsByCollection, parityShardsByCollection, opts.GlobalMaxMovesPerRack, opts.GlobalUtilizationBased)...)
+	all = append(all, detectGlobalImbalance(nodes, racks, opts.DiskType, opts.ImbalanceThreshold, dataShardsByVolume, parityShardsByVolume, opts.GlobalMaxMovesPerRack, opts.GlobalUtilizationBased)...)
 
 	out := make([]Move, 0, len(all))
 	for _, m := range all {
@@ -650,7 +675,7 @@ func balanceShardTypeAcrossNodes(vk volKey, r *rack, diskType string, dataShards
 // detectGlobalImbalance balances total EC shard load across the nodes of each
 // rack (across all volumes), using utilization ratios so heterogeneous-capacity
 // nodes are compared fairly.
-func detectGlobalImbalance(nodes map[string]*Node, racks map[string]*rack, diskType string, threshold float64, dataShardsByCollection, parityShardsByCollection map[string]int, maxMovesPerRack int, byUtilization bool) []*move {
+func detectGlobalImbalance(nodes map[string]*Node, racks map[string]*rack, diskType string, threshold float64, dataShardsByVolume, parityShardsByVolume map[volKey]int, maxMovesPerRack int, byUtilization bool) []*move {
 	var moves []*move
 
 	for _, rackID := range sortedKeys(racks) {
@@ -754,7 +779,7 @@ func detectGlobalImbalance(nodes map[string]*Node, racks map[string]*rack, diskT
 						// doesn't raise the destination machine's count past the source's.
 						// Where it isn't achievable, capacity rules and any leveling move
 						// is fine. Feasibility uses the rack's shards, not the whole volume.
-						parity := parityShardsByCollection[vk.collection]
+						parity := parityShardsByVolume[vk]
 						spreadFeasible := parity > 0 && rackMachineCount >= ceilDivide(rackVolumeShardCount(r, vk), parity)
 						if spreadFeasible && minNode.host != maxNode.host &&
 							machineVolumeCount(r, minNode.host, vk) >= machineVolumeCount(r, maxNode.host, vk) {
@@ -768,7 +793,7 @@ func detectGlobalImbalance(nodes map[string]*Node, racks map[string]*rack, diskT
 						if minInfo != nil && minInfo.shardBits.Has(sid) {
 							continue
 						}
-						dataShards := dataShardsByCollection[vk.collection]
+						dataShards := dataShardsByVolume[vk]
 						if dataShards <= 0 {
 							dataShards = erasure_coding.DataShardsCount
 						}

weed/storage/erasure_coding/ecbalancer/balancer_test.go

@@ -429,6 +429,60 @@ func TestPlanBalancesSkewedDataParityWithEvenTotals(t *testing.T) {
 	}
 }
 
+// TestPlanVolumeRatioOverridesCollection verifies the per-volume ratio (reported on
+// the heartbeat, surfaced via Options.VolumeRatio) takes precedence over the
+// collection ratio, so a mixed-ratio collection is classified and spread by each
+// volume's own data/parity split. The same 14-shard volume is planned three ways:
+// the per-volume 7+7 override must reproduce the 7+7-collection plan and differ from
+// the 10+4-collection plan. A VolumeRatio that returns 0 defers to the collection
+// ratio (the always-0 OSS case), so existing collection-keyed behavior is preserved.
+func TestPlanVolumeRatioOverridesCollection(t *testing.T) {
+	build := func() *Topology {
+		topo := NewTopology()
+		n1 := topo.AddNode("node1", "dc1", "dc1:rack1", 100)
+		n1.AddDisk(0, "", 100, 7)
+		n1.AddShards(100, "col1", 0, bits(0, 1, 2, 3, 4, 5, 6))
+		n2 := topo.AddNode("node2", "dc1", "dc1:rack2", 100)
+		n2.AddDisk(0, "", 100, 7)
+		n2.AddShards(100, "col1", 0, bits(7, 8, 9, 10, 11, 12, 13))
+		return topo
+	}
+	crossRack := func(moves []Move) int {
+		n := 0
+		for _, m := range moves {
+			if m.Phase == "cross_rack" {
+				n++
+			}
+		}
+		return n
+	}
+
+	collection104 := crossRack(Plan(build(), Options{ImbalanceThreshold: 0, Ratio: ratio(10, 4)}))
+	collection77 := crossRack(Plan(build(), Options{ImbalanceThreshold: 0, Ratio: ratio(7, 7)}))
+	if collection104 == collection77 {
+		t.Fatalf("test setup: 10+4 and 7+7 collection plans must differ, both gave %d cross-rack moves", collection104)
+	}
+
+	// Collection ratio stays 10+4, but vol100 reports 7+7 per-volume; the planner
+	// must classify/spread vol100 as 7+7 and match the 7+7-collection plan.
+	perVolume := crossRack(Plan(build(), Options{
+		ImbalanceThreshold: 0,
+		Ratio:              ratio(10, 4),
+		VolumeRatio: func(collection string, vid uint32) (int, int) {
+			if collection == "col1" && vid == 100 {
+				return 7, 7
+			}
+			return 0, 0
+		},
+	}))
+	if perVolume != collection77 {
+		t.Errorf("per-volume 7+7 override gave %d cross-rack moves, want %d (the 7+7-collection plan)", perVolume, collection77)
+	}
+	if perVolume == collection104 {
+		t.Errorf("per-volume override had no effect: %d cross-rack moves, same as the 10+4 collection plan", perVolume)
+	}
+}
+
 // TestGlobalPrefersVolumeAbsentFromDestination guards the global phase's
 // volume-diversity preference: when draining a node, move a shard of a volume the
 // destination does not hold at all before piling a second shard of an
@@ -737,8 +791,8 @@ func TestGlobalDoesNotConcentrateVolumeAcrossMachines(t *testing.T) {
 	b2 := topo.AddNode("b2", "dc1", "dc1:rack1", 10) // empty -> low util, the min node
 	b2.SetHost("boxB")
 
-	data := map[string]int{"col1": 2}
-	parity := map[string]int{"col1": 2}
+	data := map[volKey]int{{collection: "col1", vid: 100}: 2}
+	parity := map[volKey]int{{collection: "col1", vid: 100}: 2}
 	for _, m := range detectGlobalImbalance(topo.nodes, buildRacks(topo.nodes), "", 0.01, data, parity, 0, true) {
 		if m.source.host != m.target.host {
 			t.Errorf("cross-machine global move %d.%d from %s to %s concentrates the volume on a machine",

weed/storage/erasure_coding/ecbalancer/shard_ratio.go

@@ -12,3 +12,13 @@ import (
 func shardDataShards(eci *master_pb.VolumeEcShardInformationMessage) int {
 	return erasure_coding.DataShardsCount
 }
+
+// VolumeShardRatio returns the RAW per-volume (dataShards, parityShards) reported
+// on an EC shard's heartbeat, with 0 meaning "not reported". Custom per-volume
+// ratios are an enterprise feature and the OSS proto has no data_shards/parity_shards
+// fields, so this returns 0, 0 and the balancer falls back to the collection ratio
+// (the standard scheme). The enterprise build overrides this to read the per-shard
+// ratio so a mixed-ratio collection is spread by each volume's own data/parity split.
+func VolumeShardRatio(eci *master_pb.VolumeEcShardInformationMessage) (dataShards, parityShards int) {
+	return 0, 0
+}

weed/worker/tasks/ec_balance/detection.go

@@ -47,7 +47,7 @@ func Detection(
 		return nil, false, fmt.Errorf("topology info not available")
 	}
 
-	topo, nodeCount := buildBalancerTopology(topoInfo, ecConfig)
+	topo, nodeCount, volumeRatio := buildBalancerTopology(topoInfo, ecConfig)
 	if nodeCount < ecConfig.MinServerCount {
 		glog.V(1).Infof("EC balance: only %d servers, need at least %d", nodeCount, ecConfig.MinServerCount)
 		return nil, false, nil
@@ -72,6 +72,10 @@ func Detection(
 		Ratio: func(collection string) (int, int) {
 			return resolveECRatio(clusterInfo, collection)
 		},
+		// Prefer each volume's own heartbeat-reported ratio over the collection
+		// default so a mixed-ratio collection is spread per volume; 0 defers to
+		// resolveECRatio (and is the always-0 OSS case).
+		VolumeRatio: volumeRatio,
 		// Move incrementally across detection cycles rather than draining a rack
 		// in one batch; the scheduler re-evaluates each cycle.
 		GlobalMaxMovesPerRack: 10,
@@ -137,11 +141,20 @@ func Detection(
 // applying the data-center, disk-type, and collection filters. Rack keys are
 // dc:rack composites to avoid cross-DC name collisions. Per-disk free capacity
 // is split evenly from the node total because the wire collapses same-type disks.
-// Returns the topology and the number of eligible nodes (for MinServerCount).
-func buildBalancerTopology(topoInfo *master_pb.TopologyInfo, config *Config) (*ecbalancer.Topology, int) {
+// Returns the topology, the number of eligible nodes (for MinServerCount), and a
+// per-volume ratio lookup built from each shard's heartbeat (0,0 when unreported,
+// e.g. always in OSS) which Plan prefers over the collection ratio for mixed-ratio
+// clusters.
+func buildBalancerTopology(topoInfo *master_pb.TopologyInfo, config *Config) (*ecbalancer.Topology, int, func(collection string, vid uint32) (int, int)) {
 	topo := ecbalancer.NewTopology()
 	allowedCollections := wildcard.CompileWildcardMatchers(config.CollectionFilter)
 
+	type volRatioKey struct {
+		collection string
+		vid        uint32
+	}
+	volRatios := make(map[volRatioKey][2]int)
+
 	// Normalize the disk-type filter: "hdd" (and the default "") map to the
 	// HardDriveType, which the topology reports under the empty-string key. Keep a
 	// separate "filter requested" flag so a configured "hdd" still filters to HDD
@@ -222,6 +235,9 @@ func buildBalancerTopology(topoInfo *master_pb.TopologyInfo, config *Config) (*e
 							continue
 						}
 						node.AddShards(eci.Id, eci.Collection, eci.DiskId, erasure_coding.ShardBits(eci.EcIndexBits))
+						if d, p := ecbalancer.VolumeShardRatio(eci); d > 0 || p > 0 {
+							volRatios[volRatioKey{eci.Collection, eci.Id}] = [2]int{d, p}
+						}
 					}
 				}
 
@@ -230,7 +246,11 @@ func buildBalancerTopology(topoInfo *master_pb.TopologyInfo, config *Config) (*e
 		}
 	}
 
-	return topo, nodeCount
+	volumeRatio := func(collection string, vid uint32) (int, int) {
+		r := volRatios[volRatioKey{collection, vid}]
+		return r[0], r[1]
+	}
+	return topo, nodeCount, volumeRatio
 }
 
 // resolveECRatio returns the (dataShards, parityShards) for a collection from the

weed/worker/tasks/ec_balance/detection_test.go

@@ -41,7 +41,7 @@ func ecTopo(node1Collection string) *master_pb.TopologyInfo {
 
 func TestBuildBalancerTopology(t *testing.T) {
 	config := NewDefaultConfig()
-	topo, nodeCount := buildBalancerTopology(ecTopo("col1"), config)
+	topo, nodeCount, _ := buildBalancerTopology(ecTopo("col1"), config)
 	if nodeCount != 2 {
 		t.Fatalf("nodeCount = %d, want 2", nodeCount)
 	}
@@ -77,7 +77,7 @@ func TestBuildBalancerTopologyGroupsByHost(t *testing.T) {
 		}},
 	}
 
-	topo, _ := buildBalancerTopology(topoInfo, NewDefaultConfig())
+	topo, _, _ := buildBalancerTopology(topoInfo, NewDefaultConfig())
 	moves := ecbalancer.Plan(topo, ecbalancer.Options{ImbalanceThreshold: 0.01})
 
 	host := func(nodeID string) string { h, _, _ := net.SplitHostPort(nodeID); return h }
@@ -99,7 +99,7 @@ func TestBuildBalancerTopologyGroupsByHost(t *testing.T) {
 func TestBuildBalancerTopologyCollectionFilter(t *testing.T) {
 	config := NewDefaultConfig()
 	config.CollectionFilter = "other" // does not match the volume's collection
-	topo, nodeCount := buildBalancerTopology(ecTopo("col1"), config)
+	topo, nodeCount, _ := buildBalancerTopology(ecTopo("col1"), config)
 	if nodeCount != 2 {
 		t.Fatalf("nodeCount = %d, want 2", nodeCount)
 	}

weed/worker/tasks/ec_balance/multidisk_detection_test.go

@@ -295,13 +295,13 @@ func TestBuildBalancerTopologyNormalizesHddDiskType(t *testing.T) {
 	}
 	topoInfo := buildMasterTopology("c", 100, 50, specs)
 
-	if _, n := buildBalancerTopology(topoInfo, &Config{DiskType: "hdd"}); n != 2 {
+	if _, n, _ := buildBalancerTopology(topoInfo, &Config{DiskType: "hdd"}); n != 2 {
 		t.Errorf("disk_type=hdd matched %d nodes on an all-HDD cluster, want 2 (hdd must map to the empty HDD key)", n)
 	}
-	if _, n := buildBalancerTopology(topoInfo, &Config{DiskType: ""}); n != 2 {
+	if _, n, _ := buildBalancerTopology(topoInfo, &Config{DiskType: ""}); n != 2 {
 		t.Errorf("disk_type=empty matched %d nodes, want 2 (all)", n)
 	}
-	if _, n := buildBalancerTopology(topoInfo, &Config{DiskType: "ssd"}); n != 0 {
+	if _, n, _ := buildBalancerTopology(topoInfo, &Config{DiskType: "ssd"}); n != 0 {
 		t.Errorf("disk_type=ssd matched %d nodes on an all-HDD cluster, want 0", n)
 	}
 }