gitee.com/h79/goutils@v1.22.10/common/scheduler/scheduler_group.go

package scheduler

import (
	"gitee.com/h79/goutils/common/logger"
	"gitee.com/h79/goutils/common/queue"
	"gitee.com/h79/goutils/common/system"
	"sync"
	"time"
)

var _ queue.IPriority = (*scheduledItem)(nil)

type scheduledItem struct {
	Trigger  Trigger
	Job      Task
	nextTime int64
}

// GValue for implement queue.IPriority
func (it *scheduledItem) GValue() interface{} {
	return it
}

// PValue for implement queue.IPriority
func (it *scheduledItem) PValue() int64 {
	return it.nextTime
}

type groupScheduled struct {
	rm              sync.Mutex
	items           *queue.Capacity
	tm              *time.Ticker
	uniques         map[string]*scheduledItem
	feederTaskCheck system.RunningCheck
	execTaskCheck   system.RunningCheck
	duration        time.Duration
	minDuration     time.Duration
	execTime        time.Duration
	location        *time.Location
	feeder          chan *scheduledItem
	interrupt       chan struct{}
	index           int
	logEnabled      bool
}

const feederCount = 3
const defTaskDuration = time.Microsecond * 500

func newGroupScheduled(idx int, capacity int, location *time.Location, minDuration, duration time.Duration, logEnabled bool) *groupScheduled {
	if duration < defTaskDuration {
		duration = defTaskDuration
	}
	if minDuration < defTaskDuration {
		minDuration = defTaskDuration
	}
	tm := time.NewTicker(duration)
	return &groupScheduled{
		items:       queue.NewCapacity(capacity),
		feeder:      make(chan *scheduledItem, feederCount*2),
		interrupt:   make(chan struct{}, feederCount*10),
		location:    location,
		duration:    duration,
		minDuration: minDuration,
		tm:          tm,
		index:       idx,
		logEnabled:  logEnabled,
		uniques:     map[string]*scheduledItem{},
	}
}

func (gs *groupScheduled) AddTask(task Task, uniqueEnabled bool, generateFn GenerateOrInitFunc) error {
	return gs.AddTriggerTask(task, NewSimpleTrigger(time.Second*5), uniqueEnabled, generateFn)
}

func (gs *groupScheduled) AddTriggerTask(task Task, trigger Trigger, uniqueEnabled bool, generateFn GenerateOrInitFunc) error {
	if system.IsQuit() {
		logger.W("Task", "application is quited,not can add for jobId= %s", task.GetId())
		return system.ClosedError
	}
	err := gs.add(task, trigger, uniqueEnabled, generateFn)
	if err == nil {
		gs.Run()
		gs.resetTaskTime()
	}
	return err
}

func defaultUpdateFn(_ int, _ Task) (Task, Trigger, bool) {
	return nil, nil, true
}

func (gs *groupScheduled) UpdateTask(jobId string, uniqueEnabled bool, notExistNewFlag bool, generateFn GenerateOrInitFunc) bool {
	ret := gs.update(jobId, uniqueEnabled, notExistNewFlag, generateFn)
	if ret && notExistNewFlag {
		gs.Run()
		gs.resetTaskTime()
	}
	return ret
}

func (gs *groupScheduled) add(task Task, trigger Trigger, uniqueEnabled bool, generateFn GenerateOrInitFunc) error {
	gs.rm.Lock()
	defer gs.rm.Unlock()
	var item *scheduledItem
	if uniqueEnabled {
		if it, ok := gs.uniques[task.GetId()]; ok {
			item = it
		}
		if item != nil {
			if generateFn != nil {
				generateFn(gs.index, item.Job)
			}
			return ErrJobExist
		}
		//go to add
	}
	nt, err := trigger.NextFireTime(time.Now().In(gs.location).UnixNano())
	if err != nil {
		return err
	}
	item = &scheduledItem{Trigger: trigger, Job: task, nextTime: nt}
	err = gs.items.Add(item)
	if err != nil {
		return err
	}
	if uniqueEnabled {
		gs.uniques[task.GetId()] = item
	}
	if generateFn != nil {
		generateFn(gs.index, task)
	}
	return nil
}

