sift/internal/batch/batch.go

// Package batch is the batch runner core: it owns job lifecycle, idempotency and
// the per-caller / global caps. It knows nothing about HTTP, the guard or the real
// fetch pipeline — per-target work is an injected FetchFunc, and persistence is the
// dcache store. Jobs never resume after a restart, so the caller is held in memory,
// never written to the store.
package batch

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/auth"
	"git.wntrmute.dev/kyle/sift/internal/dcache"
)

type Options struct {
	Mode     string `json:"mode,omitempty"`
	Question string `json:"question,omitempty"`
	Pages    string `json:"pages,omitempty"`
	Grep     string `json:"grep,omitempty"`
	JQ       string `json:"jq,omitempty"`
	Context  *bool  `json:"context,omitempty"` // nil = /v1/fetch's default (on when there is a question)
	Recover  bool   `json:"recover,omitempty"`
}

type Target struct {
	URL     string  `json:"url"`
	Options Options `json:"options"`
}

// Result is a finished target: the /v1/fetch envelope's content, and every other envelope field
// in Meta. It is stored JSON-encoded in the target's result column.
type Result struct {
	Content string         `json:"content"`
	Meta    map[string]any `json:"meta,omitempty"`
}

// TargetError is a failure with an agent-facing hint. Status is the HTTP status the same
// failure gets from /v1/fetch; the runner ignores it (the metrics hook reads it). Any other
// error is stored with its Error() text and no hint.
type TargetError struct {
	Status    int
	Msg, Hint string
}

func (e *TargetError) Error() string { return e.Msg }

// FetchFunc does one target for the caller that submitted the job.
type FetchFunc func(ctx context.Context, c *auth.Caller, t Target) (Result, error)

type Job struct {
	ID           string
	Caller       string
	Status       string // queued | running | done | partial | failed | cancelled
	Error        string
	TargetsTotal int
	TargetsDone  int
	CreatedAt    int64 // unix seconds
}

type TargetStatus struct {
	Idx    int
	URL    string
	Status string // pending | running | done | failed
	Result *Result
	Error  string
	Hint   string
}

// Hooks observe the runner for metrics. Any nil field is skipped.
type Hooks struct {
	Job        func(outcome string)                              // created, reused, done, partial, failed, cancelled
	Target     func(err error, d time.Duration, resultBytes int) // once per stored target; err nil on success
	JobSeconds func(d time.Duration)                             // submit → terminal, once per job
	Depth      func(live int)                                    // live jobs, after every change
	Inflight   func(n int)                                       // targets being fetched, after every change
}

type Config struct {
	TTL           time.Duration // 0 → 24h
	TargetsPerJob int           // 0 → 32
	JobsPerCaller int           // 0 → 1
	JobsTotal     int           // 0 → 8
	Workers       int           // 0 → 1; values above 2 are clamped to 2
	MaxBytes      int64         // 0 → 64 MiB
	RetryWait     time.Duration // 0 → 2s: the pause before a busy target is tried again
	RetryFor      time.Duration // 0 → 10m: how long one target may keep coming back busy
	Hooks         Hooks
}

var (
	ErrNoTargets      = errors.New("no targets")
	ErrTooManyTargets = errors.New("too many targets")
	ErrCallerBusy     = errors.New("busy")
	ErrBusy           = errors.New("jobs busy")
	ErrRetryLater     = errors.New("retry later") // a FetchFunc wraps this when the target was only busy
)

// terminal statuses mark a job as finished; only a queued job can be cancelled.
func isJobTerminal(status string) bool {
	switch status {
	case "done", "partial", "failed", "cancelled":
		return true
	}
	return false
}

type Runner struct {
	db        *dcache.DB
	fetch     FetchFunc
	log       *slog.Logger
	cfg       Config
	callers   map[string]*auth.Caller // job id → submitting caller, held in memory only
	callMu    sync.Mutex
	closeOnce sync.Once

	mu       sync.Mutex // guards everything below
	cond     *sync.Cond // on mu
	live     []*liveJob // the rotation, in submit order
	cursor   int        // absolute rotation position; wraps at access, so it survives jobs being added
	closed   bool
	ctx      context.Context // the FetchFunc context; Close cancels it
	stop     context.CancelFunc
	wg       sync.WaitGroup // the workers
	inflight int
}

