mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-07 01:28:27 +08:00
258e1ccd66
This PR refines scheduler lifetime control: - Move private tri-state into lifetime package - Make scheduler block incoming "Add" task - Make scheduler Stop wait until all previously accepted task done Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
101 lines
2.4 KiB
Go
101 lines
2.4 KiB
Go
package tasks
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math/rand"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"go.uber.org/atomic"
|
|
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
|
)
|
|
|
|
func TestScheduler(t *testing.T) {
|
|
paramtable.Init()
|
|
t.Run("user-task-polling", func(t *testing.T) {
|
|
testScheduler(t, newUserTaskPollingPolicy())
|
|
})
|
|
t.Run("fifo", func(t *testing.T) {
|
|
testScheduler(t, newFIFOPolicy())
|
|
})
|
|
t.Run("scheduler_not_working", func(t *testing.T) {
|
|
scheduler := newScheduler(newFIFOPolicy())
|
|
|
|
task := newMockTask(mockTaskConfig{
|
|
nq: 1,
|
|
executeCost: 10 * time.Millisecond,
|
|
execution: func(ctx context.Context) error {
|
|
return nil
|
|
},
|
|
})
|
|
|
|
err := scheduler.Add(task)
|
|
assert.Error(t, err)
|
|
|
|
scheduler.Stop()
|
|
|
|
err = scheduler.Add(task)
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
func testScheduler(t *testing.T, policy schedulePolicy) {
|
|
// start a new scheduler
|
|
scheduler := newScheduler(policy)
|
|
scheduler.Start()
|
|
|
|
var cnt atomic.Int32
|
|
n := 100
|
|
nq := 0
|
|
userN := 10
|
|
// Test Push
|
|
for i := 1; i <= n; i++ {
|
|
username := fmt.Sprintf("user_%d", rand.Int31n(int32(userN)))
|
|
task := newMockTask(mockTaskConfig{
|
|
username: username,
|
|
nq: int64(i),
|
|
executeCost: 10 * time.Millisecond,
|
|
execution: func(ctx context.Context) error {
|
|
cnt.Inc()
|
|
return nil
|
|
},
|
|
})
|
|
nq += i
|
|
assert.NoError(t, scheduler.Add(task))
|
|
total := int(scheduler.GetWaitingTaskTotal())
|
|
nqNow := int(scheduler.GetWaitingTaskTotalNQ())
|
|
assert.LessOrEqual(t, total, i)
|
|
assert.LessOrEqual(t, nqNow, nq)
|
|
}
|
|
time.Sleep(2 * time.Second)
|
|
assert.Equal(t, cnt.Load(), int32(n))
|
|
assert.Equal(t, 0, int(scheduler.GetWaitingTaskTotal()))
|
|
assert.Equal(t, 0, int(scheduler.GetWaitingTaskTotalNQ()))
|
|
|
|
// Test Push
|
|
for i := 1; i <= n; i++ {
|
|
username := fmt.Sprintf("user_%d", rand.Int31n(int32(userN)))
|
|
task := newMockTask(mockTaskConfig{
|
|
username: username,
|
|
executeCost: 10 * time.Millisecond,
|
|
execution: func(ctx context.Context) error {
|
|
cnt.Inc()
|
|
return nil
|
|
},
|
|
})
|
|
assert.NoError(t, scheduler.Add(task))
|
|
total := int(scheduler.GetWaitingTaskTotal())
|
|
nqNow := int(scheduler.GetWaitingTaskTotalNQ())
|
|
assert.LessOrEqual(t, total, i)
|
|
assert.LessOrEqual(t, nqNow, i)
|
|
}
|
|
|
|
time.Sleep(2 * time.Second)
|
|
assert.Equal(t, cnt.Load(), int32(2*n))
|
|
assert.Equal(t, 0, int(scheduler.GetWaitingTaskTotal()))
|
|
assert.Equal(t, 0, int(scheduler.GetWaitingTaskTotalNQ()))
|
|
}
|