func (gs *groupScheduled) update(jobId string, uniqueEnabled bool, notExistNewFlag bool, generateFn GenerateOrInitFunc) bool {
	gs.rm.Lock()
	defer gs.rm.Unlock()
	var item *scheduledItem
	if uniqueEnabled {
		if it, ok := gs.uniques[jobId]; ok {
			item = it
		}
	} else {
		item = gs.find(jobId)
	}
	if item != nil {
		if generateFn != nil {
			_, _, ok := generateFn(gs.index, item.Job)
			return ok
		}
		return true
	}
	if notExistNewFlag && generateFn != nil {
		tk, trigger, _ := generateFn(gs.index, nil)
		if tk == nil || trigger == nil {
			return false
		}
		nt, err := trigger.NextFireTime(time.Now().In(gs.location).UnixNano())
		if err != nil {
			return false
		}
		item = &scheduledItem{Trigger: trigger, Job: tk, nextTime: nt}
		err = gs.items.Add(item)
		if err != nil {
			return false
		}
		if uniqueEnabled {
			gs.uniques[tk.GetId()] = item
		}
		_, _, ret := generateFn(gs.index, tk)
		return ret
	}
	return false
}

func (gs *groupScheduled) ForeachTask(update func(group int, task Task, time time.Duration, now, trigger int64) bool) {
	gs.rm.Lock()
	defer gs.rm.Unlock()
	now := time.Now().In(gs.location).UnixNano()
	gs.items.Foreach(func(v any, index int) bool {
		if item, ok := v.(*scheduledItem); ok {
			return update(gs.index, item.Job, gs.execTime, now, item.nextTime)
		}
		return false
	})
}

func (gs *groupScheduled) StopTask(jobId string) {
	gs.rm.Lock()
	defer gs.rm.Unlock()
	if task := gs.find(jobId); task != nil {
		task.Job.Cancel()
	}
}

func (gs *groupScheduled) RemoveTask(jobId string) {
	gs.rm.Lock()
	defer gs.rm.Unlock()
	if task := gs.find(jobId); task != nil {
		task.Job.Cancel()
	}
}

func (gs *groupScheduled) HasTask(jobId string) bool {
	gs.rm.Lock()
	defer gs.rm.Unlock()
	return gs.find(jobId) != nil
}

func (gs *groupScheduled) Run() {
	if system.IsQuit() {
		return
	}
	gs.execTaskCheck.GoRunning(gs.execTask)
	gs.feederTaskCheck.GoRunning(gs.runFeeder)
}

func (gs *groupScheduled) find(jobId string) *scheduledItem {
	job, idx := gs.items.Find(func(v interface{}, index int) bool {
		if item, ok := v.(*scheduledItem); ok {
			return item.Job.GetId() == jobId
		}
		return false
	})
	if idx == -1 {
		return nil
	}
	return job.(*scheduledItem)
}

func (gs *groupScheduled) peek() *scheduledItem {
	gs.rm.Lock()
	defer gs.rm.Unlock()
	var v = gs.items.Peek()
	item, ok := v.(*scheduledItem)
	if !ok {
		return nil
	}
	return item
}

func (gs *groupScheduled) calculateNextTime() time.Duration {
	job := gs.peek()
	if job != nil {
		now := time.Now().In(gs.location).UnixNano()
		if job.nextTime >= now {
			return time.Duration(job.nextTime - now)
		}
		return 0
	}
	return gs.duration
}

func (gs *groupScheduled) resetTaskTime() {
	duration := gs.calculateNextTime()
	if duration < gs.minDuration {
		duration = gs.minDuration
	}
	if gs.execTime == duration {
		return
	}
	gs.execTime = duration
	gs.tm.Reset(duration)
	if gs.logEnabled {
		logger.N("Schedule", "next run time= %d in group= %d", duration, gs.index)
	}
}

