mirror of https://github.com/milvus-io/milvus.git
101 lines
1.7 KiB
Go
101 lines
1.7 KiB
Go
package rootcoord
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/milvus-io/milvus/internal/tso"
|
|
|
|
"github.com/milvus-io/milvus/internal/allocator"
|
|
)
|
|
|
|
type IScheduler interface {
|
|
Start()
|
|
Stop()
|
|
AddTask(t taskV2) error
|
|
}
|
|
|
|
type scheduler struct {
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
wg sync.WaitGroup
|
|
|
|
idAllocator allocator.GIDAllocator
|
|
tsoAllocator tso.Allocator
|
|
|
|
taskChan chan taskV2
|
|
}
|
|
|
|
func newScheduler(ctx context.Context, idAllocator allocator.GIDAllocator, tsoAllocator tso.Allocator) *scheduler {
|
|
ctx1, cancel := context.WithCancel(ctx)
|
|
// TODO
|
|
n := 1024 * 10
|
|
return &scheduler{
|
|
ctx: ctx1,
|
|
cancel: cancel,
|
|
idAllocator: idAllocator,
|
|
tsoAllocator: tsoAllocator,
|
|
taskChan: make(chan taskV2, n),
|
|
}
|
|
}
|
|
|
|
func (s *scheduler) Start() {
|
|
s.wg.Add(1)
|
|
go s.taskLoop()
|
|
}
|
|
|
|
func (s *scheduler) Stop() {
|
|
s.cancel()
|
|
s.wg.Wait()
|
|
}
|
|
|
|
func (s *scheduler) taskLoop() {
|
|
defer s.wg.Done()
|
|
for {
|
|
select {
|
|
case <-s.ctx.Done():
|
|
return
|
|
case task := <-s.taskChan:
|
|
if err := task.Prepare(task.GetCtx()); err != nil {
|
|
task.NotifyDone(err)
|
|
continue
|
|
}
|
|
err := task.Execute(task.GetCtx())
|
|
task.NotifyDone(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *scheduler) setID(task taskV2) error {
|
|
id, err := s.idAllocator.AllocOne()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
task.SetID(id)
|
|
return nil
|
|
}
|
|
|
|
func (s *scheduler) setTs(task taskV2) error {
|
|
ts, err := s.tsoAllocator.GenerateTSO(1)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
task.SetTs(ts)
|
|
return nil
|
|
}
|
|
|
|
func (s *scheduler) enqueue(task taskV2) {
|
|
s.taskChan <- task
|
|
}
|
|
|
|
func (s *scheduler) AddTask(task taskV2) error {
|
|
if err := s.setID(task); err != nil {
|
|
return err
|
|
}
|
|
if err := s.setTs(task); err != nil {
|
|
return err
|
|
}
|
|
s.enqueue(task)
|
|
return nil
|
|
}
|