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 crontabtype 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.CronsvcCtx *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": //SayHelloscheduler.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 userfunc 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
}