common/pkg/cron/cron.go

package cron

import (
	"context"
	"time"
)

type Cron struct {
	time   time.Time
	ctx    context.Context
	cancel context.CancelFunc
	f      func()
}

func NewCron() *Cron {
	now := time.Now()
	l, _ := time.LoadLocation("Asia/Shanghai")
	c := &Cron{
		time: time.Date(now.Year(), now.Month(), now.Day(), 01, 0, 0, 0, l),
	}
	return c
}

func (c *Cron) SetFunc(f func()) *Cron {
	c.f = f
	return c
}

func (c *Cron) SetContext(ctx context.Context) *Cron {
	c.ctx = ctx
	return c
}

func (c *Cron) SetTimeHHmm(HHmm string) *Cron {
	t, _ := time.Parse("15:04", HHmm)
	now := time.Now()
	l, _ := time.LoadLocation("Asia/Shanghai")
	c.time = time.Date(now.Year(), now.Month(), now.Day(), t.Hour(), t.Minute(), 0, 0, l)
	return c
}

func (c *Cron) Run(ctx context.Context) {
	c.ctx, c.cancel = context.WithCancel(ctx)
	for {
		if time.Now().After(c.time) {
			c.time = c.time.Add(time.Hour * 24)
		}
		select {
		case <-time.After(time.Until(c.time)):
			if c.f != nil {
				c.f()
			}
		case <-c.ctx.Done():
			return
		}
	}
}

func (c *Cron) Stop() {
	c.cancel()
}