scheduler updates

This commit is contained in:
Jose Cely 2026-06-23 18:12:36 -05:00
parent d156e3785d
commit 987e8b98d0
2 changed files with 274 additions and 39 deletions

View file

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

View file

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