Changeset 0.32.0 (#206)

server/scheduler/scheduler.go

@@ -5,16 +5,22 @@
 // v0.27.2: Adds global config checks (enabled, max_concurrent), full cron
 //          parsing via robfig/cron/v3, and completion invocation via executor.
 // v0.28.0: Adopt queued runs (webhook triggers), action tasks, C3/C4 audit fixes.
+// v0.32.0: SKIP LOCKED atomic claim — every replica polls, PG serializes.
+//          Replaces ListDue with ClaimDueTask, CreateRunExclusive for
+//          belt-and-suspenders uniqueness. Startup jitter staggers replicas.
 package scheduler
 
 import (
 	"context"
+	"database/sql"
+	"errors"
 	"log"
+	"math/rand"
 	"time"
 
-	"git.gobha.me/xcaliber/chat-switchboard/models"
-	"git.gobha.me/xcaliber/chat-switchboard/store"
-	"git.gobha.me/xcaliber/chat-switchboard/taskutil"
+	"chat-switchboard/models"
+	"chat-switchboard/store"
+	"chat-switchboard/taskutil"
 )
 
 // Scheduler polls for due tasks and dispatches execution.
@@ -45,7 +51,12 @@ func (s *Scheduler) Run() {
 		return
 	}
 	s.running = true
-	log.Printf("[scheduler] Started (poll interval: %s)", s.interval)
+
+	// v0.32.0: Startup jitter — stagger replica polling to reduce lock contention.
+	// Not strictly necessary with SKIP LOCKED but reduces unnecessary work.
+	jitter := time.Duration(rand.Intn(15000)) * time.Millisecond
+	time.Sleep(jitter)
+	log.Printf("[scheduler] Started (jitter=%s, interval=%s)", jitter, s.interval)
 
 	ticker := time.NewTicker(s.interval)
 	defer ticker.Stop()
@@ -69,42 +80,53 @@ func (s *Scheduler) Stop() {
 	}
 }
 
