github.com/matrixorigin/matrixone@v1.2.0/pkg/vm/engine/tae/tasks/scheduler.go (about)

     1  // Copyright 2021 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package tasks
    16  
    17  import (
    18  	"context"
    19  
    20  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    21  
    22  	"github.com/matrixorigin/matrixone/pkg/logutil"
    23  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/common"
    24  	iops "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/tasks/ops/base"
    25  	ops "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/tasks/worker"
    26  )
    27  
    28  var (
    29  	ErrDispatcherNotFound = moerr.NewInternalErrorNoCtx("tae sched: dispatcher not found")
    30  	ErrSchedule           = moerr.NewInternalErrorNoCtx("tae sched: cannot schedule")
    31  )
    32  
    33  type Scheduler interface {
    34  	Start()
    35  	Stop()
    36  	Schedule(Task) error
    37  }
    38  
    39  type TaskScheduler interface {
    40  	Scheduler
    41  	ScheduleTxnTask(ctx *Context, taskType TaskType, factory TxnTaskFactory) (Task, error)
    42  	ScheduleMultiScopedTxnTask(ctx *Context, taskType TaskType, scopes []common.ID, factory TxnTaskFactory) (Task, error)
    43  	ScheduleMultiScopedFn(ctx *Context, taskType TaskType, scopes []common.ID, fn FuncT) (Task, error)
    44  	ScheduleFn(ctx *Context, taskType TaskType, fn func() error) (Task, error)
    45  	ScheduleScopedFn(ctx *Context, taskType TaskType, scope *common.ID, fn func() error) (Task, error)
    46  
    47  	CheckAsyncScopes(scopes []common.ID) error
    48  
    49  	GetCheckpointedLSN() uint64
    50  	GetPenddingLSNCnt() uint64
    51  }
    52  
    53  type BaseScheduler struct {
    54  	ops.OpWorker
    55  	idAlloc     *common.IdAlloctor
    56  	Dispatchers map[TaskType]Dispatcher
    57  }
    58  
    59  func NewBaseScheduler(ctx context.Context, name string) *BaseScheduler {
    60  	scheduler := &BaseScheduler{
    61  		OpWorker:    *ops.NewOpWorker(ctx, name),
    62  		idAlloc:     common.NewIdAlloctor(1),
    63  		Dispatchers: make(map[TaskType]Dispatcher),
    64  	}
    65  	scheduler.ExecFunc = scheduler.doDispatch
    66  	return scheduler
    67  }
    68  
    69  func (s *BaseScheduler) RegisterDispatcher(t TaskType, dispatcher Dispatcher) {
    70  	s.Dispatchers[t] = dispatcher
    71  }
    72  
    73  func (s *BaseScheduler) Schedule(task Task) error {
    74  	// task.AttachID(s.idAlloc())
    75  	if !s.SendOp(task) {
    76  		return ErrSchedule
    77  	}
    78  	return nil
    79  }
    80  
    81  func (s *BaseScheduler) doDispatch(op iops.IOp) {
    82  	task := op.(Task)
    83  	dispatcher := s.Dispatchers[task.Type()]
    84  	if dispatcher == nil {
    85  		logutil.Errorf("No dispatcher found for %d[T] Task", task.Type())
    86  		panic(ErrDispatcherNotFound)
    87  	}
    88  	dispatcher.Dispatch(task)
    89  }
    90  
    91  func (s *BaseScheduler) Stop() {
    92  	s.OpWorker.Stop()
    93  	for _, d := range s.Dispatchers {
    94  		d.Close()
    95  	}
    96  }