func New(db *dcache.DB, fetch FetchFunc, cfg Config, log *slog.Logger) *Runner {
	if cfg.TTL == 0 {
		cfg.TTL = 24 * time.Hour
	}
	if cfg.TargetsPerJob == 0 {
		cfg.TargetsPerJob = 32
	}
	if cfg.JobsPerCaller == 0 {
		cfg.JobsPerCaller = 1
	}
	if cfg.JobsTotal == 0 {
		cfg.JobsTotal = 8
	}
	if cfg.Workers == 0 {
		cfg.Workers = 1
	}
	if cfg.Workers > 2 {
		cfg.Workers = 2
	}
	if cfg.MaxBytes == 0 {
		cfg.MaxBytes = 64 << 20
	}
	if cfg.RetryWait == 0 {
		cfg.RetryWait = 2 * time.Second
	}
	if cfg.RetryFor == 0 {
		cfg.RetryFor = 10 * time.Minute
	}
	ctx, stop := context.WithCancel(context.Background())
	r := &Runner{
		db:      db,
		fetch:   fetch,
		log:     log,
		cfg:     cfg,
		callers: make(map[string]*auth.Caller),
		ctx:     ctx,
		stop:    stop,
	}
	r.cond = sync.NewCond(&r.mu)
	r.wg.Add(cfg.Workers)
	for i := 0; i < cfg.Workers; i++ {
		go r.worker()
	}
	return r
}

// Submit stores the job and its targets before it returns. The order matters: an empty
// request, then idempotency (which never touches the caps), then the caps, then the write.
func (r *Runner) Submit(ctx context.Context, c *auth.Caller, callerClass string, targets []Target) (id string, reused bool, err error) {
	if len(targets) == 0 {
		return "", false, ErrNoTargets
	}

	raw, err := json.Marshal(targets)
	if err != nil {
		return "", false, fmt.Errorf("batch: marshal targets: %w", err)
	}
	sum := sha256.Sum256([]byte(c.Name + "\x00" + string(raw)))
	idemKey := hex.EncodeToString(sum[:])

	now := time.Now().Unix()
	if existing, ok := r.db.BatchJobByIDem(c.Name, idemKey, now, r.cfg.TTL); ok {
		r.callJob("reused")
		return existing, true, nil
	}

	if len(targets) > r.cfg.TargetsPerJob {
		return "", false, ErrTooManyTargets
	}
	if r.db.LiveBatchJobs(c.Name) >= r.cfg.JobsPerCaller {
		return "", false, ErrCallerBusy
	}
	if r.db.LiveBatchJobsTotal() >= r.cfg.JobsTotal {
		return "", false, ErrBusy
	}

	id, err = randomID()
	if err != nil {
		return "", false, fmt.Errorf("batch: job id: %w", err)
	}

	if err := r.db.CreateBatchJob(dcache.BatchJob{
		ID:           id,
		Caller:       c.Name,
		CallerClass:  callerClass,
		IDemKey:      idemKey,
		Status:       "queued",
		TargetsTotal: len(targets),
		CreatedAt:    now,
		Lease:        now,
	}); err != nil {
		r.log.Warn("batch: create job", "job", id, "err", err)
		return "", false, fmt.Errorf("batch: create job: %w", err)
	}

	btargets := make([]dcache.BatchTarget, len(targets))
	for i := range targets {
		opts, _ := json.Marshal(targets[i].Options)
		btargets[i] = dcache.BatchTarget{
			JobID:   id,
			Idx:     i,
			URL:     targets[i].URL,
			Options: string(opts),
			Status:  "pending",
		}
	}
	if err := r.db.PutBatchTargets(id, btargets); err != nil {
		r.log.Warn("batch: store targets", "job", id, "err", err)
		return "", false, fmt.Errorf("batch: store targets: %w", err)
	}

	r.callMu.Lock()
	r.callers[id] = c
	r.callMu.Unlock()

	r.mu.Lock()
	r.live = append(r.live, &liveJob{
		id:      id,
		caller:  c,
		targets: targets,
		created: time.Unix(now, 0),
	})
	r.callDepth(len(r.live))
	r.cond.Broadcast()
	r.mu.Unlock()

	r.callJob("created")
	r.db.SweepBatch(now, r.cfg.TTL)
	return id, false, nil
}

