milvus/internal/proxy/task_scheduler_test.go

609 lines
15 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 proxy
import (
"context"
"fmt"
"math/rand"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
"github.com/milvus-io/milvus/pkg/mq/msgstream"
"github.com/milvus-io/milvus/pkg/util/funcutil"
)
func TestBaseTaskQueue(t *testing.T) {
var err error
var unissuedTask task
var activeTask task
tsoAllocatorIns := newMockTsoAllocator()
queue := newBaseTaskQueue(tsoAllocatorIns)
assert.NotNil(t, queue)
assert.True(t, queue.utEmpty())
assert.False(t, queue.utFull())
st := newDefaultMockTask()
stID := st.ID()
// no task in queue
unissuedTask = queue.FrontUnissuedTask()
assert.Nil(t, unissuedTask)
unissuedTask = queue.getTaskByReqID(stID)
assert.Nil(t, unissuedTask)
unissuedTask = queue.PopUnissuedTask()
assert.Nil(t, unissuedTask)
// task enqueue, only one task in queue
err = queue.Enqueue(st)
assert.NoError(t, err)
assert.False(t, queue.utEmpty())
assert.False(t, queue.utFull())
assert.Equal(t, 1, queue.unissuedTasks.Len())
assert.Equal(t, 1, len(queue.utChan()))
unissuedTask = queue.FrontUnissuedTask()
assert.NotNil(t, unissuedTask)
unissuedTask = queue.getTaskByReqID(unissuedTask.ID())
assert.NotNil(t, unissuedTask)
unissuedTask = queue.PopUnissuedTask()
assert.NotNil(t, unissuedTask)
assert.True(t, queue.utEmpty())
assert.False(t, queue.utFull())
// test active list, no task in queue
activeTask = queue.getTaskByReqID(unissuedTask.ID())
assert.Nil(t, activeTask)
activeTask = queue.PopActiveTask(unissuedTask.ID())
assert.Nil(t, activeTask)
// test active list, no task in unissued list, only one task in active list
queue.AddActiveTask(unissuedTask)
activeTask = queue.getTaskByReqID(unissuedTask.ID())
assert.NotNil(t, activeTask)
activeTask = queue.PopActiveTask(unissuedTask.ID())
assert.NotNil(t, activeTask)
// test utFull
queue.setMaxTaskNum(10) // not accurate, full also means utBufChan block
for i := 0; i < int(queue.getMaxTaskNum()); i++ {
err = queue.Enqueue(newDefaultMockTask())
assert.NoError(t, err)
}
assert.True(t, queue.utFull())
err = queue.Enqueue(newDefaultMockTask())
assert.Error(t, err)
}
func TestDdTaskQueue(t *testing.T) {
var err error
var unissuedTask task
var activeTask task
tsoAllocatorIns := newMockTsoAllocator()
queue := newDdTaskQueue(tsoAllocatorIns)
assert.NotNil(t, queue)
assert.True(t, queue.utEmpty())
assert.False(t, queue.utFull())
st := newDefaultMockDdlTask()
stID := st.ID()
// no task in queue
unissuedTask = queue.FrontUnissuedTask()
assert.Nil(t, unissuedTask)
unissuedTask = queue.getTaskByReqID(stID)
assert.Nil(t, unissuedTask)
unissuedTask = queue.PopUnissuedTask()
assert.Nil(t, unissuedTask)
// task enqueue, only one task in queue
err = queue.Enqueue(st)
assert.NoError(t, err)
assert.False(t, queue.utEmpty())
assert.False(t, queue.utFull())
assert.Equal(t, 1, queue.unissuedTasks.Len())
assert.Equal(t, 1, len(queue.utChan()))
unissuedTask = queue.FrontUnissuedTask()
assert.NotNil(t, unissuedTask)
unissuedTask = queue.getTaskByReqID(unissuedTask.ID())
assert.NotNil(t, unissuedTask)
unissuedTask = queue.PopUnissuedTask()
assert.NotNil(t, unissuedTask)
assert.True(t, queue.utEmpty())
assert.False(t, queue.utFull())
// test active list, no task in queue
activeTask = queue.getTaskByReqID(unissuedTask.ID())
assert.Nil(t, activeTask)
activeTask = queue.PopActiveTask(unissuedTask.ID())
assert.Nil(t, activeTask)
// test active list, no task in unissued list, only one task in active list
queue.AddActiveTask(unissuedTask)
activeTask = queue.getTaskByReqID(unissuedTask.ID())
assert.NotNil(t, activeTask)
activeTask = queue.PopActiveTask(unissuedTask.ID())
assert.NotNil(t, activeTask)
// test utFull
queue.setMaxTaskNum(10) // not accurate, full also means utBufChan block
for i := 0; i < int(queue.getMaxTaskNum()); i++ {
err = queue.Enqueue(newDefaultMockDdlTask())
assert.NoError(t, err)
}
assert.True(t, queue.utFull())
err = queue.Enqueue(newDefaultMockDdlTask())
assert.Error(t, err)
}
// test the logic of queue
func TestDmTaskQueue_Basic(t *testing.T) {
var err error
var unissuedTask task
var activeTask task
tsoAllocatorIns := newMockTsoAllocator()
queue := newDmTaskQueue(tsoAllocatorIns)
assert.NotNil(t, queue)
assert.True(t, queue.utEmpty())
assert.False(t, queue.utFull())
st := newDefaultMockDmlTask()
stID := st.ID()
// no task in queue
unissuedTask = queue.FrontUnissuedTask()
assert.Nil(t, unissuedTask)
unissuedTask = queue.getTaskByReqID(stID)
assert.Nil(t, unissuedTask)
unissuedTask = queue.PopUnissuedTask()
assert.Nil(t, unissuedTask)
// task enqueue, only one task in queue
err = queue.Enqueue(st)
assert.NoError(t, err)
assert.False(t, queue.utEmpty())
assert.False(t, queue.utFull())
assert.Equal(t, 1, queue.unissuedTasks.Len())
assert.Equal(t, 1, len(queue.utChan()))
unissuedTask = queue.FrontUnissuedTask()
assert.NotNil(t, unissuedTask)
unissuedTask = queue.getTaskByReqID(unissuedTask.ID())
assert.NotNil(t, unissuedTask)
unissuedTask = queue.PopUnissuedTask()
assert.NotNil(t, unissuedTask)
assert.True(t, queue.utEmpty())
assert.False(t, queue.utFull())
// test active list, no task in queue
activeTask = queue.getTaskByReqID(unissuedTask.ID())
assert.Nil(t, activeTask)
activeTask = queue.PopActiveTask(unissuedTask.ID())
assert.Nil(t, activeTask)
// test active list, no task in unissued list, only one task in active list
queue.AddActiveTask(unissuedTask)
activeTask = queue.getTaskByReqID(unissuedTask.ID())
assert.NotNil(t, activeTask)
activeTask = queue.PopActiveTask(unissuedTask.ID())
assert.NotNil(t, activeTask)
// test utFull
queue.setMaxTaskNum(10) // not accurate, full also means utBufChan block
for i := 0; i < int(queue.getMaxTaskNum()); i++ {
err = queue.Enqueue(newDefaultMockDmlTask())
assert.NoError(t, err)
}
assert.True(t, queue.utFull())
err = queue.Enqueue(newDefaultMockDmlTask())
assert.Error(t, err)
}
// test the timestamp statistics
func TestDmTaskQueue_TimestampStatistics(t *testing.T) {
var err error
var unissuedTask task
tsoAllocatorIns := newMockTsoAllocator()
queue := newDmTaskQueue(tsoAllocatorIns)
assert.NotNil(t, queue)
st := newDefaultMockDmlTask()
stPChans := st.pchans
err = queue.Enqueue(st)
assert.NoError(t, err)
stats, err := queue.getPChanStatsInfo()
assert.NoError(t, err)
assert.Equal(t, len(stPChans), len(stats))
unissuedTask = queue.FrontUnissuedTask()
assert.NotNil(t, unissuedTask)
for _, stat := range stats {
assert.Equal(t, unissuedTask.BeginTs(), stat.minTs)
assert.Equal(t, unissuedTask.EndTs(), stat.maxTs)
}
unissuedTask = queue.PopUnissuedTask()
assert.NotNil(t, unissuedTask)
assert.True(t, queue.utEmpty())
queue.AddActiveTask(unissuedTask)
queue.PopActiveTask(unissuedTask.ID())
stats, err = queue.getPChanStatsInfo()
assert.NoError(t, err)
assert.Zero(t, len(stats))
}
// test the timestamp statistics
func TestDmTaskQueue_TimestampStatistics2(t *testing.T) {
tsoAllocatorIns := newMockTsoAllocator()
queue := newDmTaskQueue(tsoAllocatorIns)
assert.NotNil(t, queue)
prefix := funcutil.GenRandomStr()
insertNum := 100
var processWg sync.WaitGroup
processWg.Add(1)
processCtx, processCancel := context.WithCancel(context.TODO())
processCount := insertNum
var processCountMut sync.RWMutex
go func() {
defer processWg.Done()
var workerWg sync.WaitGroup
workerWg.Add(insertNum)
for processCtx.Err() == nil {
if queue.utEmpty() {
continue
}
utTask := queue.PopUnissuedTask()
go func(ut task) {
defer workerWg.Done()
assert.NotNil(t, ut)
queue.AddActiveTask(ut)
dur := time.Duration(50+rand.Int()%10) * time.Millisecond
time.Sleep(dur)
queue.PopActiveTask(ut.ID())
processCountMut.Lock()
defer processCountMut.Unlock()
processCount--
}(utTask)
}
workerWg.Wait()
}()
var currPChanStats map[pChan]*pChanStatistics
var wgSchedule sync.WaitGroup
scheduleCtx, scheduleCancel := context.WithCancel(context.TODO())
schedule := func() {
defer wgSchedule.Done()
ticker := time.NewTicker(time.Millisecond * 10)
defer ticker.Stop()
for {
select {
case <-scheduleCtx.Done():
return
case <-ticker.C:
stats, err := queue.getPChanStatsInfo()
assert.NoError(t, err)
if currPChanStats == nil {
currPChanStats = stats
} else {
// assure minTs and maxTs will not go back
for p, stat := range stats {
curInfo, ok := currPChanStats[p]
if ok {
fmt.Println("stat.minTs", stat.minTs, " ", "curInfo.minTs:", curInfo.minTs)
fmt.Println("stat.maxTs", stat.maxTs, " ", "curInfo.minTs:", curInfo.maxTs)
assert.True(t, stat.minTs >= curInfo.minTs)
curInfo.minTs = stat.minTs
assert.True(t, stat.maxTs >= curInfo.maxTs)
curInfo.maxTs = stat.maxTs
}
}
}
}
}
}
wgSchedule.Add(1)
go schedule()
var wg sync.WaitGroup
wg.Add(insertNum)
for i := 0; i < insertNum; i++ {
go func() {
defer wg.Done()
time.Sleep(time.Millisecond)
st := newDefaultMockDmlTask()
vChannels := make([]string, 2)
vChannels[0] = prefix + "_1"
vChannels[1] = prefix + "_2"
st.vchans = vChannels
st.pchans = vChannels
err := queue.Enqueue(st)
assert.NoError(t, err)
}()
}
wg.Wait()
// time.Sleep(time.Millisecond*100)
needLoop := true
for needLoop {
processCountMut.RLock()
needLoop = processCount != 0
processCountMut.RUnlock()
}
processCancel()
processWg.Wait()
scheduleCancel()
wgSchedule.Wait()
stats, err := queue.getPChanStatsInfo()
assert.NoError(t, err)
assert.Zero(t, len(stats))
}
func TestDqTaskQueue(t *testing.T) {
var err error
var unissuedTask task
var activeTask task
tsoAllocatorIns := newMockTsoAllocator()
queue := newDqTaskQueue(tsoAllocatorIns)
assert.NotNil(t, queue)
assert.True(t, queue.utEmpty())
assert.False(t, queue.utFull())
st := newDefaultMockDqlTask()
stID := st.ID()
// no task in queue
unissuedTask = queue.FrontUnissuedTask()
assert.Nil(t, unissuedTask)
unissuedTask = queue.getTaskByReqID(stID)
assert.Nil(t, unissuedTask)
unissuedTask = queue.PopUnissuedTask()
assert.Nil(t, unissuedTask)
// task enqueue, only one task in queue
err = queue.Enqueue(st)
assert.NoError(t, err)
assert.False(t, queue.utEmpty())
assert.False(t, queue.utFull())
assert.Equal(t, 1, queue.unissuedTasks.Len())
assert.Equal(t, 1, len(queue.utChan()))
unissuedTask = queue.FrontUnissuedTask()
assert.NotNil(t, unissuedTask)
unissuedTask = queue.getTaskByReqID(unissuedTask.ID())
assert.NotNil(t, unissuedTask)
unissuedTask = queue.PopUnissuedTask()
assert.NotNil(t, unissuedTask)
assert.True(t, queue.utEmpty())
assert.False(t, queue.utFull())
// test active list, no task in queue
activeTask = queue.getTaskByReqID(unissuedTask.ID())
assert.Nil(t, activeTask)
activeTask = queue.PopActiveTask(unissuedTask.ID())
assert.Nil(t, activeTask)
// test active list, no task in unissued list, only one task in active list
queue.AddActiveTask(unissuedTask)
activeTask = queue.getTaskByReqID(unissuedTask.ID())
assert.NotNil(t, activeTask)
activeTask = queue.PopActiveTask(unissuedTask.ID())
assert.NotNil(t, activeTask)
// test utFull
queue.setMaxTaskNum(10) // not accurate, full also means utBufChan block
for i := 0; i < int(queue.getMaxTaskNum()); i++ {
err = queue.Enqueue(newDefaultMockDqlTask())
assert.NoError(t, err)
}
assert.True(t, queue.utFull())
err = queue.Enqueue(newDefaultMockDqlTask())
assert.Error(t, err)
}
func TestTaskScheduler(t *testing.T) {
var err error
ctx := context.Background()
tsoAllocatorIns := newMockTsoAllocator()
factory := newSimpleMockMsgStreamFactory()
sched, err := newTaskScheduler(ctx, tsoAllocatorIns, factory)
assert.NoError(t, err)
assert.NotNil(t, sched)
err = sched.Start()
assert.NoError(t, err)
defer sched.Close()
stats, err := sched.getPChanStatistics()
assert.NoError(t, err)
assert.Equal(t, 0, len(stats))
ddNum := rand.Int() % 10
dmNum := rand.Int() % 10
dqNum := rand.Int() % 10
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < ddNum; i++ {
wg.Add(1)
go func() {
defer wg.Done()
err := sched.ddQueue.Enqueue(newDefaultMockDdlTask())
assert.NoError(t, err)
}()
}
}()
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < dmNum; i++ {
wg.Add(1)
go func() {
defer wg.Done()
err := sched.dmQueue.Enqueue(newDefaultMockDmlTask())
assert.NoError(t, err)
}()
}
}()
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < dqNum; i++ {
wg.Add(1)
go func() {
defer wg.Done()
err := sched.dqQueue.Enqueue(newDefaultMockDqlTask())
assert.NoError(t, err)
}()
}
}()
wg.Wait()
}
func TestTaskScheduler_concurrentPushAndPop(t *testing.T) {
collectionID := UniqueID(0)
collectionName := "col-0"
channels := []pChan{"mock-chan-0", "mock-chan-1"}
cache := NewMockCache(t)
cache.On("GetCollectionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(collectionID, nil)
globalMetaCache = cache
tsoAllocatorIns := newMockTsoAllocator()
factory := newSimpleMockMsgStreamFactory()
scheduler, err := newTaskScheduler(context.Background(), tsoAllocatorIns, factory)
assert.NoError(t, err)
run := func(wg *sync.WaitGroup) {
defer wg.Done()
chMgr := newMockChannelsMgr()
chMgr.getChannelsFunc = func(collectionID UniqueID) ([]pChan, error) {
return channels, nil
}
it := &insertTask{
ctx: context.Background(),
insertMsg: &msgstream.InsertMsg{
InsertRequest: msgpb.InsertRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
},
},
chMgr: chMgr,
}
err := scheduler.dmQueue.Enqueue(it)
assert.NoError(t, err)
task := scheduler.scheduleDmTask()
scheduler.dmQueue.AddActiveTask(task)
chMgr.getChannelsFunc = func(collectionID UniqueID) ([]pChan, error) {
return nil, fmt.Errorf("mock err")
}
scheduler.dmQueue.PopActiveTask(task.ID()) // assert no panic
}
wg := &sync.WaitGroup{}
for i := 0; i < 100; i++ {
wg.Add(1)
go run(wg)
}
wg.Wait()
}