-// poll checks for due tasks and dispatches them.
+// poll claims due tasks one at a time and dispatches them.
+// v0.32.0: Each call to ClaimDueTask atomically locks and claims one task
+// via FOR UPDATE SKIP LOCKED (PG). Multiple replicas poll concurrently;
+// PG ensures each task is handed to exactly one replica.
 func (s *Scheduler) poll() {
 	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
 	defer cancel()
 
-	// v0.27.2: Check global config — tasks may be disabled at runtime.
+	// v0.32.0: Piggyback cleanup of expired tickets and stale rate limit
+	// counters on the scheduler tick. Cheap no-ops when tables are empty.
+	if s.stores.Tickets != nil {
+		if n, err := s.stores.Tickets.Reap(ctx); err == nil && n > 0 {
+			log.Printf("[scheduler] Reaped %d expired WS tickets", n)
+		}
+	}
+	if s.stores.RateLimits != nil {
+		_ = s.stores.RateLimits.Cleanup(ctx, 5*time.Minute)
+	}
+
+	// Check global config — tasks may be disabled at runtime.
 	cfg := taskutil.LoadTaskConfig(ctx, s.stores.GlobalConfig)
 	if !cfg.Enabled {
 		return
 	}
 
-	due, err := s.stores.Tasks.ListDue(ctx, cfg.MaxConcurrent)
-	if err != nil {
-		log.Printf("[scheduler] Failed to list due tasks: %v", err)
-		return
-	}
-
-	for _, task := range due {
-		go s.execute(ctx, task)
+	// Claim tasks one at a time until none remain or max_concurrent reached.
+	claimed := 0
+	for claimed < cfg.MaxConcurrent {
+		task, err := s.stores.Tasks.ClaimDueTask(ctx)
+		if err != nil {
+			// sql.ErrNoRows = nothing due; any other error = log and stop.
+			if !errors.Is(err, sql.ErrNoRows) {
+				log.Printf("[scheduler] ClaimDueTask error: %v", err)
+			}
+			break
+		}
+		claimed++
+		go s.execute(task)
 	}
 }
 
-// execute runs a single task.
-func (s *Scheduler) execute(parentCtx context.Context, task models.Task) {
-	ctx, cancel := context.WithTimeout(parentCtx, time.Duration(task.MaxWallClock)*time.Second)
+// execute runs a single task. The task has already been claimed
+// (next_run_at set to NULL by ClaimDueTask).
+func (s *Scheduler) execute(task *models.Task) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(task.MaxWallClock)*time.Second)
 	defer cancel()
 
-	// Skip if previous run still active
-	active, _ := s.stores.Tasks.GetActiveRun(ctx, task.ID)
-	if active != nil {
-		log.Printf("[scheduler] Skipping task %s (%s) — previous run still active", task.ID, task.Name)
-		// Still advance next_run_at to avoid re-polling the same task
-		s.advanceNextRun(ctx, task)
-		return
-	}
-
 	// v0.28.0: Check for queued run (from webhook trigger) — adopt it instead
 	// of creating a new one so the trigger_payload is preserved.
 	run, _ := s.stores.Tasks.GetQueuedRun(ctx, task.ID)
@@ -114,29 +136,32 @@ func (s *Scheduler) execute(parentCtx context.Context, task models.Task) {
 		run.Status = "running"
 		log.Printf("[scheduler] Adopting queued run %s for task %s (%s)", run.ID, task.ID, task.Name)
 	} else {
-		// Normal cron-triggered execution — create a new run
-		run = &models.TaskRun{
-			TaskID: task.ID,
-			Status: "running",
-		}
-		if err := s.stores.Tasks.CreateRun(ctx, run); err != nil {
-			log.Printf("[scheduler] Failed to create run for task %s: %v", task.ID, err)
-			s.advanceNextRun(ctx, task)
+		// v0.32.0: Conditional insert — prevents double-execution if another
+		// replica somehow also processes this task (belt-and-suspenders).
+		var err error
+		run, err = s.stores.Tasks.CreateRunExclusive(ctx, task.ID)
+		if err != nil {
+			if errors.Is(err, sql.ErrNoRows) {
+				log.Printf("[scheduler] Skipping task %s (%s) — run already exists", task.ID, task.Name)
+			} else {
+				log.Printf("[scheduler] Failed to create run for task %s: %v", task.ID, err)
+			}
+			s.advanceNextRun(ctx, *task)
 			return
 		}
 	}
 
 	log.Printf("[scheduler] Executing task %s (%s) type=%s", task.ID, task.Name, task.TaskType)
 
-	// Mark execution start (C4 fix: SetLastRun is now separate from SetNextRun)
+	// Mark execution start
 	_ = s.stores.Tasks.SetLastRun(ctx, task.ID)
 
 	// Create or reuse service channel
-	channelID, err := s.ensureServiceChannel(ctx, task)
+	channelID, err := s.ensureServiceChannel(ctx, *task)
 	if err != nil {
 		log.Printf("[scheduler] Failed to create service channel for task %s: %v", task.ID, err)
 		_ = s.stores.Tasks.UpdateRun(ctx, run.ID, "failed", 0, 0, 0, "channel creation failed: "+err.Error())
-		s.advanceNextRun(ctx, task)
+		s.advanceNextRun(ctx, *task)
 		return
 	}
 
@@ -154,9 +179,9 @@ func (s *Scheduler) execute(parentCtx context.Context, task models.Task) {
 		})
 	}
 
-	// v0.27.2: Invoke completion via executor
+	// Invoke completion via executor
 	if s.executor != nil {
-		s.executor.Execute(ctx, task, run, channelID)
+		s.executor.Execute(ctx, *task, run, channelID)
 	} else {
 		// No executor — mark completed (channel + prompt persisted)
 		_ = s.stores.Tasks.UpdateRun(ctx, run.ID, "completed", 0, 0, 0, "")
@@ -165,7 +190,7 @@ func (s *Scheduler) execute(parentCtx context.Context, task models.Task) {
 
 	log.Printf("[scheduler] Task %s finished (channel %s)", task.ID, channelID)
 
-	s.advanceNextRun(ctx, task)
+	s.advanceNextRun(ctx, *task)
 }
 
 // ensureServiceChannel creates a new service channel or reuses an existing one.
@@ -213,7 +238,7 @@ func (s *Scheduler) advanceNextRun(ctx context.Context, task models.Task) {
 		return
 	}
 
-	// v0.27.2: Full cron parsing via robfig/cron/v3 (replaces hand-rolled parser).
+	// Full cron parsing via robfig/cron/v3.
 	next := taskutil.NextRunFromSchedule(task.Schedule, task.Timezone)
 	if next == nil {
 		log.Printf("[scheduler] Failed to compute next run for task %s (schedule: %q) — deactivating", task.ID, task.Schedule)