scheduler updates
This commit is contained in:
parent
d156e3785d
commit
987e8b98d0
2 changed files with 274 additions and 39 deletions
36
scheduler.go
36
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)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue