BountyBoard/internal/jobs/jobs.go

// Package jobs is a small persisted background job queue (§12): jobs live in
// the `jobs` collection so restarts don't lose work; workers recover from
// panics and retry with exponential backoff up to 3 attempts.
package jobs

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"runtime/debug"
	"time"

	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"

	"bountyboard/internal/metrics"
	"bountyboard/internal/store"
	"bountyboard/internal/ulid"
)

const (
	StatusQueued    = "queued"
	StatusRunning   = "running"
	StatusSucceeded = "succeeded"
	StatusFailed    = "failed"

	maxAttempts = 3
	// a "running" job untouched this long is presumed lost to a crash and
	// re-queued
	staleAfter = 5 * time.Minute
)

type Job struct {
	ID        string    `bson:"_id" json:"id"`
	Kind      string    `bson:"kind" json:"kind"`
	Payload   bson.Raw  `bson:"payload" json:"-"`
	Status    string    `bson:"status" json:"status"`
	Attempts  int       `bson:"attempts" json:"attempts"`
	RunAfter  time.Time `bson:"runAfter" json:"runAfter"`
	LastError string    `bson:"lastError,omitempty" json:"lastError,omitempty"`
	CreatedAt time.Time `bson:"createdAt" json:"createdAt"`
	UpdatedAt time.Time `bson:"updatedAt" json:"updatedAt"`
}

// Handler processes one job. err == nil marks the job succeeded; an error
// schedules a retry (with backoff) until attempts run out — the final
// failure is reported via OnExhausted if the handler set one.
type Handler func(ctx context.Context, job *Job) error

type Runner struct {
	st       *store.Store
	log      *slog.Logger
	reg      *metrics.Registry
	handlers map[string]Handler
	workers  int
}

func NewRunner(st *store.Store, log *slog.Logger, reg *metrics.Registry, workers int) *Runner {
	if workers < 1 {
		workers = 2
	}
	return &Runner{
		st: st, log: log, reg: reg,
		handlers: map[string]Handler{}, workers: workers,
	}
}

func (r *Runner) coll() *mongo.Collection { return r.st.DB.Collection("jobs") }

// Register installs the handler for a job kind (call before Run).
func (r *Runner) Register(kind string, h Handler) { r.handlers[kind] = h }

// Enqueue persists a job for asynchronous execution.
func (r *Runner) Enqueue(ctx context.Context, kind string, payload any) (string, error) {
	raw, err := bson.Marshal(payload)
	if err != nil {
		return "", fmt.Errorf("encode job payload: %w", err)
	}
	now := time.Now().UTC()
	job := Job{
		ID: ulid.New(), Kind: kind, Payload: raw,
		Status: StatusQueued, RunAfter: now, CreatedAt: now, UpdatedAt: now,
	}
	if _, err := r.coll().InsertOne(ctx, job); err != nil {
		return "", fmt.Errorf("enqueue %s: %w", kind, err)
	}
	r.reg.Inc("jobs_enqueued_total", 1)
	return job.ID, nil
}

// DecodePayload unmarshals the job payload into dst.
func DecodePayload(job *Job, dst any) error {
	if err := bson.Unmarshal(job.Payload, dst); err != nil {
		return fmt.Errorf("decode payload of job %s (%s): %w", job.ID, job.Kind, err)
	}
	return nil
}

// Run starts the worker pool and blocks until ctx is canceled.
func (r *Runner) Run(ctx context.Context) {
	for i := 0; i < r.workers; i++ {
		go r.workLoop(ctx)
	}
	// stale-job janitor
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			r.requeueStale(ctx)
			r.updateGauges(ctx)
		}
	}
}

func (r *Runner) workLoop(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		job, err := r.claim(ctx)
		if err != nil {
			if !errors.Is(err, context.Canceled) {
				r.log.Error("claim job", "err", err)
			}
			sleep(ctx, 2*time.Second)
			continue
		}
		if job == nil {
			sleep(ctx, time.Second)
			continue
		}
		r.execute(ctx, job)
	}
}

func sleep(ctx context.Context, d time.Duration) {
	select {
	case <-ctx.Done():
	case <-time.After(d):
	}
}

// claim atomically moves one due job to running.
func (r *Runner) claim(ctx context.Context) (*Job, error) {
	now := time.Now().UTC()
	var job Job
	err := r.coll().FindOneAndUpdate(ctx,
		bson.M{"status": StatusQueued, "runAfter": bson.M{"$lte": now}},
		bson.M{"$set": bson.M{"status": StatusRunning, "updatedAt": now},
			"$inc": bson.M{"attempts": 1}},
	).Decode(&job)
	if errors.Is(err, mongo.ErrNoDocuments) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	job.Attempts++ // FindOneAndUpdate returned the pre-update doc
	job.Status = StatusRunning
	return &job, nil
}

