web_huakui_server/pub/common/cronx/cron.go

package cronx

import (
	"context"
	"hkdj-web-server/pub/common/log"
	"sync"
	"time"

	"github.com/robfig/cron/v3"
	"go.uber.org/zap"
)

type CronJob interface {
	Run() error
	Desc() string
	Name() string
}

type CronJobAndSpec struct {
	CronJob CronJob
	Spec    string
}

type BaseJob struct {
	ctx  context.Context
	spec string
	name string
	desc string
	l    *zap.SugaredLogger
}

type cronFunc struct {
	*BaseJob
	fn func(ctx context.Context, l *zap.SugaredLogger) error
}

func addError(fn func(ctx context.Context, l *zap.SugaredLogger)) func(context.Context, *zap.SugaredLogger) error {
	return func(ctx context.Context, l *zap.SugaredLogger) error {
		l.Info("start")
		defer l.Info("end")
		fn(ctx, l)
		return nil
	}
}

func (f cronFunc) Run() error { return f.fn(f.ctx, f.l) }
func NewCronFunc(ctx context.Context, name, desc string, fn func(ctx context.Context, l *zap.SugaredLogger)) cronFunc {
	return cronFunc{
		BaseJob: NewBaseJob(ctx, name, desc),
		fn:      addError(fn),
	}
}

func NewCronFuncWithError(ctx context.Context, name, desc string, fn func(ctx context.Context, l *zap.SugaredLogger) error) cronFunc {
	return cronFunc{
		BaseJob: NewBaseJob(ctx, name, desc),
		fn:      fn,
	}
}

func NewBaseJob(ctx context.Context, name, desc string) *BaseJob {
	return &BaseJob{
		ctx:  ctx,
		name: name,
		desc: desc,
		l:    log.ZapLogS.Named(name),
	}
}

func (j *BaseJob) Name() string             { return j.name }
func (j *BaseJob) Desc() string             { return j.desc }
func (j *BaseJob) L() *zap.SugaredLogger    { return j.l }
func (j *BaseJob) Context() context.Context { return j.ctx }

type CronLogger struct {
	l *zap.SugaredLogger
}

func NewCronLogger() *CronLogger {
	return &CronLogger{
		l: log.ZapLogS.Named("cron"),
	}
}

func (l *CronLogger) Info(msg string, kvs ...interface{}) {
	l.l.Infow(msg, kvs)
}
func (l *CronLogger) Error(err error, msg string, kvs ...interface{}) {
	l.l.Errorw(err.Error()+msg, kvs)
}

func StartJobs(jobs ...CronJobAndSpec) (closer func()) {
	var once sync.Once
	l := NewCronLogger()
	c := cron.New(cron.WithSeconds(), cron.WithLocation(time.Local), cron.WithChain(cron.Recover(l)))
	for i := range jobs {
		j := jobs[i].CronJob
		spec := jobs[i].Spec
		name := j.Name()
		_, err := c.AddFunc(spec, func() {
			defer log.Infof("%s end", name)
			log.Infof("%s start", name)
			if err := j.Run(); err != nil {
				log.Errorf("%s run failed: %s", name, err.Error())
			}
		})
		if err != nil {
			log.Fatalf("%s add failed: %s", err)
		}
	}
	c.Start()
	return func() {
		once.Do(func() {
			c.Stop()
		})
	}
}

type CronFuncType func(ctx context.Context, l *zap.SugaredLogger)

func TraceTime(fn CronFuncType) CronFuncType {
	return func(ctx context.Context, l *zap.SugaredLogger) {
		now := time.Now()
		defer func() {
			elapse := time.Since(now)
			if elapse > time.Second {
				l.Infof("takes %s", elapse.Truncate(time.Millisecond))
			} else {
				l.Infof("takes %d ms", elapse/time.Millisecond)
			}
		}()
		l.Info("start run")
		fn(ctx, l)
		l.Info("end run")
	}
}

// MustSecondSpec is a helper method to validate spec at compile time
func MustSecondSpec(spec string) string {
	parser := cron.NewParser(
		cron.Second | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor,
	)
	_, err := parser.Parse(spec)
	if err != nil {
		panic(err)
	}
	return spec
}