func (gs *groupScheduled) execTask() {
	defer gs.tm.Stop()
	for {
		gs.resetTaskTime()

		select {
		case <-system.Closed():
			gs.quited()
			return

		case <-gs.tm.C:
			for i := 0; i < feederCount-1; i++ {
				gs.runTask()
			}

		case <-gs.interrupt:
			break
		}
	}
}

func (gs *groupScheduled) quited() {
	var opts = With(2)
	for i := 0; i < len(gs.feeder); i++ {
		item, ok := <-gs.feeder
		if !ok {
			return
		}
		_, _ = item.Job.Execute(opts...)
	}
	gs.rm.Lock()
	defer gs.rm.Unlock()
	for {
		var v = gs.items.Pop()
		var item, ok = v.(*scheduledItem)
		if !ok {
			break
		}
		_, _ = item.Job.Execute(opts...)
	}
}

func (gs *groupScheduled) addTask(item *scheduledItem) error {
	gs.rm.Lock()
	defer gs.rm.Unlock()
	now := time.Now().In(gs.location).UnixNano()
	nextTime, err := item.Trigger.NextFireTime(now)
	if err != nil {
		return err
	}
	if now-item.nextTime > (time.Second * 5).Nanoseconds() {
		if gs.logEnabled {
			logger.W("Schedule", "the task is block too long, jobId= '%s' in group=%d", item.Job.GetId(), gs.index)
		}
	}
	item.nextTime = nextTime
	return gs.items.Add(item)
}

func (gs *groupScheduled) runFeeder() {
	for {
		select {
		case item := <-gs.feeder:
			err := gs.addTask(item)
			if err != nil {
				gs.removeUnique(item.Job.GetId())
				logger.E("Schedule", "add task failure for feeder, the task lost, jobId= '%s' in group=%d, err= %s", item.Job.GetId(), gs.index, err)
				break
			}
			gs.interrupt <- struct{}{}

		case <-system.Closed():
			return
		}
	}
}

func (gs *groupScheduled) pop(now int64) (*scheduledItem, bool) {
	gs.rm.Lock()
	defer gs.rm.Unlock()

	var v = gs.items.Peek()
	var item, ok = v.(*scheduledItem)
	if !ok {
		return nil, false
	}
	if gs.logEnabled {
		logger.N("Schedule", "pop the task,jobId= '%s' in group= %d, time= (%d,%d), ex= %v", item.Job.GetId(), gs.index, item.nextTime, now, item.nextTime <= now)
	}
	if item.nextTime > now {
		return nil, false
	}
	var it = gs.items.Pop()
	if it != item {
		panic("the value not equal")
	}
	return item, true
}

func (gs *groupScheduled) runTask() bool {
	// 没到时间，不用执行
	now := time.Now().In(gs.location).UnixNano()
	item, ok := gs.pop(now)
	if !ok {
		return false
	}
	_, err := item.Job.Execute()
	if err != nil {
		gs.removeUnique(item.Job.GetId())
		if gs.logEnabled {
			logger.E("Schedule", "the task execute error, jobId= '%s' in group=%d, err= %s", item.Job.GetId(), gs.index, err)
		}
		return true
	}
	if item.Job.GetState().IsQuit() {
		gs.removeUnique(item.Job.GetId())
		if gs.logEnabled {
			logger.E("Schedule", "the task quited, state= '%s', jobId= '%s' in group= %d", item.Job.GetState(), item.Job.GetId(), gs.index)
		}
		return true
	}
	if system.IsQuit() {
		gs.removeUnique(item.Job.GetId())
		if gs.logEnabled {
			logger.E("Schedule", "system quited, state= '%s', jobId= '%s' in group= %d", item.Job.GetState(), item.Job.GetId(), gs.index)
		}
		return false
	}
	item.nextTime = now
	gs.feeder <- item
	return true
}

func (gs *groupScheduled) removeUnique(jobId string) {
	gs.rm.Lock()
	defer gs.rm.Unlock()
	delete(gs.uniques, jobId)
}