github.com/ngicks/gokugen@v0.0.5/scheduler/scheduler.go (about)

     1  package scheduler
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/ngicks/gokugen/common"
    10  )
    11  
    12  // Scheduler is an in-memory scheduler backed by min heap.
    13  type Scheduler struct {
    14  	workingState
    15  	endState
    16  	wg sync.WaitGroup
    17  
    18  	taskTimer     *TaskTimer
    19  	cancellerLoop *CancellerLoop
    20  	dispatchLoop  *DispatchLoop
    21  	workerPool    *WorkerPool
    22  	taskCh        chan *Task
    23  }
    24  
    25  func NewScheduler(initialWorkerNum, queueMax uint) *Scheduler {
    26  	return newScheduler(initialWorkerNum, queueMax, common.GetNowImpl{})
    27  }
    28  
    29  func newScheduler(initialWorkerNum, queueMax uint, getNow common.GetNower) *Scheduler {
    30  	taskCh := make(chan *Task)
    31  	taskTimer := NewTaskTimer(queueMax, getNow, common.NewTimerImpl())
    32  
    33  	s := &Scheduler{
    34  		taskTimer:     taskTimer,
    35  		cancellerLoop: NewCancellerLoop(taskTimer, getNow, time.Minute),
    36  		dispatchLoop:  NewDispatchLoop(taskTimer, getNow),
    37  		workerPool:    NewWorkerPool(BuildWorkerConstructor(taskCh, nil, nil)),
    38  		taskCh:        taskCh,
    39  	}
    40  	s.AddWorker(uint32(initialWorkerNum))
    41  	return s
    42  }
    43  
    44  func (s *Scheduler) Schedule(task *Task) (*TaskController, error) {
    45  	if s.IsEnded() {
    46  		return nil, ErrAlreadyEnded
    47  	}
    48  
    49  	err := s.dispatchLoop.PushTask(task)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	return &TaskController{t: task}, nil
    54  }
    55  
    56  type LoopError struct {
    57  	cancellerLoopErr error
    58  	dispatchLoopErr  error
    59  }
    60  
    61  func (e LoopError) Error() string {
    62  	return fmt.Sprintf(
    63  		"cancellerLoopErr: %s, dispatchLoopErr: %s",
    64  		e.cancellerLoopErr,
    65  		e.dispatchLoopErr,
    66  	)
    67  }
    68  
    69  // Start starts needed loops.
    70  // Start creates one goroutine for periodical-removal-of-cancelled-task.
    71  // Start blocks until ctx is cancelled, and other loops to return.
    72  func (s *Scheduler) Start(ctx context.Context) error {
    73  	if s.IsEnded() {
    74  		return ErrAlreadyEnded
    75  	}
    76  	if !s.setWorking() {
    77  		return ErrAlreadyStarted
    78  	}
    79  	defer s.setWorking(false)
    80  
    81  	err := new(LoopError)
    82  	s.wg.Add(1)
    83  	s.taskTimer.Start()
    84  	go func() {
    85  		err.cancellerLoopErr = s.cancellerLoop.Start(ctx)
    86  		s.wg.Done()
    87  	}()
    88  	err.dispatchLoopErr = s.dispatchLoop.Start(ctx, s.taskCh)
    89  	s.wg.Wait()
    90  	s.taskTimer.Stop()
    91  
    92  	if err.cancellerLoopErr == nil && err.dispatchLoopErr == nil {
    93  		return nil
    94  	}
    95  	return err
    96  }
    97  
    98  func (s *Scheduler) AddWorker(delta uint32) (workerNum int) {
    99  	return int(s.workerPool.Add(delta))
   100  }
   101  
   102  func (s *Scheduler) RemoveWorker(delta uint32) (aliveWorkerNum int, sleepingWorkerNum int) {
   103  	return s.workerPool.Remove(delta)
   104  }
   105  
   106  func (s *Scheduler) ActiveWorkerNum() int64 {
   107  	return s.workerPool.ActiveWorkerNum()
   108  }
   109  
   110  // End remove all workers and let this scheduler to step into ended-state where no new Start is allowed.
   111  // End also cancel tasks if they are working on in any work.
   112  // Calling this method *before* cancelling of ctx passed to Start will cause blocking forever.
   113  func (s *Scheduler) End() {
   114  	s.setEnded()
   115  	// wait for the Start loop to be done.
   116  	s.wg.Wait()
   117  	s.workerPool.Kill()
   118  	s.workerPool.Wait()
   119  }