go-zero中定时任务的用法

文章目录

使用扩展
定义调度器
测试方法

使用扩展

在go-zero框架中使用定时任务调度的写法示例，首先需要用到的扩展：go get -u github.com/robfig/cron/v3

扩展网址：robfig/cron: a cron library for go (github.com)

定义调度器

在 gozero/internal 目录下新建 crontab目录，然后首先需要创建一个配置文件：gozero/internal/crontab/config.go ，在这里我写上两个示例的定时任务，分别为“20秒打印一次SayHello”，以及“每一分钟查询一次当前用户总数”：

package crontab

type TaskConfig struct {
	Name      string `json:"name"`      //任务名称
	Scheduler string `json:"scheduler"` //@every表达式中的单位可以是s（秒）、m（分钟）、h（小时）等
	Enabled   bool   // 是否启用
}

func DefaultTask() []TaskConfig {
	return []TaskConfig{
		{
			Name:      "SayHello",
			Scheduler: "@every 20s", //每20秒钟执行一次
			Enabled:   true,
		},
		{
			Name:      "StatisticsUserCount",
			Scheduler: "@every 1m", //每1分钟执行一次
			Enabled:   true,
		},
	}
}

然后，定义定时任务调度器文件：gozero/internal/crontab/scheduler.go，这个文件中，先定义好调度器的结构体和注册方法以及启动方法。代码如下：

type Scheduler struct {
	cron   *cron.Cron
	svcCtx *svc.ServiceContext
}

// NewScheduler creates a new scheduler instance.
func NewScheduler(svcCtx *svc.ServiceContext) *Scheduler {
	return &Scheduler{
		cron:   cron.New(),
		svcCtx: svcCtx,
	}
}

// RegisterTask registers a task with the scheduler.
func (s *Scheduler) RegisterTask(name, schedule string, task func(ctx context.Context, svcCtx *svc.ServiceContext)) {
	_, err := s.cron.AddFunc(schedule, func() {
		task(context.Background(), s.svcCtx)
	})
	if err != nil {
		log.Fatalf("Failed to register task %s: %v", name, err)
	}
	log.Printf("Registered task %s with schedule %s", name, schedule)
}

func (s *Scheduler) Start() {
	s.cron.Start()
	log.Println("Scheduler started")
}

func (s *Scheduler) Stop() {
	s.cron.Stop()
	log.Println("Scheduler stopped")
}

上面的 Start() 和 Stop() 方法，分别实现了 vendor/github.com/robfig/cron/v3/cron.go 中的 Start() 和 Stop() 方法。

接下来，需要初始化定时任务调度器，对于config.go中配置好的定时任务，分别实现不同的业务逻辑。

// InitScheduler 初始化定时任务调度器
func InitScheduler(svcCtx *svc.ServiceContext) (*Scheduler, error) {
	log.Println("Initializing scheduler...")
	scheduler := NewScheduler(svcCtx)
	if scheduler == nil {
		return nil, fmt.Errorf("failed to create scheduler")
	}
	taskConfigs := DefaultTask()
	if taskConfigs == nil {
		return nil, fmt.Errorf("failed to load task configurations")
	}
	for _, taskConfig := range taskConfigs {
		if !taskConfig.Enabled {
			log.Printf("Task %s is disabled, skipping...", taskConfig.Name)
			continue
		}

		log.Printf("Registering task %s...", taskConfig.Name)

		switch taskConfig.Name {
		case "SayHello": //SayHello
			scheduler.RegisterTask(taskConfig.Name, taskConfig.Scheduler, hello.SayHello)
		case "StatisticsUserCount": //统计用户总数
			scheduler.RegisterTask(taskConfig.Name, taskConfig.Scheduler, user.StatisticsUserCount)
		case "xxxx":
			// todo 注册其他定时任务
		default:
			log.Printf("Unknown task: %s", taskConfig.Name)
		}
	}
	return scheduler, nil
}

接下来，在入口文件gozero/gozero.go 的main方法中调用上面定义好的调度器：

func main() {
	//....

	// 定时任务调度
	scheduler, err := crontab.InitScheduler(ctx)
	if err != nil {
		fmt.Printf("Failed to initialize scheduler: %v", err)
	}
	scheduler.Start()

	fmt.Printf("Starting server at %s:%d...\n", c.Host, c.Port)
	server.Start()
}

测试方法

接下来，在 gozero/internal/crontab/task 目录下，新建两个测试方法，测试一下上面配置的两个定时任务：

SayHello方法： gozero/internal/crontab/task/hello/say_hello.go

func SayHello(ctx context.Context, svcCtx *svc.ServiceContext) {
	fmt.Println("SayHello running at:", time.Now())
	return
}

查询用户总数的方法：gozero/internal/crontab/task/user/statistics_user_count.go

package user

func StatisticsUserCount(ctx context.Context, svcCtx *svc.ServiceContext) {
	fmt.Println("StatisticsUserCount running at:", time.Now())

	listLogic := admin.NewUserListLogic(ctx, svcCtx)
	count, err := listLogic.GetUserCount() //调用逻辑层查询数据库
	if err != nil {
		fmt.Println("StatisticsUserCount error:", err)
		return
	}
	if count == 0 {
		fmt.Println("StatisticsUserCount count is 0")
		return
	}
	fmt.Println("StatisticsUserCount count:", count)
	return
}