diff --git a/weed/storage/store_disk_space_test.go b/weed/storage/store_disk_space_test.go index 13ed30dae..884b8dda1 100644 --- a/weed/storage/store_disk_space_test.go +++ b/weed/storage/store_disk_space_test.go @@ -95,20 +95,25 @@ func TestHasFreeDiskLocation(t *testing.T) { } } +func newTestLocation(maxCount int32, isDiskLow bool, volCount int) *DiskLocation { + location := &DiskLocation{ + volumes: make(map[needle.VolumeId]*Volume), + ecVolumes: make(map[needle.VolumeId]*erasure_coding.EcVolume), + MaxVolumeCount: maxCount, + DiskType: types.ToDiskType("hdd"), + isDiskSpaceLow: isDiskLow, + } + for i := 1; i <= volCount; i++ { + location.volumes[needle.VolumeId(i)] = &Volume{} + } + return location +} + func TestCollectHeartbeatRespectsLowDiskSpace(t *testing.T) { diskType := types.ToDiskType("hdd") t.Run("low disk space", func(t *testing.T) { - location := &DiskLocation{ - volumes: make(map[needle.VolumeId]*Volume), - ecVolumes: make(map[needle.VolumeId]*erasure_coding.EcVolume), - MaxVolumeCount: 10, - DiskType: diskType, - isDiskSpaceLow: true, - } - for i := 1; i <= 3; i++ { - location.volumes[needle.VolumeId(i)] = &Volume{} - } + location := newTestLocation(10, true, 3) store := &Store{Locations: []*DiskLocation{location}} hb := store.CollectHeartbeat() @@ -118,16 +123,7 @@ func TestCollectHeartbeatRespectsLowDiskSpace(t *testing.T) { }) t.Run("normal disk space", func(t *testing.T) { - location := &DiskLocation{ - volumes: make(map[needle.VolumeId]*Volume), - ecVolumes: make(map[needle.VolumeId]*erasure_coding.EcVolume), - MaxVolumeCount: 10, - DiskType: diskType, - isDiskSpaceLow: false, - } - for i := 1; i <= 3; i++ { - location.volumes[needle.VolumeId(i)] = &Volume{} - } + location := newTestLocation(10, false, 3) store := &Store{Locations: []*DiskLocation{location}} hb := store.CollectHeartbeat() @@ -136,17 +132,18 @@ func TestCollectHeartbeatRespectsLowDiskSpace(t *testing.T) { } }) - t.Run("low disk space with ec shards", func(t *testing.T) { - location := &DiskLocation{ - volumes: make(map[needle.VolumeId]*Volume), - ecVolumes: make(map[needle.VolumeId]*erasure_coding.EcVolume), - MaxVolumeCount: 10, - DiskType: diskType, - isDiskSpaceLow: true, - } - for i := 1; i <= 3; i++ { - location.volumes[needle.VolumeId(i)] = &Volume{} + t.Run("low disk space zero volumes", func(t *testing.T) { + location := newTestLocation(10, true, 0) + store := &Store{Locations: []*DiskLocation{location}} + + hb := store.CollectHeartbeat() + if got := hb.MaxVolumeCounts[string(diskType)]; got != 0 { + t.Errorf("expected zero volumes to report zero capacity, got %d", got) } + }) + + t.Run("low disk space with ec shards", func(t *testing.T) { + location := newTestLocation(10, true, 3) ecVolume := &erasure_coding.EcVolume{VolumeId: 1} const shardCount = 15 @@ -164,4 +161,34 @@ func TestCollectHeartbeatRespectsLowDiskSpace(t *testing.T) { t.Errorf("expected low disk space to include ec shard contribution, got %d want %d", got, expectedSlots) } }) + + t.Run("low disk space with multiple ec volumes", func(t *testing.T) { + location := newTestLocation(10, true, 2) + + totalShardCount := 0 + + addEcVolume := func(vid needle.VolumeId, shardCount int) { + ecVolume := &erasure_coding.EcVolume{VolumeId: vid} + for i := 0; i < shardCount; i++ { + ecVolume.Shards = append(ecVolume.Shards, &erasure_coding.EcVolumeShard{ + ShardId: erasure_coding.ShardId(i), + }) + } + location.ecVolumes[vid] = ecVolume + totalShardCount += shardCount + } + + addEcVolume(1, 12) + addEcVolume(2, 6) + + store := &Store{Locations: []*DiskLocation{location}} + + hb := store.CollectHeartbeat() + expectedSlots := len(location.volumes) + expectedSlots += (totalShardCount + erasure_coding.DataShardsCount - 1) / erasure_coding.DataShardsCount + + if got := hb.MaxVolumeCounts[string(diskType)]; got != uint32(expectedSlots) { + t.Errorf("expected multiple ec volumes to be counted, got %d want %d", got, expectedSlots) + } + }) }