scheduler updates

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)
 }
 

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 {
+func NewSemaphore(startGreen bool) Semaphore {
-	s := &Semaphore{green: make(chan bool, 1)}
+	s := &inMemorySemaphore{green: make(chan bool, 1)}
 	s.green <- startGreen
 	return s
 }
 
-// SetGreen allows task execution to proceed.
+func (s *inMemorySemaphore) SetGreen() {
-func (s *Semaphore) SetGreen() {
 	<-s.green
 	s.green <- true
 }
 
-// SetRed blocks task execution.
+func (s *inMemorySemaphore) SetRed() {
-func (s *Semaphore) SetRed() {
 	<-s.green
 	s.green <- false
 }
 
-// IsGreen reports whether execution is currently allowed.
+func (s *inMemorySemaphore) IsGreen() bool {
-func (s *Semaphore) 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))
+	// Determine how many runs to fire this tick.
-
+	desiredRuns := item.MaxRuns
-	startedAt := time.Now()
+	if desiredRuns < 1 {
-	handler := s.registry.get(item.TaskType)
+		desiredRuns = 1 // -1 (infinite): one batch per tick
-	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())))
 	}
 
-	if err := s.store.RecordResult(ctx, item, startedAt, execErr); err != nil {
+	// Cap by the rate limit so one item cannot exhaust the tick budget.
-		logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] failed to record result for id=%d: %v", pos, total, item.ID, err))
+	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))
 	}
 }