diff --git a/.env.example b/.env.example index 899de8f..f0d7af6 100644 --- a/.env.example +++ b/.env.example @@ -2,3 +2,8 @@ API_ADDR=:8080 DATABASE_URL=postgres://dtq:dtq@localhost:5432/dtq?sslmode=disable REDIS_URL=redis://localhost:6379/0 + +# How long a 'queued' job (with no last_error) is allowed to sit before the +# sweeper re-pushes it to its Redis stage queue. Recovers from BLPOP-to-claim +# races where a worker crashed mid-claim. Default 60 seconds. +STALE_QUEUED_THRESHOLD_SEC=60 diff --git a/cmd/sweeper/main.go b/cmd/sweeper/main.go index 7825331..81aa49c 100644 --- a/cmd/sweeper/main.go +++ b/cmd/sweeper/main.go @@ -5,6 +5,7 @@ import ( "log/slog" "os" "os/signal" + "strconv" "syscall" "time" @@ -46,10 +47,18 @@ func main() { redis := goredis.NewClient(opts) defer func() { _ = redis.Close() }() + staleThresholdSec := 60 + if v := os.Getenv("STALE_QUEUED_THRESHOLD_SEC"); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + staleThresholdSec = n + } + } + sw := sweeper.New(sweeper.Config{ - Store: store.New(pool), - Queue: queue.New(redis), - Interval: 5 * time.Second, + Store: store.New(pool), + Queue: queue.New(redis), + Interval: 5 * time.Second, + StaleThreshold: time.Duration(staleThresholdSec) * time.Second, }) slog.Info("sweeper starting") diff --git a/internal/store/store.go b/internal/store/store.go index 80c2ba2..bc32cfb 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -165,6 +165,25 @@ func (s *Store) AdvanceJob(ctx context.Context, id uuid.UUID, nextStage string) return nil } +func (s *Store) PoolForTest() *pgxpool.Pool { return s.pool } + +func (s *Store) ListStaleQueuedJobs(ctx context.Context, threshold time.Duration) ([]Job, error) { + const q = ` + SELECT id, user_id, gmail_message_id, stage, status, payload, is_synthetic, + worker_id, attempts, max_attempts, last_error, next_run_at, + claimed_at, completed_at, created_at, updated_at + FROM pipeline_jobs + WHERE status = $1 AND last_error IS NULL + AND created_at < now() - make_interval(secs => $2)` + + rows, err := s.pool.Query(ctx, q, StatusQueued, threshold.Seconds()) + if err != nil { + return nil, fmt.Errorf("list stale queued: %w", err) + } + defer rows.Close() + return scanJobs(rows) +} + func (s *Store) GetJob(ctx context.Context, id uuid.UUID) (Job, error) { const q = ` SELECT id, user_id, gmail_message_id, stage, status, payload, is_synthetic, diff --git a/internal/store/store_integration_test.go b/internal/store/store_integration_test.go index 5de723d..c81ce06 100644 --- a/internal/store/store_integration_test.go +++ b/internal/store/store_integration_test.go @@ -315,6 +315,43 @@ func TestAdvanceJob_TransitionsToNextStage(t *testing.T) { } } +func TestListStaleQueuedJobs_ReturnsOldQueuedWithoutLastError(t *testing.T) { + s := newStore(t) + ctx := context.Background() + + jStale, _ := s.EnqueueJob(ctx, store.NewJob{Stage: "fetch"}) + if _, err := s.GetJob(ctx, jStale.ID); err != nil { + t.Fatal(err) + } + if _, err := s.PoolForTest().Exec(ctx, //nolint:gosec + `UPDATE pipeline_jobs SET created_at = now() - interval '5 minutes' WHERE id = $1`, + jStale.ID, + ); err != nil { + t.Fatal(err) + } + + _, _ = s.EnqueueJob(ctx, store.NewJob{Stage: "fetch"}) + + jRetry, _ := s.EnqueueJob(ctx, store.NewJob{Stage: "fetch"}) + _ = s.ClaimJob(ctx, jRetry.ID, "w1") + _ = s.MarkFailed(ctx, jRetry.ID, "boom", time.Now().Add(-1*time.Hour)) + _, _ = s.PoolForTest().Exec(ctx, //nolint:gosec + `UPDATE pipeline_jobs SET created_at = now() - interval '5 minutes' WHERE id = $1`, + jRetry.ID, + ) + + jobs, err := s.ListStaleQueuedJobs(ctx, 1*time.Minute) + if err != nil { + t.Fatal(err) + } + if len(jobs) != 1 { + t.Fatalf("got %d, want 1; got: %v", len(jobs), jobIDs(jobs)) + } + if jobs[0].ID != jStale.ID { + t.Fatalf("got %s, want %s", jobs[0].ID, jStale.ID) + } +} + func jobIDs(js []store.Job) []string { out := make([]string, len(js)) for i, j := range js { diff --git a/internal/sweeper/sweeper.go b/internal/sweeper/sweeper.go index e4a5d04..a45d491 100644 --- a/internal/sweeper/sweeper.go +++ b/internal/sweeper/sweeper.go @@ -11,9 +11,10 @@ import ( ) type Config struct { - Store *store.Store - Queue *queue.Queue - Interval time.Duration + Store *store.Store + Queue *queue.Queue + Interval time.Duration + StaleThreshold time.Duration } type Sweeper struct { @@ -24,6 +25,9 @@ func New(cfg Config) *Sweeper { if cfg.Interval == 0 { cfg.Interval = 5 * time.Second } + if cfg.StaleThreshold == 0 { + cfg.StaleThreshold = 60 * time.Second + } return &Sweeper{cfg: cfg} } @@ -50,6 +54,9 @@ func (s *Sweeper) SweepOnce(ctx context.Context) error { if err := s.promoteDelayed(ctx); err != nil { return fmt.Errorf("promote delayed: %w", err) } + if err := s.repushStale(ctx); err != nil { + return fmt.Errorf("repush stale: %w", err) + } return nil } @@ -79,6 +86,21 @@ func (s *Sweeper) reviveOrphans(ctx context.Context) error { return nil } +func (s *Sweeper) repushStale(ctx context.Context) error { + stale, err := s.cfg.Store.ListStaleQueuedJobs(ctx, s.cfg.StaleThreshold) + if err != nil { + return err + } + for _, j := range stale { + if err := s.cfg.Queue.Push(ctx, j.Stage, j.ID.String()); err != nil { + slog.Warn("repush stale failed", "job_id", j.ID, "err", err) + continue + } + slog.Info("repushed stale queued job", "job_id", j.ID, "stage", j.Stage) + } + return nil +} + func (s *Sweeper) promoteDelayed(ctx context.Context) error { ready, err := s.cfg.Store.ListReadyRetryJobs(ctx) if err != nil { diff --git a/internal/sweeper/sweeper_integration_test.go b/internal/sweeper/sweeper_integration_test.go index fdd2d04..a1dc01b 100644 --- a/internal/sweeper/sweeper_integration_test.go +++ b/internal/sweeper/sweeper_integration_test.go @@ -101,6 +101,51 @@ func TestSweepOnce_PromotesDelayedRetryToRedisQueue(t *testing.T) { } } +func TestSweepOnce_RepushesStaleQueuedJob(t *testing.T) { + s, q, sw := setup(t) + ctx := context.Background() + + job, _ := s.EnqueueJob(ctx, store.NewJob{Stage: "fetch"}) + _, _ = s.PoolForTest().Exec(ctx, //nolint:gosec + `UPDATE pipeline_jobs SET created_at = now() - interval '5 minutes' WHERE id = $1`, + job.ID, + ) + + depthBefore, _ := q.Depth(ctx, "fetch") + if err := sw.SweepOnce(ctx); err != nil { + t.Fatalf("sweep: %v", err) + } + depthAfter, _ := q.Depth(ctx, "fetch") + + if depthAfter != depthBefore+1 { + t.Fatalf("depth: before=%d after=%d, want +1", depthBefore, depthAfter) + } + popped, err := q.BlockingPop(ctx, "fetch", 1*time.Second) + if err != nil { + t.Fatalf("pop: %v", err) + } + if popped != job.ID.String() { + t.Fatalf("popped %q, want %q", popped, job.ID.String()) + } +} + +func TestSweepOnce_LeavesFreshlyQueuedJobAlone(t *testing.T) { + s, q, sw := setup(t) + ctx := context.Background() + + _, _ = s.EnqueueJob(ctx, store.NewJob{Stage: "fetch"}) + depthBefore, _ := q.Depth(ctx, "fetch") + + if err := sw.SweepOnce(ctx); err != nil { + t.Fatalf("sweep: %v", err) + } + depthAfter, _ := q.Depth(ctx, "fetch") + + if depthAfter != depthBefore { + t.Fatalf("fresh job got re-pushed: before=%d after=%d", depthBefore, depthAfter) + } +} + func TestSweepOnce_LeavesFutureRetryInPostgresOnly(t *testing.T) { s, q, sw := setup(t) ctx := context.Background()