seaweedfs/test/erasure_coding/multidisk_shardloss_test.go

package erasure_coding

import (
	"context"
	"fmt"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/seaweedfs/seaweedfs/weed/shell"
	"github.com/seaweedfs/seaweedfs/weed/storage/needle"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"google.golang.org/grpc"
)

// TestMultiDiskECBalanceNoShardLoss is the end-to-end regression for issue 9593.
// It runs a real cluster of multi-disk volume servers (3 servers x 4 disks),
// EC-encodes a volume, then runs ec.balance, asserting hard invariants the older
// integration tests only logged:
//
//   - after encode the full set of 14 EC shards exists,
//   - ec.balance never loses a shard (still 14 distinct shards afterwards),
//   - shards end up spread across more than one disk per node, and
//   - cluster.status counts physical disks (not one per node) and matches the
//     real on-disk distribution.
func TestMultiDiskECBalanceNoShardLoss(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping multi-disk EC integration test in short mode")
	}

	testDir := t.TempDir()
	ctx, cancel := context.WithTimeout(context.Background(), 240*time.Second)
	defer cancel()

	cluster, err := startMultiDiskCluster(ctx, testDir)
	require.NoError(t, err)
	defer cluster.Stop()

	require.NoError(t, waitForServer("127.0.0.1:9334", 30*time.Second))
	for i := 0; i < 3; i++ {
		require.NoError(t, waitForServer(fmt.Sprintf("127.0.0.1:809%d", i), 30*time.Second))
	}
	t.Log("waiting for multi-disk volume servers to register...")
	time.Sleep(10 * time.Second)

	commandEnv := shell.NewCommandEnv(&shell.ShellOptions{
		Masters:        stringPtr("127.0.0.1:9334"),
		GrpcDialOption: grpc.WithInsecure(),
		FilerGroup:     stringPtr("default"),
	})
	connectToMasterAndSync(ctx, t, commandEnv)

	// Upload enough small files that the volume holds real data to encode.
	var volumeId needle.VolumeId
	for retry := 0; retry < 5; retry++ {
		volumeId, err = uploadTestDataToMaster([]byte(strings.Repeat("multidisk-ec-9593 ", 64)), "127.0.0.1:9334")
		if err == nil {
			break
		}
		time.Sleep(3 * time.Second)
	}
	require.NoError(t, err, "failed to upload test data")
	for i := 0; i < 40; i++ {
		if _, e := uploadTestDataToMaster([]byte(strings.Repeat("filler ", 128)), "127.0.0.1:9334"); e != nil {
			break
		}
	}
	t.Logf("using volume %d", volumeId)
	time.Sleep(3 * time.Second)

	// Populate every server's disks with volumes so the balancer can see and
	// target each physical disk. The master only enumerates disks that already
	// hold a volume or EC shard — an empty disk leaves no trace in the topology
	// (heartbeats aggregate capacity per disk type, not per physical disk). So
	// without pre-populating, the post-encode balance would collapse each node's
	// shards onto the single disk that happened to hold data, and whether the
	// fillers spread across disks is environment-dependent (master volume-growth
	// timing). Growing a few volumes per server makes the multi-disk layout
	// deterministic: the volume server places each new volume on its least-loaded
	// disk, so a handful of grows touches every disk.
	for i := 0; i < 3; i++ {
		server := fmt.Sprintf("127.0.0.1:809%d", i)
		out, growErr := captureCommandOutput(t, shell.Commands[findCommandIndex("volume.grow")],
			[]string{"-collection", "test", "-dataNode", server, "-count", "4"}, commandEnv)
		require.NoError(t, growErr, "volume.grow on %s failed: %s", server, out)
	}
	// Let the freshly grown volumes reach the master via heartbeat before encoding
	// so collectEcNodes sees every disk.
	time.Sleep(5 * time.Second)

	locked, unlock := tryLockWithTimeout(t, commandEnv, 15*time.Second)
	require.True(t, locked, "could not acquire shell lock")
	defer unlock()

	// EC-encode the volume.
	out, err := captureCommandOutput(t, shell.Commands[findCommandIndex("ec.encode")],
		[]string{"-volumeId", fmt.Sprintf("%d", volumeId), "-collection", "test", "-force"}, commandEnv)
	t.Logf("ec.encode output:\n%s", out)
	require.NoError(t, err, "ec.encode failed")

	// All 14 shards must exist after encoding.
	require.Eventually(t, func() bool {
		return len(collectDistinctShardIDs(testDir, uint32(volumeId))) == erasureShardCount
	}, 30*time.Second, time.Second, "expected all %d EC shards after encode, got %v",
		erasureShardCount, collectDistinctShardIDs(testDir, uint32(volumeId)))

	beforeBalance := collectDistinctShardIDs(testDir, uint32(volumeId))
	t.Logf("after encode: %d distinct shards on %d disks", len(beforeBalance), disksWithShards(testDir, uint32(volumeId)))

	// Run ec.balance.
	out, err = captureCommandOutput(t, shell.Commands[findCommandIndex("ec.balance")],
		[]string{"-collection", "test", "-force"}, commandEnv)
	t.Logf("ec.balance output:\n%s", out)
	require.NoError(t, err, "ec.balance failed")
	time.Sleep(3 * time.Second)

	// The core regression: ec.balance must not lose any shard.
	afterBalance := collectDistinctShardIDs(testDir, uint32(volumeId))
	require.Equal(t, erasureShardCount, len(afterBalance),
		"ec.balance lost shards on multi-disk nodes: had %v, now %v", sortedKeysOf(beforeBalance), sortedKeysOf(afterBalance))

	// Shards must be spread across more than one physical disk per node overall.
	usedDisks := disksWithShards(testDir, uint32(volumeId))
	assert.Greater(t, usedDisks, 3, "EC shards should span more than one disk per node (got %d disks across 3 nodes)", usedDisks)

	// cluster.status must count physical disks, not collapse to one per node: it
	// must report at least the disks actually holding this volume's shards (which
	// is already >3 across the 3 nodes). Before the fix it reported 3 (node count).
	require.Eventually(t, func() bool {
		n, ok := clusterStatusDiskCount(t, commandEnv)
		return ok && n >= usedDisks
	}, 30*time.Second, 2*time.Second, "cluster.status never reported the >=%d physical disks holding shards (multi-disk count)", usedDisks)

	n, _ := clusterStatusDiskCount(t, commandEnv)
	t.Logf("cluster.status reports %d physical disks (>= %d holding this volume's shards)", n, usedDisks)
}

const erasureShardCount = 14 // 10 data + 4 parity

// collectDistinctShardIDs returns the set of EC shard ids present for a volume
// across every disk of every server in the multi-disk test layout.
func collectDistinctShardIDs(testDir string, volumeId uint32) map[int]bool {
	ids := map[int]bool{}
	for server := 0; server < 3; server++ {
		for disk := 0; disk < 4; disk++ {
			diskDir := filepath.Join(testDir, fmt.Sprintf("server%d_disk%d", server, disk))
			files, err := listECShardFiles(diskDir, volumeId)
			if err != nil {
				continue
			}
			for _, f := range files {
				i := strings.LastIndex(f, ".ec")
				if i < 0 {
					continue
				}
				if n, err := strconv.Atoi(f[i+3:]); err == nil && n >= 0 && n < erasureShardCount {
					ids[n] = true
				}
			}
		}
	}
	return ids
}

// disksWithShards counts how many physical disks hold at least one shard.
func disksWithShards(testDir string, volumeId uint32) int {
	n := 0
	for _, disks := range countShardsPerDisk(testDir, volumeId) {
		for _, c := range disks {
			if c > 0 {
				n++
			}
		}
	}
	return n
}

var diskCountRe = regexp.MustCompile(`(\d+)\s+disks?`)

// clusterStatusDiskCount runs cluster.status and parses the reported disk count.
func clusterStatusDiskCount(t *testing.T, commandEnv *shell.CommandEnv) (int, bool) {
	t.Helper()
	out, err := captureCommandOutput(t, shell.Commands[findCommandIndex("cluster.status")], []string{}, commandEnv)
	if err != nil {
		return 0, false
	}
	m := diskCountRe.FindStringSubmatch(out)
	if m == nil {
		return 0, false
	}
	n, err := strconv.Atoi(m[1])
	return n, err == nil
}

func sortedKeysOf(m map[int]bool) []int {
	out := make([]int, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	for i := 1; i < len(out); i++ {
		for j := i; j > 0 && out[j-1] > out[j]; j-- {
			out[j-1], out[j] = out[j], out[j-1]
		}
	}
	return out
}