[cluster] Fix unbalanced initial shard allocation for sharded placement

src/cluster/placement/algo/sharded_helper.go

@@ -491,8 +491,8 @@ func (ph *helper) returnInitializingShardsToSource(
 
 func (ph *helper) mostUnderLoadedInstance() (placement.Instance, bool) {
 	var (
-		res        placement.Instance
-		maxLoadGap int
+		res              placement.Instance
+		maxLoadGap       int
 		totalLoadSurplus int
 	)
 
@@ -656,12 +656,19 @@ func (h *instanceHeap) Less(i, j int) bool {
 	leftLoadOnJ := h.targetLoadForInstance(instanceJ.ID()) - loadOnInstance(instanceJ)
 	// If both instance has tokens to be filled, prefer the one from bigger isolation group
 	// since it tends to be more picky in accepting shards
-	if leftLoadOnI > 0 && leftLoadOnJ > 0 {
-		if instanceI.IsolationGroup() != instanceJ.IsolationGroup() {
-			return h.igToWeightMap[instanceI.IsolationGroup()] > h.igToWeightMap[instanceJ.IsolationGroup()]
+	if leftLoadOnI > 0 && leftLoadOnJ > 0 && instanceI.IsolationGroup() != instanceJ.IsolationGroup() {
+		var (
+			igWeightI = h.igToWeightMap[instanceI.IsolationGroup()]
+			igWeightJ = h.igToWeightMap[instanceJ.IsolationGroup()]
+		)
+		if igWeightI != igWeightJ {
+			return igWeightI > igWeightJ
 		}
 	}
 	// compare left capacity on both instances
+	if leftLoadOnI == leftLoadOnJ {
+		return instanceI.ID() < instanceJ.ID()
+	}
 	if h.capacityAscending {
 		return leftLoadOnI > leftLoadOnJ
 	}

src/cluster/placement/algo/sharded_test.go

@@ -1306,6 +1306,42 @@ func TestBalanceShardsForShardedWhenImbalanced(t *testing.T) {
 	assert.Equal(t, expectedInstances, balancedPlacement.Instances())
 }
 
+func TestInitialPlacementIsBalanced(t *testing.T) {
+	instances := make([]placement.Instance, 0)
+
+	for i := 0; i < 10; i++ {
+		instances = append(instances,
+			placement.NewEmptyInstance(fmt.Sprintf("instance-0-%03d", i), "iso0", "zone", "endpoint", 1),
+			placement.NewEmptyInstance(fmt.Sprintf("instance-1-%03d", i), "iso1", "zone", "endpoint", 1),
+		)
+	}
+
+	ids := make([]uint32, 1024)
+	for i := range ids {
+		ids[i] = uint32(i)
+	}
+
+	algo := newShardedAlgorithm(placement.NewOptions())
+	p, err := algo.InitialPlacement(instances, ids, 2)
+	require.NoError(t, err)
+
+	var (
+		min = math.MaxInt32
+		max = math.MinInt32
+	)
+	for _, instance := range p.Instances() {
+		n := instance.Shards().NumShards()
+		if n < min {
+			min = n
+		}
+		if n > max {
+			max = n
+		}
+	}
+	require.True(t, max-min <= 1,
+		"The number of shards per instance differs by more than 1, min=%d max=%d", min, max)
+}
+
 func verifyAllShardsInAvailableState(t *testing.T, p placement.Placement) {
 	for _, instance := range p.Instances() {
 		s := instance.Shards()