web_hgtv_server/pub/common/cronmgr/crommgr.go

package cronmgr

import (
	"context"
	"errors"
	"sync"
	"videoweb/pub/common/log"

	"github.com/robfig/cron/v3"
)

var c *cron.Cron

// runIngs 正在运行中的任务
var runIngs map[*context.Context]struct{}
var r_lock sync.Mutex

// RegisterCron 注册一个定时任务
func RegisterCron(spec string, cmd func(), errfunc func(err interface{})) (cron.EntryID, error) {
	if len(spec) == 0 || cmd == nil {
		return -1, errors.New("spec or cmd is nil")
	}
	return registerCron(spec, func() {
		defer func() {
			if err := recover(); err != nil {
				if errfunc != nil {
					errfunc(err)
				}
				log.Errorf("cron task error.%v", err)
			}
		}()
		cmd()
	})
}

// RegisterCronAsyncOne 注册一个协程执行的定时任务 限制了cmd同时只有一个协程在执行
func RegisterCronAsyncOne(spec string, cmd func(), errfunc func(err interface{})) (cron.EntryID, error) {
	if len(spec) == 0 || cmd == nil {
		return -1, errors.New("spec or cmd is nil")
	}
	ctx := context.Background()
	return registerCron(spec, func() {
		go func(c *context.Context) {
			defer func() {
				if err := recover(); err != nil {
					if errfunc != nil {
						errfunc(err)
					}
					log.Errorf("cron task error.%v", err)
				}
			}()
			if taskRunIng(c) {
				cmd()
				taskRunEnd(c)
			} else {
				log.Warnf("Limits cmd to only one coroutine executing at the same time")
			}
		}(&ctx)
	})
}

// RegisterCronAsync 注册一个协程执行的定时任务
func RegisterCronAsync(spec string, cmd func(), errfunc func(err interface{})) (cron.EntryID, error) {
	if len(spec) == 0 || cmd == nil {
		return -1, errors.New("spec or cmd is nil")
	}
	return registerCron(spec, func() {
		go func() {
			defer func() {
				if err := recover(); err != nil {
					if errfunc != nil {
						errfunc(err)
					}
					log.Errorf("cron task error.%v", err)
				}
			}()
			cmd()
		}()
	})
}
func registerCron(spec string, cmd func()) (cron.EntryID, error) {
	if c == nil {
		c = cron.New(cron.WithSeconds())
		runIngs = make(map[*context.Context]struct{})
	}
	id, err := c.AddFunc(spec, cmd)
	if err != nil {
		log.Warnf("add cron [%s] fail", spec)
		return -1, err
	}
	c.Start()
	log.Infof("add cron [%d] spec:[%s].", id, spec)
	return id, nil
}

// UnregisterCron 注销一个定时任务
func UnregisterCron(id cron.EntryID) bool {
	if id > 0 {
		c.Remove(id)
		log.Infof("remove cron id:[%d]", id)
	}
	log.Infof("remove cron [%d] not found", id)
	return false
}

func taskRunIng(ctx *context.Context) bool {
	r_lock.Lock()
	defer r_lock.Unlock()
	if _, ok := runIngs[ctx]; ok {
		return false
	}
	runIngs[ctx] = struct{}{}
	return true
}
func taskRunEnd(ctx *context.Context) {
	r_lock.Lock()
	defer r_lock.Unlock()
	delete(runIngs, ctx)
}