milvus/internal/querynodev2/delegator/distribution_test.go

694 lines
14 KiB
Go
Raw Normal View History

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package delegator
import (
"testing"
"time"
"github.com/stretchr/testify/suite"
)
type DistributionSuite struct {
suite.Suite
dist *distribution
}
func (s *DistributionSuite) SetupTest() {
s.dist = NewDistribution()
s.Equal(initialTargetVersion, s.dist.getTargetVersion())
}
func (s *DistributionSuite) TearDownTest() {
s.dist = nil
}
func (s *DistributionSuite) TestAddDistribution() {
type testCase struct {
tag string
input []SegmentEntry
growing []SegmentEntry
expected []SnapshotItem
expectedSignalClosed bool
}
cases := []testCase{
{
tag: "one_node",
input: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
},
{
NodeID: 1,
SegmentID: 2,
},
},
expected: []SnapshotItem{
{
NodeID: 1,
Segments: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
TargetVersion: unreadableTargetVersion,
},
{
NodeID: 1,
SegmentID: 2,
TargetVersion: unreadableTargetVersion,
},
},
},
},
expectedSignalClosed: true,
},
{
tag: "duplicate segment",
input: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
},
{
NodeID: 1,
SegmentID: 1,
},
},
expected: []SnapshotItem{
{
NodeID: 1,
Segments: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
TargetVersion: unreadableTargetVersion,
},
},
},
},
expectedSignalClosed: true,
},
{
tag: "multiple_nodes",
input: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
},
{
NodeID: 2,
SegmentID: 2,
},
{
NodeID: 1,
SegmentID: 3,
},
},
expected: []SnapshotItem{
{
NodeID: 1,
Segments: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
TargetVersion: unreadableTargetVersion,
},
{
NodeID: 1,
SegmentID: 3,
TargetVersion: unreadableTargetVersion,
},
},
},
{
NodeID: 2,
Segments: []SegmentEntry{
{
NodeID: 2,
SegmentID: 2,
TargetVersion: unreadableTargetVersion,
},
},
},
},
expectedSignalClosed: true,
},
{
tag: "remove_growing",
growing: []SegmentEntry{
{NodeID: 1, SegmentID: 1},
},
input: []SegmentEntry{
{NodeID: 1, SegmentID: 1},
{NodeID: 1, SegmentID: 2},
},
expected: []SnapshotItem{
{
NodeID: 1,
Segments: []SegmentEntry{
{NodeID: 1, SegmentID: 1, TargetVersion: unreadableTargetVersion},
{NodeID: 1, SegmentID: 2, TargetVersion: unreadableTargetVersion},
},
},
},
expectedSignalClosed: false,
},
}
for _, tc := range cases {
s.Run(tc.tag, func() {
s.SetupTest()
defer s.TearDownTest()
s.dist.AddGrowing(tc.growing...)
_, _, version := s.dist.GetSegments(false)
s.dist.AddDistributions(tc.input...)
sealed, _ := s.dist.PeekSegments(false)
s.compareSnapshotItems(tc.expected, sealed)
s.dist.FinishUsage(version)
})
}
}
func (s *DistributionSuite) isClosedCh(ch chan struct{}) bool {
select {
case <-ch:
return true
default:
return false
}
}
func (s *DistributionSuite) compareSnapshotItems(target, value []SnapshotItem) {
if !s.Equal(len(target), len(value)) {
return
}
mapNodeItem := make(map[int64]SnapshotItem)
for _, valueItem := range value {
mapNodeItem[valueItem.NodeID] = valueItem
}
for _, targetItem := range target {
valueItem, ok := mapNodeItem[targetItem.NodeID]
if !s.True(ok) {
return
}
s.ElementsMatch(targetItem.Segments, valueItem.Segments)
}
}
func (s *DistributionSuite) TestAddGrowing() {
type testCase struct {
tag string
input []SegmentEntry
expected []SegmentEntry
}
cases := []testCase{
{
tag: "nil input",
input: nil,
expected: []SegmentEntry{},
},
{
tag: "normal case",
input: []SegmentEntry{
{SegmentID: 1, PartitionID: 1},
{SegmentID: 2, PartitionID: 2},
},
expected: []SegmentEntry{
{SegmentID: 1, PartitionID: 1},
{SegmentID: 2, PartitionID: 2},
},
},
}
for _, tc := range cases {
s.Run(tc.tag, func() {
s.SetupTest()
defer s.TearDownTest()
s.dist.AddGrowing(tc.input...)
_, growing, version := s.dist.GetSegments(false)
defer s.dist.FinishUsage(version)
s.ElementsMatch(tc.expected, growing)
})
}
}
func (s *DistributionSuite) TestRemoveDistribution() {
type testCase struct {
tag string
presetSealed []SegmentEntry
presetGrowing []SegmentEntry
removalSealed []SegmentEntry
removalGrowing []SegmentEntry
withMockRead bool
expectSealed []SnapshotItem
expectGrowing []SegmentEntry
}
cases := []testCase{
{
tag: "remove with no read",
presetSealed: []SegmentEntry{
{NodeID: 1, SegmentID: 1},
{NodeID: 2, SegmentID: 2},
{NodeID: 1, SegmentID: 3},
},
presetGrowing: []SegmentEntry{
{SegmentID: 4},
{SegmentID: 5},
},
removalSealed: []SegmentEntry{
{NodeID: 1, SegmentID: 1},
},
removalGrowing: []SegmentEntry{
{SegmentID: 5},
},
withMockRead: false,
expectSealed: []SnapshotItem{
{
NodeID: 1,
Segments: []SegmentEntry{
{NodeID: 1, SegmentID: 3, TargetVersion: unreadableTargetVersion},
},
},
{
NodeID: 2,
Segments: []SegmentEntry{
{NodeID: 2, SegmentID: 2, TargetVersion: unreadableTargetVersion},
},
},
},
expectGrowing: []SegmentEntry{{SegmentID: 4}},
},
{
tag: "remove with wrong nodeID",
presetSealed: []SegmentEntry{
{NodeID: 1, SegmentID: 1},
{NodeID: 2, SegmentID: 2},
{NodeID: 1, SegmentID: 3},
},
presetGrowing: []SegmentEntry{
{SegmentID: 4},
{SegmentID: 5},
},
removalSealed: []SegmentEntry{
{NodeID: 2, SegmentID: 1},
},
removalGrowing: nil,
withMockRead: false,
expectSealed: []SnapshotItem{
{
NodeID: 1,
Segments: []SegmentEntry{
{NodeID: 1, SegmentID: 1, TargetVersion: unreadableTargetVersion},
{NodeID: 1, SegmentID: 3, TargetVersion: unreadableTargetVersion},
},
},
{
NodeID: 2,
Segments: []SegmentEntry{
{NodeID: 2, SegmentID: 2, TargetVersion: unreadableTargetVersion},
},
},
},
expectGrowing: []SegmentEntry{{SegmentID: 4}, {SegmentID: 5}},
},
{
tag: "remove with wildcardNodeID",
presetSealed: []SegmentEntry{
{NodeID: 1, SegmentID: 1},
{NodeID: 2, SegmentID: 2},
{NodeID: 1, SegmentID: 3},
},
presetGrowing: []SegmentEntry{
{SegmentID: 4},
{SegmentID: 5},
},
removalSealed: []SegmentEntry{
{NodeID: wildcardNodeID, SegmentID: 1},
},
removalGrowing: nil,
withMockRead: false,
expectSealed: []SnapshotItem{
{
NodeID: 1,
Segments: []SegmentEntry{
{NodeID: 1, SegmentID: 3, TargetVersion: unreadableTargetVersion},
},
},
{
NodeID: 2,
Segments: []SegmentEntry{
{NodeID: 2, SegmentID: 2, TargetVersion: unreadableTargetVersion},
},
},
},
expectGrowing: []SegmentEntry{{SegmentID: 4}, {SegmentID: 5}},
},
{
tag: "remove with read",
presetSealed: []SegmentEntry{
{NodeID: 1, SegmentID: 1},
{NodeID: 2, SegmentID: 2},
{NodeID: 1, SegmentID: 3},
},
presetGrowing: []SegmentEntry{
{SegmentID: 4},
{SegmentID: 5},
},
removalSealed: []SegmentEntry{
{NodeID: 1, SegmentID: 1},
},
removalGrowing: []SegmentEntry{
{SegmentID: 5},
},
withMockRead: true,
expectSealed: []SnapshotItem{
{
NodeID: 1,
Segments: []SegmentEntry{
{
NodeID: 1,
SegmentID: 3,
TargetVersion: unreadableTargetVersion,
},
},
},
{
NodeID: 2,
Segments: []SegmentEntry{
{
NodeID: 2,
SegmentID: 2,
TargetVersion: unreadableTargetVersion,
},
},
},
},
expectGrowing: []SegmentEntry{{SegmentID: 4}},
},
}
for _, tc := range cases {
s.Run(tc.tag, func() {
s.SetupTest()
defer s.TearDownTest()
s.dist.AddGrowing(tc.presetGrowing...)
s.dist.AddDistributions(tc.presetSealed...)
var version int64
if tc.withMockRead {
_, _, version = s.dist.GetSegments(false)
}
ch := s.dist.RemoveDistributions(tc.removalSealed, tc.removalGrowing)
if tc.withMockRead {
// check ch not closed
select {
case <-ch:
s.Fail("ch closed with running read")
default:
}
s.dist.FinishUsage(version)
}
// check ch close very soon
timeout := time.NewTimer(time.Second)
defer timeout.Stop()
select {
case <-timeout.C:
s.Fail("ch not closed after 1 second")
case <-ch:
}
sealed, growing, version := s.dist.GetSegments(false)
defer s.dist.FinishUsage(version)
s.compareSnapshotItems(tc.expectSealed, sealed)
s.ElementsMatch(tc.expectGrowing, growing)
})
}
}
func (s *DistributionSuite) TestPeek() {
type testCase struct {
tag string
input []SegmentEntry
expected []SnapshotItem
}
cases := []testCase{
{
tag: "one node",
input: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
},
{
NodeID: 1,
SegmentID: 2,
},
},
expected: []SnapshotItem{
{
NodeID: 1,
Segments: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
TargetVersion: unreadableTargetVersion,
},
{
NodeID: 1,
SegmentID: 2,
TargetVersion: unreadableTargetVersion,
},
},
},
},
},
{
tag: "multiple nodes",
input: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
},
{
NodeID: 2,
SegmentID: 2,
},
{
NodeID: 1,
SegmentID: 3,
},
},
expected: []SnapshotItem{
{
NodeID: 1,
Segments: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
TargetVersion: unreadableTargetVersion,
},
{
NodeID: 1,
SegmentID: 3,
TargetVersion: unreadableTargetVersion,
},
},
},
{
NodeID: 2,
Segments: []SegmentEntry{
{
NodeID: 2,
SegmentID: 2,
TargetVersion: unreadableTargetVersion,
},
},
},
},
},
}
for _, tc := range cases {
s.Run(tc.tag, func() {
s.SetupTest()
defer s.TearDownTest()
// peek during lock
s.dist.AddDistributions(tc.input...)
s.dist.mut.Lock()
sealed, _ := s.dist.PeekSegments(false)
s.compareSnapshotItems(tc.expected, sealed)
s.dist.mut.Unlock()
})
}
}
func (s *DistributionSuite) TestAddOfflines() {
type testCase struct {
tag string
input []SegmentEntry
offlines []int64
serviceable bool
}
cases := []testCase{
{
tag: "offlineHits",
input: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
},
{
NodeID: 1,
SegmentID: 2,
},
},
offlines: []int64{2},
serviceable: false,
},
{
tag: "offlineMissed",
input: []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
},
{
NodeID: 2,
SegmentID: 2,
},
{
NodeID: 1,
SegmentID: 3,
},
},
serviceable: true,
},
}
for _, tc := range cases {
s.Run(tc.tag, func() {
s.SetupTest()
defer s.TearDownTest()
s.dist.AddDistributions(tc.input...)
s.dist.AddOfflines(tc.offlines...)
s.Equal(tc.serviceable, s.dist.Serviceable())
})
}
}
func (s *DistributionSuite) Test_SyncTargetVersion() {
growing := []SegmentEntry{
{
NodeID: 1,
SegmentID: 1,
PartitionID: 1,
TargetVersion: 1,
},
{
NodeID: 1,
SegmentID: 2,
PartitionID: 1,
TargetVersion: 1,
},
{
NodeID: 1,
SegmentID: 3,
PartitionID: 1,
TargetVersion: 1,
},
}
sealed := []SegmentEntry{
{
NodeID: 1,
SegmentID: 4,
PartitionID: 1,
TargetVersion: 1,
},
{
NodeID: 1,
SegmentID: 5,
PartitionID: 1,
TargetVersion: 1,
},
{
NodeID: 1,
SegmentID: 6,
PartitionID: 1,
TargetVersion: 1,
},
}
s.dist.AddGrowing(growing...)
s.dist.AddDistributions(sealed...)
s.dist.SyncTargetVersion(2, []int64{2, 3}, []int64{6}, []int64{})
s1, s2, _ := s.dist.GetSegments(true)
s.Len(s1[0].Segments, 1)
s.Len(s2, 2)
s1, s2, _ = s.dist.GetSegments(false)
s.Len(s1[0].Segments, 3)
s.Len(s2, 3)
s.dist.serviceable.Store(true)
s.dist.SyncTargetVersion(2, []int64{222}, []int64{}, []int64{})
s.True(s.dist.Serviceable())
s.dist.SyncTargetVersion(2, []int64{}, []int64{333}, []int64{})
s.False(s.dist.Serviceable())
s.dist.SyncTargetVersion(2, []int64{}, []int64{333}, []int64{1, 2, 3})
_, segments, _ := s.dist.GetSegments(true)
s.Len(segments, 0)
}
func TestDistributionSuite(t *testing.T) {
suite.Run(t, new(DistributionSuite))
}