forked from gsn/predictor
97 lines
2 KiB
Go
97 lines
2 KiB
Go
package scheduler
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"git.intra.yksa.space/gsn/predictor/internal/pkg/errcodes"
|
|
"git.intra.yksa.space/gsn/predictor/internal/pkg/log"
|
|
"github.com/go-co-op/gocron"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type Job interface {
|
|
GetInterval() time.Duration
|
|
GetTimeout() time.Duration
|
|
GetCount() int
|
|
GetAsync() bool
|
|
Execute(context.Context) error
|
|
}
|
|
|
|
type Scheduler struct {
|
|
scheduler *gocron.Scheduler
|
|
}
|
|
|
|
func New() *Scheduler {
|
|
scheduler := gocron.NewScheduler(time.UTC)
|
|
return &Scheduler{
|
|
scheduler: scheduler,
|
|
}
|
|
}
|
|
|
|
func (s *Scheduler) AddJob(job Job) error {
|
|
interval := job.GetInterval()
|
|
timeout := job.GetTimeout()
|
|
count := job.GetCount()
|
|
async := job.GetAsync()
|
|
|
|
// Validate job parameters
|
|
if !async && count != 1 {
|
|
return errcodes.ErrSchedulerInvalidJob
|
|
}
|
|
if timeout > interval {
|
|
return errcodes.ErrSchedulerTimeoutTooLong
|
|
}
|
|
|
|
// Create job function with timeout
|
|
jobFunc := func() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
|
defer cancel()
|
|
logger := log.Ctx(ctx)
|
|
if err := job.Execute(ctx); err != nil {
|
|
logger.Error("job execution failed",
|
|
zap.Error(err),
|
|
zap.Duration("interval", interval),
|
|
zap.Duration("timeout", timeout))
|
|
} else {
|
|
logger.Debug("job executed successfully",
|
|
zap.Duration("interval", interval),
|
|
zap.Duration("timeout", timeout))
|
|
}
|
|
}
|
|
|
|
// Add job to scheduler
|
|
schedulerJob := s.scheduler.Every(interval)
|
|
|
|
if !async {
|
|
schedulerJob = schedulerJob.SingletonMode()
|
|
}
|
|
|
|
if count > 0 {
|
|
schedulerJob = schedulerJob.LimitRunsTo(count)
|
|
}
|
|
|
|
schedulerJob.Do(jobFunc)
|
|
|
|
log.Ctx(context.Background()).Info("job added to scheduler",
|
|
zap.Duration("interval", interval),
|
|
zap.Duration("timeout", timeout),
|
|
zap.Int("count", count),
|
|
zap.Bool("async", async))
|
|
|
|
return nil
|
|
}
|
|
|
|
func (s *Scheduler) Start() {
|
|
s.scheduler.StartAsync()
|
|
log.Ctx(context.Background()).Info("scheduler started")
|
|
}
|
|
|
|
func (s *Scheduler) Stop() {
|
|
s.scheduler.Stop()
|
|
log.Ctx(context.Background()).Info("scheduler stopped")
|
|
}
|
|
|
|
func (s *Scheduler) IsRunning() bool {
|
|
return s.scheduler.IsRunning()
|
|
}
|