mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 13:28:49 +08:00
26f06dd732
Signed-off-by: SimFG <bang.fu@zilliz.com>
230 lines
5.4 KiB
Go
230 lines
5.4 KiB
Go
// 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 rootcoord
|
|
|
|
import (
|
|
"context"
|
|
"math/rand"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"github.com/milvus-io/milvus/pkg/util/retry"
|
|
)
|
|
|
|
type mockChildStep struct{}
|
|
|
|
func (m *mockChildStep) Execute(ctx context.Context) ([]nestedStep, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (m *mockChildStep) Desc() string {
|
|
return "mock child step"
|
|
}
|
|
|
|
func (m *mockChildStep) Weight() stepPriority {
|
|
return stepPriorityLow
|
|
}
|
|
|
|
func newMockChildStep() *mockChildStep {
|
|
return &mockChildStep{}
|
|
}
|
|
|
|
type mockStepWithChild struct{}
|
|
|
|
func (m *mockStepWithChild) Execute(ctx context.Context) ([]nestedStep, error) {
|
|
return []nestedStep{newMockChildStep()}, nil
|
|
}
|
|
|
|
func (m *mockStepWithChild) Desc() string {
|
|
return "mock step with child"
|
|
}
|
|
|
|
func (m *mockStepWithChild) Weight() stepPriority {
|
|
return stepPriorityLow
|
|
}
|
|
|
|
func newMockStepWithChild() *mockStepWithChild {
|
|
return &mockStepWithChild{}
|
|
}
|
|
|
|
func Test_stepStack_Execute(t *testing.T) {
|
|
t.Run("normal case", func(t *testing.T) {
|
|
steps := []nestedStep{
|
|
newMockStepWithChild(),
|
|
newMockChildStep(),
|
|
}
|
|
s := &stepStack{steps: steps}
|
|
unfinished := s.Execute(context.Background())
|
|
assert.Nil(t, unfinished)
|
|
})
|
|
|
|
t.Run("error case", func(t *testing.T) {
|
|
steps := []nestedStep{
|
|
newMockNormalStep(),
|
|
newMockNormalStep(),
|
|
newMockFailStep(),
|
|
newMockNormalStep(),
|
|
}
|
|
s := &stepStack{steps: steps}
|
|
unfinished := s.Execute(context.Background())
|
|
assert.Equal(t, 3, len(unfinished.steps))
|
|
})
|
|
|
|
t.Run("Unrecoverable", func(t *testing.T) {
|
|
failStep := newMockFailStep()
|
|
failStep.err = retry.Unrecoverable(errors.New("error mock Execute"))
|
|
steps := []nestedStep{
|
|
failStep,
|
|
}
|
|
s := &stepStack{steps: steps}
|
|
unfinished := s.Execute(context.Background())
|
|
assert.Nil(t, unfinished)
|
|
})
|
|
}
|
|
|
|
func Test_randomSelect(t *testing.T) {
|
|
s0 := &stepStack{steps: []nestedStep{}}
|
|
s1 := &stepStack{steps: []nestedStep{
|
|
newMockNormalStep(),
|
|
}}
|
|
s2 := &stepStack{steps: []nestedStep{
|
|
newMockNormalStep(),
|
|
newMockNormalStep(),
|
|
}}
|
|
s3 := &stepStack{steps: []nestedStep{
|
|
newMockNormalStep(),
|
|
newMockNormalStep(),
|
|
newMockNormalStep(),
|
|
}}
|
|
s4 := &stepStack{steps: []nestedStep{
|
|
newMockNormalStep(),
|
|
newMockNormalStep(),
|
|
newMockNormalStep(),
|
|
newMockNormalStep(),
|
|
}}
|
|
m := map[*stepStack]struct{}{
|
|
s0: {},
|
|
s1: {},
|
|
s2: {},
|
|
s3: {},
|
|
s4: {},
|
|
}
|
|
selected := randomSelect(0, m)
|
|
assert.Equal(t, defaultBgExecutingParallel, len(selected))
|
|
for _, s := range selected {
|
|
_, ok := m[s]
|
|
assert.True(t, ok)
|
|
}
|
|
selected = randomSelect(2, m)
|
|
assert.Equal(t, 2, len(selected))
|
|
for _, s := range selected {
|
|
_, ok := m[s]
|
|
assert.True(t, ok)
|
|
}
|
|
}
|
|
|
|
func Test_bgStepExecutor_scheduleLoop(t *testing.T) {
|
|
bg := newBgStepExecutor(context.Background(),
|
|
withSelectStepPolicy(randomSelectPolicy(defaultBgExecutingParallel)),
|
|
withBgInterval(time.Millisecond*10))
|
|
bg.Start()
|
|
n := 20
|
|
records := make([]int, 0, n)
|
|
steps := make([]*stepStack, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
var s *stepStack
|
|
r := rand.Int() % 3
|
|
records = append(records, r)
|
|
switch r {
|
|
case 0:
|
|
s = nil
|
|
case 1:
|
|
failStep := newMockFailStep()
|
|
failStep.err = retry.Unrecoverable(errors.New("error mock Execute"))
|
|
s = &stepStack{steps: []nestedStep{
|
|
newMockNormalStep(),
|
|
failStep,
|
|
newMockNormalStep(),
|
|
}}
|
|
case 2:
|
|
s = &stepStack{steps: []nestedStep{
|
|
newMockNormalStep(),
|
|
newMockNormalStep(),
|
|
newMockNormalStep(),
|
|
}}
|
|
default:
|
|
}
|
|
steps = append(steps, s)
|
|
bg.AddSteps(s)
|
|
}
|
|
for i, r := range records {
|
|
switch r {
|
|
case 0:
|
|
assert.Nil(t, steps[i])
|
|
case 1:
|
|
<-steps[i].steps[1].(*mockFailStep).calledChan
|
|
assert.True(t, steps[i].steps[1].(*mockFailStep).called)
|
|
case 2:
|
|
default:
|
|
}
|
|
}
|
|
bg.Stop()
|
|
}
|
|
|
|
func Test_selectByPriorityPolicy(t *testing.T) {
|
|
policy := selectByPriorityPolicy(4)
|
|
|
|
t.Run("select all", func(t *testing.T) {
|
|
m := map[*stepStack]struct{}{
|
|
{steps: []nestedStep{}}: {},
|
|
{steps: []nestedStep{}}: {},
|
|
}
|
|
selected := policy(m)
|
|
assert.Equal(t, 2, len(selected))
|
|
})
|
|
|
|
t.Run("select by priority", func(t *testing.T) {
|
|
steps := []nestedStep{
|
|
&releaseCollectionStep{},
|
|
&releaseCollectionStep{},
|
|
&releaseCollectionStep{},
|
|
&releaseCollectionStep{},
|
|
&releaseCollectionStep{},
|
|
}
|
|
s1 := &stepStack{steps: steps[0:1]}
|
|
s2 := &stepStack{steps: steps[0:2]}
|
|
s3 := &stepStack{steps: steps[0:3]}
|
|
s4 := &stepStack{steps: steps[0:4]}
|
|
s5 := &stepStack{steps: steps[0:5]}
|
|
m := map[*stepStack]struct{}{
|
|
s1: {},
|
|
s2: {},
|
|
s3: {},
|
|
s4: {},
|
|
s5: {},
|
|
}
|
|
selected := policy(m)
|
|
assert.Equal(t, 4, len(selected))
|
|
for i := 1; i < len(selected); i++ {
|
|
assert.True(t, selected[i].totalPriority() <= selected[i-1].totalPriority())
|
|
}
|
|
})
|
|
}
|