func (r *Runner) execute(ctx context.Context, job *Job) {
	handler, ok := r.handlers[job.Kind]
	if !ok {
		r.finish(ctx, job, fmt.Errorf("no handler registered for kind %q", job.Kind), false)
		return
	}
	var err error
	func() {
		defer func() {
			if rec := recover(); rec != nil {
				err = fmt.Errorf("panic: %v", rec)
				r.log.Error("job panic", "job", job.ID, "kind", job.Kind,
					"panic", fmt.Sprint(rec), "stack", string(debug.Stack()))
			}
		}()
		jctx, cancel := context.WithTimeout(ctx, 10*time.Minute)
		defer cancel()
		err = handler(jctx, job)
	}()
	r.finish(ctx, job, err, true)
}

func (r *Runner) finish(ctx context.Context, job *Job, err error, retryable bool) {
	// Bookkeeping must survive shutdown/cancellation: otherwise a canceled
	// run leaves the job stuck in "running" until stale recovery.
	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
	defer cancel()
	now := time.Now().UTC()
	switch {
	case err == nil:
		_, uerr := r.coll().UpdateOne(ctx, bson.M{"_id": job.ID}, bson.M{
			"$set": bson.M{"status": StatusSucceeded, "lastError": "", "updatedAt": now}})
		if uerr != nil {
			r.log.Error("mark job succeeded", "job", job.ID, "err", uerr)
		}
		r.reg.Inc("jobs_succeeded_total", 1)
	case retryable && job.Attempts < maxAttempts:
		backoff := time.Duration(1<<job.Attempts) * 5 * time.Second // 10s, 20s
		_, uerr := r.coll().UpdateOne(ctx, bson.M{"_id": job.ID}, bson.M{
			"$set": bson.M{"status": StatusQueued, "lastError": err.Error(),
				"runAfter": now.Add(backoff), "updatedAt": now}})
		if uerr != nil {
			r.log.Error("requeue job", "job", job.ID, "err", uerr)
		}
		r.log.Warn("job failed, will retry", "job", job.ID, "kind", job.Kind,
			"attempt", job.Attempts, "err", err)
		r.reg.Inc("jobs_retried_total", 1)
	default:
		_, uerr := r.coll().UpdateOne(ctx, bson.M{"_id": job.ID}, bson.M{
			"$set": bson.M{"status": StatusFailed, "lastError": err.Error(), "updatedAt": now}})
		if uerr != nil {
			r.log.Error("mark job failed", "job", job.ID, "err", uerr)
		}
		r.log.Error("job exhausted retries", "job", job.ID, "kind", job.Kind, "err", err)
		r.reg.Inc("jobs_failed_total", 1)
	}
}

// IsFinalAttempt lets handlers run compensation logic (revert status, notify)
// only when no more retries will happen.
func IsFinalAttempt(job *Job) bool { return job.Attempts >= maxAttempts }

func (r *Runner) requeueStale(ctx context.Context) {
	cutoff := time.Now().UTC().Add(-staleAfter)
	res, err := r.coll().UpdateMany(ctx,
		bson.M{"status": StatusRunning, "updatedAt": bson.M{"$lt": cutoff}},
		bson.M{"$set": bson.M{"status": StatusQueued, "updatedAt": time.Now().UTC()}})
	if err != nil {
		r.log.Error("requeue stale jobs", "err", err)
		return
	}
	if res.ModifiedCount > 0 {
		r.log.Warn("requeued stale running jobs", "count", res.ModifiedCount)
	}
}

func (r *Runner) updateGauges(ctx context.Context) {
	for _, status := range []string{StatusQueued, StatusRunning, StatusFailed} {
		n, err := r.coll().CountDocuments(ctx, bson.M{"status": status})
		if err == nil {
			r.reg.Set("jobs_"+status, n)
		}
	}
}

// Statuses feeds the admin service-status panel.
func (r *Runner) Statuses() any {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	out := map[string]int64{}
	for _, status := range []string{StatusQueued, StatusRunning, StatusSucceeded, StatusFailed} {
		n, err := r.coll().CountDocuments(ctx, bson.M{"status": status})
		if err != nil {
			return map[string]string{"error": err.Error()}
		}
		out[status] = n
	}
	return out
}