diff --git a/scheduler.go b/scheduler.go index 90f6f57..dd81035 100644 --- a/scheduler.go +++ b/scheduler.go @@ -27,15 +27,15 @@ var globalSchedulerMux *Schedulermux type Schedulermux struct { registry *scheduler.Registry store *scheduler.Store - semaphore *scheduler.Semaphore + semaphore *scheduler.PersistedSemaphore sched *scheduler.Scheduler } -// SchedulerInit initializes the scheduler internals (registry and semaphore). +// SchedulerInit initializes the scheduler internals (registry). // The store must be set separately via SetStore before starting. +// The semaphore is created lazily during RunScheduler once the DB is available. func (s *Schedulermux) SchedulerInit() { s.registry = scheduler.NewRegistry() - s.semaphore = scheduler.NewSemaphore(true) globalSchedulerMux = s } @@ -56,7 +56,8 @@ func (s *Schedulermux) GetStore() *scheduler.Store { } // RunScheduler starts the scheduler loop with the given configuration. -// It creates the internal Scheduler instance and launches it in a goroutine. +// It auto-migrates the scheduler_meta table, recovers the persisted semaphore +// state, creates the in-memory semaphore, and launches the scheduler goroutine. func (s *Schedulermux) RunScheduler(config SchedulerConfig) { if s.store == nil { log.Fatal("scheduler: store is not set — call SetStore before RunScheduler") @@ -72,6 +73,24 @@ func (s *Schedulermux) RunScheduler(config SchedulerConfig) { rateLimit = DefaultSchedulerRateLimit } + ctx := context.Background() + + // Auto-migrate the scheduler_meta table for semaphore persistence. + _ = s.store.DB().AutoMigrate(&scheduler.SchedulerMeta{}) + + // Recover the persisted semaphore state (defaults to green). + wasGreen := s.store.GetSemaphore(ctx) + s.semaphore = scheduler.NewPersistedSemaphore(s.store, wasGreen) + logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] semaphore recovered: %s", map[bool]string{true: "green", false: "red"}[wasGreen])) + + // Recover stuck items from a previous crash before starting. + recovered, err := s.store.RecoverStuck(ctx) + if err != nil { + logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] recover stuck error: %v", err)) + } else if recovered > 0 { + logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] recovered %d stuck task(s) back to pending", recovered)) + } + s.sched = scheduler.NewScheduler( scheduler.SchedulerConfig{ Interval: interval, @@ -82,15 +101,6 @@ func (s *Schedulermux) RunScheduler(config SchedulerConfig) { s.registry, ) - // Recover stuck items from a previous crash before starting. - ctx := context.Background() - recovered, err := s.store.RecoverStuck(ctx) - if err != nil { - logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] recover stuck error: %v", err)) - } else if recovered > 0 { - logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] recovered %d stuck task(s) back to pending", recovered)) - } - s.sched.Start(ctx) } diff --git a/scheduler/scheduler.go b/scheduler/scheduler.go index d458a6e..c314009 100644 --- a/scheduler/scheduler.go +++ b/scheduler/scheduler.go @@ -72,6 +72,17 @@ type ProcessedItem struct { // TableName overrides the GORM default. func (ProcessedItem) TableName() string { return "processed_items" } +// SchedulerMeta stores key-value metadata for the scheduler, such as the +// semaphore state. This table persists across restarts so the scheduler can +// recover its previous state. +type SchedulerMeta struct { + Key string `gorm:"primaryKey;size:64"` + Value string `gorm:"size:16;not null"` +} + +// TableName overrides the GORM default. +func (SchedulerMeta) TableName() string { return "scheduler_meta" } + // ─── Handler ───────────────────────────────────────────────────────────────── // HandlerFunc processes a single queue item. @@ -124,6 +135,11 @@ func NewStore(db *gorm.DB) (*Store, error) { return &Store{db: db}, nil } +// DB returns the underlying *gorm.DB used by the store. +func (s *Store) DB() *gorm.DB { + return s.db +} + // Enqueue inserts a new pending task. // maxRuns controls how many times the task executes: 1 = once (default), -1 = repeat indefinitely. // thread controls concurrency: 0 = sequential per type, 1 = concurrent per type. @@ -242,6 +258,44 @@ func (s *Store) RecordResult(ctx context.Context, item *QueueItem, startedAt tim }) } +// SemaphoreKey is the key used in scheduler_meta for the semaphore state. +const SemaphoreKey = "semaphore" + +// SemaphoreValueGreen is the stored value for a green semaphore. +const SemaphoreValueGreen = "green" + +// SemaphoreValueRed is the stored value for a red semaphore. +const SemaphoreValueRed = "red" + +// SetSemaphore persists the semaphore state to the database. +func (s *Store) SetSemaphore(ctx context.Context, green bool) error { + value := SemaphoreValueRed + if green { + value = SemaphoreValueGreen + } + return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var count int64 + tx.Model(&SchedulerMeta{}).Where("key = ?", SemaphoreKey).Count(&count) + if count == 0 { + return tx.Create(&SchedulerMeta{Key: SemaphoreKey, Value: value}).Error + } + return tx.Model(&SchedulerMeta{}). + Where("key = ?", SemaphoreKey). + Update("value", value).Error + }) +} + +// GetSemaphore reads the persisted semaphore state from the database. +// Returns true if green, false if red or not found. +func (s *Store) GetSemaphore(ctx context.Context) bool { + var meta SchedulerMeta + err := s.db.WithContext(ctx).Where("key = ?", SemaphoreKey).First(&meta).Error + if err != nil || meta.Value != SemaphoreValueGreen { + return false + } + return true +} + // PendingCount returns the number of tasks currently in pending or processing state. func (s *Store) PendingCount(ctx context.Context) (int64, error) { var count int64 @@ -275,45 +329,83 @@ func (s *Store) RecoverStuck(ctx context.Context) (int64, error) { // ─── Semaphore ─────────────────────────────────────────────────────────────── -// Semaphore controls whether the scheduler is allowed to execute tasks. +// Semaphore is the interface for the scheduler kill switch. // Green (true) means tasks can run; Red (false) means tasks are blocked. -type Semaphore struct { +type Semaphore interface { + SetGreen() + SetRed() + IsGreen() bool +} + +// inMemorySemaphore is the default channel-based implementation. +type inMemorySemaphore struct { green chan bool } // NewSemaphore creates a Semaphore. Pass true to start green, false to start red. -func NewSemaphore(startGreen bool) *Semaphore { - s := &Semaphore{green: make(chan bool, 1)} +func NewSemaphore(startGreen bool) Semaphore { + s := &inMemorySemaphore{green: make(chan bool, 1)} s.green <- startGreen return s } -// SetGreen allows task execution to proceed. -func (s *Semaphore) SetGreen() { +func (s *inMemorySemaphore) SetGreen() { <-s.green s.green <- true } -// SetRed blocks task execution. -func (s *Semaphore) SetRed() { +func (s *inMemorySemaphore) SetRed() { <-s.green s.green <- false } -// IsGreen reports whether execution is currently allowed. -func (s *Semaphore) IsGreen() bool { +func (s *inMemorySemaphore) IsGreen() bool { val := <-s.green s.green <- val return val } +// PersistedSemaphore implements Semaphore and persists every state change to +// the database via the Store, so the state survives application restarts. +type PersistedSemaphore struct { + inner Semaphore + store *Store + ctx context.Context +} + +// NewPersistedSemaphore creates a PersistedSemaphore with an initial state. +func NewPersistedSemaphore(store *Store, startGreen bool) *PersistedSemaphore { + return &PersistedSemaphore{ + inner: NewSemaphore(startGreen), + store: store, + ctx: context.Background(), + } +} + +// SetGreen allows task execution and persists the state. +func (ps *PersistedSemaphore) SetGreen() { + ps.inner.SetGreen() + _ = ps.store.SetSemaphore(ps.ctx, true) +} + +// SetRed blocks task execution and persists the state. +func (ps *PersistedSemaphore) SetRed() { + ps.inner.SetRed() + _ = ps.store.SetSemaphore(ps.ctx, false) +} + +// IsGreen reports whether execution is currently allowed. +func (ps *PersistedSemaphore) IsGreen() bool { + return ps.inner.IsGreen() +} + // ─── Scheduler ─────────────────────────────────────────────────────────────── // Scheduler polls the database every interval and executes pending tasks. type Scheduler struct { interval time.Duration rateLimit int // max tasks per tick; 0 = unlimited - semaphore *Semaphore + semaphore Semaphore store *Store registry *Registry @@ -328,7 +420,7 @@ type SchedulerConfig struct { } // NewScheduler creates a Scheduler. -func NewScheduler(cfg SchedulerConfig, sem *Semaphore, st *Store, reg *Registry) *Scheduler { +func NewScheduler(cfg SchedulerConfig, sem Semaphore, st *Store, reg *Registry) *Scheduler { return &Scheduler{ interval: cfg.Interval, rateLimit: cfg.RateLimit, @@ -424,23 +516,156 @@ func (s *Scheduler) tick(ctx context.Context, t time.Time) { } } -// execute runs one task and records the result regardless of outcome. +// execute runs a task's handler one or more times within the same tick. +// When Thread=0 (sequential), runs execute one after another in a loop. +// When Thread=1 (concurrent), runs fire simultaneously in goroutines. +// The number of runs per tick is capped by the rate limit, and the item +// is re-queued with remaining runs or deleted when exhausted. func (s *Scheduler) execute(ctx context.Context, item *QueueItem, pos, total int) { - logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] starting task id=%d type=%s", pos, total, item.ID, item.TaskType)) - - startedAt := time.Now() - handler := s.registry.get(item.TaskType) - execErr := handler(ctx, item) - - if execErr != nil { - logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] task id=%d FAILED: %v", pos, total, item.ID, execErr)) - } else { - logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] task id=%d OK (%.0fms)", - pos, total, item.ID, float64(time.Since(startedAt).Milliseconds()))) + // Determine how many runs to fire this tick. + desiredRuns := item.MaxRuns + if desiredRuns < 1 { + desiredRuns = 1 // -1 (infinite): one batch per tick } - if err := s.store.RecordResult(ctx, item, startedAt, execErr); err != nil { - logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] failed to record result for id=%d: %v", pos, total, item.ID, err)) + // Cap by the rate limit so one item cannot exhaust the tick budget. + budget := s.rateLimit + if budget <= 0 { + budget = 1000 + } + if desiredRuns > budget { + desiredRuns = budget + } + + remainingRuns := 0 + if item.MaxRuns == -1 { + remainingRuns = -1 + } else { + remainingRuns = item.MaxRuns - desiredRuns + if remainingRuns < 0 { + remainingRuns = 0 + } + } + + isConcurrent := item.Thread != 0 + if isConcurrent { + logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] task id=%d type=%s firing %d/%d concurrent run(s)", pos, total, item.ID, item.TaskType, desiredRuns, item.MaxRuns)) + } else { + logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] task id=%d type=%s firing %d/%d sequential run(s)", pos, total, item.ID, item.TaskType, desiredRuns, item.MaxRuns)) + } + + if isConcurrent { + // ── Concurrent: all runs in parallel goroutines ── + var wg sync.WaitGroup + wg.Add(desiredRuns) + + for i := 0; i < desiredRuns; i++ { + go func(runIdx int) { + defer wg.Done() + startedAt := time.Now() + handler := s.registry.get(item.TaskType) + execErr := handler(ctx, item) + + finishedAt := time.Now() + elapsed := finishedAt.Sub(startedAt).Milliseconds() + + status := ProcessedOK + errMsg := "" + if execErr != nil { + status = ProcessedError + errMsg = execErr.Error() + } + + processed := &ProcessedItem{ + QueueItemID: item.ID, + TaskType: item.TaskType, + Payload: item.Payload, + Status: status, + ErrorMsg: errMsg, + StartedAt: startedAt, + FinishedAt: finishedAt, + DurationMs: elapsed, + } + if err := s.store.db.WithContext(ctx).Create(processed).Error; err != nil { + logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] failed to record processed_item for id=%d run=%d: %v", pos, total, item.ID, runIdx+1, err)) + } + + if execErr != nil { + logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] task id=%d run=%d FAILED after %dms: %v", pos, total, item.ID, runIdx+1, elapsed, execErr)) + } else { + logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] task id=%d run=%d OK (%dms)", pos, total, item.ID, runIdx+1, elapsed)) + } + }(i) + } + + wg.Wait() + } else { + // ── Sequential: all runs one after another ── + for i := 0; i < desiredRuns; i++ { + startedAt := time.Now() + handler := s.registry.get(item.TaskType) + execErr := handler(ctx, item) + + finishedAt := time.Now() + elapsed := finishedAt.Sub(startedAt).Milliseconds() + + status := ProcessedOK + errMsg := "" + if execErr != nil { + status = ProcessedError + errMsg = execErr.Error() + } + + processed := &ProcessedItem{ + QueueItemID: item.ID, + TaskType: item.TaskType, + Payload: item.Payload, + Status: status, + ErrorMsg: errMsg, + StartedAt: startedAt, + FinishedAt: finishedAt, + DurationMs: elapsed, + } + if err := s.store.db.WithContext(ctx).Create(processed).Error; err != nil { + logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] failed to record processed_item for id=%d run=%d: %v", pos, total, item.ID, i+1, err)) + } + + if execErr != nil { + logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] task id=%d run=%d FAILED after %dms: %v", pos, total, item.ID, i+1, elapsed, execErr)) + } else { + logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] task id=%d run=%d OK (%dms)", pos, total, item.ID, i+1, elapsed)) + } + } + } + + // ── After all runs: decide the queue item's fate ── + if remainingRuns == -1 { + if err := s.store.db.WithContext(ctx). + Model(&QueueItem{}). + Where("id = ?", item.ID). + Updates(map[string]interface{}{ + "status": StatusPending, + "max_runs": -1, + }).Error; err != nil { + logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] failed to re-queue infinite item id=%d: %v", pos, total, item.ID, err)) + } + logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] task id=%d re-queued for next tick (infinite)", pos, total, item.ID)) + } else if remainingRuns > 0 { + if err := s.store.db.WithContext(ctx). + Model(&QueueItem{}). + Where("id = ?", item.ID). + Updates(map[string]interface{}{ + "status": StatusPending, + "max_runs": remainingRuns, + }).Error; err != nil { + logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] failed to re-queue item id=%d with %d remaining runs: %v", pos, total, item.ID, remainingRuns, err)) + } + logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] task id=%d re-queued with %d run(s) remaining", pos, total, item.ID, remainingRuns)) + } else { + if err := s.store.db.WithContext(ctx).Delete(&QueueItem{}, item.ID).Error; err != nil { + logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] failed to delete finished item id=%d: %v", pos, total, item.ID, err)) + } + logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] task id=%d all runs exhausted, deleted", pos, total, item.ID)) } }