package engine import ( "context" "fmt" "sync" "sync/atomic" "time" "github.com/robfig/cron/v3" "go.uber.org/zap" ) // Job 定时任务接口 type Job interface { // ID 返回任务唯一标识 ID() string // Start 启动任务 Start(ctx context.Context) error // Stop 停止任务 Stop() error // IsRunning 检查任务是否正在运行 IsRunning() bool // NextRun 返回下次执行时间 NextRun() time.Time } // JobFunc 任务执行函数类型 type JobFunc func(ctx context.Context) error // Scheduler 定时任务调度器 type Scheduler struct { jobs map[string]Job mu sync.RWMutex logger *zap.Logger ctx context.Context cancel context.CancelFunc wg sync.WaitGroup running int32 } // NewScheduler 创建新的调度器 func NewScheduler(logger *zap.Logger) *Scheduler { ctx, cancel := context.WithCancel(context.Background()) return &Scheduler{ jobs: make(map[string]Job), logger: logger.Named("scheduler"), ctx: ctx, cancel: cancel, } } // Start 启动调度器 func (s *Scheduler) Start() error { if !atomic.CompareAndSwapInt32(&s.running, 0, 1) { return fmt.Errorf("scheduler is already running") } s.mu.RLock() defer s.mu.RUnlock() // 启动所有任务 for id, job := range s.jobs { if err := job.Start(s.ctx); err != nil { s.logger.Error("Failed to start job", zap.String("job_id", id), zap.Error(err)) continue } s.logger.Info("Job started", zap.String("job_id", id)) } s.logger.Info("Scheduler started", zap.Int("job_count", len(s.jobs))) return nil } // Stop 停止调度器 func (s *Scheduler) Stop() error { if !atomic.CompareAndSwapInt32(&s.running, 1, 0) { return fmt.Errorf("scheduler is not running") } s.cancel() s.mu.RLock() defer s.mu.RUnlock() // 停止所有任务 for id, job := range s.jobs { if err := job.Stop(); err != nil { s.logger.Error("Failed to stop job", zap.String("job_id", id), zap.Error(err)) continue } s.logger.Info("Job stopped", zap.String("job_id", id)) } s.wg.Wait() s.logger.Info("Scheduler stopped") return nil } // AddJob 添加任务 func (s *Scheduler) AddJob(job Job) error { s.mu.Lock() defer s.mu.Unlock() id := job.ID() if _, exists := s.jobs[id]; exists { return fmt.Errorf("job with id %s already exists", id) } s.jobs[id] = job // 如果调度器正在运行,立即启动任务 if atomic.LoadInt32(&s.running) == 1 { if err := job.Start(s.ctx); err != nil { delete(s.jobs, id) return fmt.Errorf("failed to start job: %w", err) } s.logger.Info("Job added and started", zap.String("job_id", id)) } else { s.logger.Info("Job added", zap.String("job_id", id)) } return nil } // RemoveJob 移除任务 func (s *Scheduler) RemoveJob(id string) error { s.mu.Lock() defer s.mu.Unlock() job, exists := s.jobs[id] if !exists { return fmt.Errorf("job with id %s not found", id) } if err := job.Stop(); err != nil { s.logger.Error("Failed to stop job during removal", zap.String("job_id", id), zap.Error(err)) } delete(s.jobs, id) s.logger.Info("Job removed", zap.String("job_id", id)) return nil } // GetJob 获取任务 func (s *Scheduler) GetJob(id string) (Job, bool) { s.mu.RLock() defer s.mu.RUnlock() job, exists := s.jobs[id] return job, exists } // GetAllJobs 获取所有任务 func (s *Scheduler) GetAllJobs() map[string]Job { s.mu.RLock() defer s.mu.RUnlock() result := make(map[string]Job, len(s.jobs)) for id, job := range s.jobs { result[id] = job } return result } // IsRunning 检查调度器是否正在运行 func (s *Scheduler) IsRunning() bool { return atomic.LoadInt32(&s.running) == 1 } // ============================================================================ // Job 实现 // ============================================================================ // CronJob 基于 Cron 表达式的任务 type CronJob struct { id string spec string handler JobFunc cron *cron.Cron logger *zap.Logger running int32 nextRun time.Time mu sync.RWMutex } // NewCronJob 创建 Cron 任务 func NewCronJob(id, spec string, handler JobFunc, logger *zap.Logger) (*CronJob, error) { parser := cron.NewParser(cron.Second | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor) c := cron.New(cron.WithParser(parser), cron.WithChain(cron.Recover(cron.DefaultLogger))) job := &CronJob{ id: id, spec: spec, handler: handler, cron: c, logger: logger.Named("cron-job").With(zap.String("job_id", id)), } // 添加任务到 cron _, err := c.AddFunc(spec, func() { ctx := context.Background() if err := handler(ctx); err != nil { job.logger.Error("Cron job execution failed", zap.Error(err)) } // 更新下次执行时间 entries := c.Entries() job.mu.Lock() if len(entries) > 0 { // 找到最近的执行时间 next := entries[0].Next for _, entry := range entries { if entry.Next.Before(next) { next = entry.Next } } job.nextRun = next } job.mu.Unlock() }) if err != nil { return nil, fmt.Errorf("invalid cron spec: %w", err) } // 计算初始下次执行时间(需要先启动 cron 才能计算) // 这里先设置为零值,在 Start 时再计算 job.mu.Lock() job.nextRun = time.Time{} job.mu.Unlock() return job, nil } func (j *CronJob) ID() string { return j.id } func (j *CronJob) Start(ctx context.Context) error { if !atomic.CompareAndSwapInt32(&j.running, 0, 1) { return fmt.Errorf("job is already running") } j.cron.Start() j.logger.Info("Cron job started", zap.String("spec", j.spec)) // 更新下次执行时间 entries := j.cron.Entries() if len(entries) > 0 { j.mu.Lock() // 找到最近的执行时间 next := entries[0].Next for _, entry := range entries { if entry.Next.Before(next) { next = entry.Next } } j.nextRun = next j.mu.Unlock() } return nil } func (j *CronJob) Stop() error { if !atomic.CompareAndSwapInt32(&j.running, 1, 0) { return fmt.Errorf("job is not running") } ctx := j.cron.Stop() <-ctx.Done() j.logger.Info("Cron job stopped") return nil } func (j *CronJob) IsRunning() bool { return atomic.LoadInt32(&j.running) == 1 } func (j *CronJob) NextRun() time.Time { j.mu.RLock() defer j.mu.RUnlock() return j.nextRun } // IntervalJob 固定间隔的任务 type IntervalJob struct { id string interval time.Duration handler JobFunc logger *zap.Logger running int32 nextRun time.Time mu sync.RWMutex ticker *time.Ticker ctx context.Context cancel context.CancelFunc wg sync.WaitGroup } // NewIntervalJob 创建固定间隔任务 func NewIntervalJob(id string, interval time.Duration, handler JobFunc, logger *zap.Logger) *IntervalJob { return &IntervalJob{ id: id, interval: interval, handler: handler, logger: logger.Named("interval-job").With(zap.String("job_id", id)), } } func (j *IntervalJob) ID() string { return j.id } func (j *IntervalJob) Start(ctx context.Context) error { if !atomic.CompareAndSwapInt32(&j.running, 0, 1) { return fmt.Errorf("job is already running") } j.ctx, j.cancel = context.WithCancel(ctx) j.ticker = time.NewTicker(j.interval) j.mu.Lock() j.nextRun = time.Now().Add(j.interval) j.mu.Unlock() j.wg.Add(1) go j.run() j.logger.Info("Interval job started", zap.Duration("interval", j.interval)) return nil } func (j *IntervalJob) run() { defer j.wg.Done() // 立即执行一次(可选,根据需求调整) // if err := j.handler(j.ctx); err != nil { // j.logger.Error("Interval job execution failed", zap.Error(err)) // } for { select { case <-j.ticker.C: j.mu.Lock() j.nextRun = time.Now().Add(j.interval) j.mu.Unlock() if err := j.handler(j.ctx); err != nil { j.logger.Error("Interval job execution failed", zap.Error(err)) } case <-j.ctx.Done(): return } } } func (j *IntervalJob) Stop() error { if !atomic.CompareAndSwapInt32(&j.running, 1, 0) { return fmt.Errorf("job is not running") } if j.cancel != nil { j.cancel() } if j.ticker != nil { j.ticker.Stop() } j.wg.Wait() j.logger.Info("Interval job stopped") return nil } func (j *IntervalJob) IsRunning() bool { return atomic.LoadInt32(&j.running) == 1 } func (j *IntervalJob) NextRun() time.Time { j.mu.RLock() defer j.mu.RUnlock() return j.nextRun } // OnceJob 单次延迟执行的任务 type OnceJob struct { id string delay time.Duration handler JobFunc logger *zap.Logger running int32 nextRun time.Time mu sync.RWMutex timer *time.Timer ctx context.Context cancel context.CancelFunc wg sync.WaitGroup } // NewOnceJob 创建单次延迟执行任务 func NewOnceJob(id string, delay time.Duration, handler JobFunc, logger *zap.Logger) *OnceJob { return &OnceJob{ id: id, delay: delay, handler: handler, logger: logger.Named("once-job").With(zap.String("job_id", id)), } } func (j *OnceJob) ID() string { return j.id } func (j *OnceJob) Start(ctx context.Context) error { if !atomic.CompareAndSwapInt32(&j.running, 0, 1) { return fmt.Errorf("job is already running") } j.ctx, j.cancel = context.WithCancel(ctx) j.timer = time.NewTimer(j.delay) j.mu.Lock() j.nextRun = time.Now().Add(j.delay) j.mu.Unlock() j.wg.Add(1) go j.run() j.logger.Info("Once job started", zap.Duration("delay", j.delay)) return nil } func (j *OnceJob) run() { defer j.wg.Done() select { case <-j.timer.C: if err := j.handler(j.ctx); err != nil { j.logger.Error("Once job execution failed", zap.Error(err)) } atomic.StoreInt32(&j.running, 0) case <-j.ctx.Done(): if !j.timer.Stop() { <-j.timer.C } return } } func (j *OnceJob) Stop() error { if !atomic.CompareAndSwapInt32(&j.running, 1, 0) { return fmt.Errorf("job is not running") } if j.cancel != nil { j.cancel() } if j.timer != nil { if !j.timer.Stop() { <-j.timer.C } } j.wg.Wait() j.logger.Info("Once job stopped") return nil } func (j *OnceJob) IsRunning() bool { return atomic.LoadInt32(&j.running) == 1 } func (j *OnceJob) NextRun() time.Time { j.mu.RLock() defer j.mu.RUnlock() return j.nextRun } // ============================================================================ // 全局调度器 API(链式风格,延迟注册) // ============================================================================ var ( globalJobRegistry = make([]JobBuilder, 0) globalJobMu sync.RWMutex jobCounter int64 ) // JobBuilder 任务构建器接口(延迟注册) type JobBuilder interface { // Build 构建任务实例(由依赖注入系统调用) Build(logger *zap.Logger) (Job, error) } // generateJobID 生成任务 ID func generateJobID(prefix string) string { counter := atomic.AddInt64(&jobCounter, 1) return fmt.Sprintf("%s_%d", prefix, counter) } // CronJobBuilder Cron 任务构建器 type CronJobBuilder struct { id string spec string handler JobFunc } // Cron 创建 Cron 任务构建器(在 init 函数中调用) func Cron(spec string) *CronJobBuilder { return &CronJobBuilder{ id: generateJobID("cron"), spec: spec, } } // Handle 设置处理函数并注册到全局注册表(延迟注册) func (b *CronJobBuilder) Handle(handler JobFunc) { b.handler = handler if b.handler == nil { panic("scheduler: handler cannot be nil") } globalJobMu.Lock() defer globalJobMu.Unlock() globalJobRegistry = append(globalJobRegistry, b) } // Build 构建 Cron 任务 func (b *CronJobBuilder) Build(logger *zap.Logger) (Job, error) { return NewCronJob(b.id, b.spec, b.handler, logger) } // IntervalJobBuilder 固定间隔任务构建器 type IntervalJobBuilder struct { id string interval time.Duration handler JobFunc } // Interval 创建固定间隔任务构建器(在 init 函数中调用) func Interval(interval time.Duration) *IntervalJobBuilder { return &IntervalJobBuilder{ id: generateJobID("interval"), interval: interval, } } // Handle 设置处理函数并注册到全局注册表(延迟注册) func (b *IntervalJobBuilder) Handle(handler JobFunc) { b.handler = handler if b.handler == nil { panic("scheduler: handler cannot be nil") } globalJobMu.Lock() defer globalJobMu.Unlock() globalJobRegistry = append(globalJobRegistry, b) } // Build 构建固定间隔任务 func (b *IntervalJobBuilder) Build(logger *zap.Logger) (Job, error) { return NewIntervalJob(b.id, b.interval, b.handler, logger), nil } // OnceJobBuilder 单次延迟任务构建器 type OnceJobBuilder struct { id string delay time.Duration handler JobFunc } // Once 创建单次延迟任务构建器(在 init 函数中调用) func Once(delay time.Duration) *OnceJobBuilder { return &OnceJobBuilder{ id: generateJobID("once"), delay: delay, } } // Handle 设置处理函数并注册到全局注册表(延迟注册) func (b *OnceJobBuilder) Handle(handler JobFunc) { b.handler = handler if b.handler == nil { panic("scheduler: handler cannot be nil") } globalJobMu.Lock() defer globalJobMu.Unlock() globalJobRegistry = append(globalJobRegistry, b) } // Build 构建单次延迟任务 func (b *OnceJobBuilder) Build(logger *zap.Logger) (Job, error) { return NewOnceJob(b.id, b.delay, b.handler, logger), nil } // LoadAllJobs 加载所有已注册的任务(由依赖注入系统调用) func LoadAllJobs(scheduler *Scheduler, logger *zap.Logger) error { globalJobMu.RLock() defer globalJobMu.RUnlock() for i, builder := range globalJobRegistry { job, err := builder.Build(logger) if err != nil { logger.Error("Failed to build job", zap.Int("index", i), zap.Error(err)) continue } if err := scheduler.AddJob(job); err != nil { logger.Error("Failed to add job to scheduler", zap.String("job_id", job.ID()), zap.Error(err)) continue } logger.Debug("Job loaded", zap.String("job_id", job.ID())) } logger.Info("All scheduled jobs loaded", zap.Int("job_count", len(globalJobRegistry))) return nil }