// Cancel cancels a job. A job still in the rotation (queued or with a target in flight) is marked
// cancelled at once: its in-flight target, if any, finishes and is stored without overwriting
// cancelled, and no further target starts. A job not in the rotation (unknown or already terminal)
// keeps task 02's behaviour.
func (r *Runner) Cancel(id string) bool {
	r.mu.Lock()
	var j *liveJob
	for _, lj := range r.live {
		if lj.id == id {
			j = lj
			break
		}
	}
	if j != nil {
		j.cancelled = true
		if err := r.db.SetBatchJobStatus(id, "cancelled", "", time.Now().Unix()); err != nil {
			r.log.Warn("batch: cancel job", "job", id, "err", err)
		}
		r.callJob("cancelled")
		if !j.running {
			r.removeLocked(j)
			r.forgetCaller(id)
			r.callDepth(len(r.live))
		}
		r.cond.Broadcast()
		r.mu.Unlock()
		return true
	}
	r.mu.Unlock()

	jb, ok := r.db.BatchJob(id)
	if !ok || isJobTerminal(jb.Status) {
		return false
	}
	if err := r.db.SetBatchJobStatus(id, "cancelled", "", time.Now().Unix()); err != nil {
		r.log.Warn("batch: cancel job", "job", id, "err", err)
		return false
	}
	r.callMu.Lock()
	delete(r.callers, id)
	r.callMu.Unlock()
	r.callJob("cancelled")
	return true
}

// Status maps the store rows for a job and its targets; Result is decoded from the result
// column. It writes nothing.
func (r *Runner) Status(id string) (Job, []TargetStatus, bool) {
	jb, ok := r.db.BatchJob(id)
	if !ok {
		return Job{}, nil, false
	}
	j := Job{
		ID:           jb.ID,
		Caller:       jb.Caller,
		Status:       jb.Status,
		Error:        jb.Error,
		TargetsTotal: jb.TargetsTotal,
		TargetsDone:  jb.TargetsDone,
		CreatedAt:    jb.CreatedAt,
	}
	tbs, _ := r.db.BatchTargets(id)
	ts := make([]TargetStatus, 0, len(tbs))
	for _, bt := range tbs {
		var res *Result
		if bt.Result != nil {
			res = new(Result)
			if err := json.Unmarshal([]byte(*bt.Result), res); err != nil {
				r.log.Warn("batch: decode result", "job", id, "err", err)
			}
		}
		ts = append(ts, TargetStatus{
			Idx:    bt.Idx,
			URL:    bt.URL,
			Status: bt.Status,
			Result: res,
			Error:  strOrEmpty(bt.Error),
			Hint:   strOrEmpty(bt.Hint),
		})
	}
	return j, ts, true
}

// Touch bumps the job's lease to now so a slow owner does not lose it.
func (r *Runner) Touch(id string) {
	now := time.Now().Unix()
	if err := r.db.TouchBatchJob(id, now); err != nil {
		r.log.Warn("batch: touch job", "job", id, "err", err)
	}
}

// Close stops the workers, cancels their context, and waits for an in-flight fetch to return.
// Safe to call more than once. The in-memory callers are dropped.
func (r *Runner) Close() {
	r.closeOnce.Do(func() {
		r.callMu.Lock()
		r.callers = nil
		r.callMu.Unlock()

		r.mu.Lock()
		r.closed = true
		r.stop()
		r.cond.Broadcast()
		r.mu.Unlock()
		r.wg.Wait()
	})
}

func (r *Runner) callJob(outcome string) {
	if r.cfg.Hooks.Job != nil {
		r.cfg.Hooks.Job(outcome)
	}
}

func randomID() (string, error) {
	var b [16]byte
	if _, err := rand.Read(b[:]); err != nil {
		return "", err
	}
	return hex.EncodeToString(b[:]), nil
}

func strOrEmpty(s *string) string {
	if s == nil {
		return ""
	}
	return *s
}