sift/internal/batch/worker.go

package batch

import (
	"context"
	"encoding/json"
	"errors"
	"time"

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

// ErrInternal is stored for a target whose fetch panicked; the panic value is never echoed.
var ErrInternal = errors.New("internal error")

// liveJob is a submitted job that the workers are draining. It lives in Runner.live until the
// last target is stored or the job is cancelled.
type liveJob struct {
	id           string
	caller       *auth.Caller
	targets      []Target // as submitted
	created      time.Time
	next         int  // index of the next target to fetch
	running      bool // a target of this job is in flight
	started      bool // set to "running" in the store already
	cancelled    bool
	retryAt      time.Time // zero, or when the busy target may be tried again
	busySince    time.Time // zero, or when the current target first came back busy
	done, failed int
}

// worker drains the rotation: it takes the next live job with a pending target and no target in
// flight, fetches and stores that one target, and moves on. Each job is serial — its target i+1
// never starts before i is stored.
func (r *Runner) worker() {
	defer r.wg.Done()

	for {
		r.mu.Lock()
		if r.closed { // stop taking targets at once; live jobs are reaped on the next start
			r.mu.Unlock()
			return
		}
		now := time.Now()
		j := r.pick(now)
		if j == nil {
			r.scheduleRetryLocked(now)
			r.cond.Wait() // re-locks on wake; release before the next top-of-loop Lock
			r.mu.Unlock()
			continue
		}

		// Take the target: mark the job running and bump the in-flight count under the lock,
		// then fetch outside it.
		j.running = true
		r.inflight++
		r.callInflight(r.inflight)
		if !j.started {
			if err := r.db.SetBatchJobStatus(j.id, "running", "", time.Now().Unix()); err != nil {
				r.log.Warn("batch: set running", "job", j.id, "err", err)
			}
			j.started = true
		}
		idx := j.next
		target := j.targets[idx]
		startedAt := time.Now()
		r.mu.Unlock()

		res, err := r.runTarget(r.ctx, j, idx, target)

		r.mu.Lock()
		now = time.Now() // after the fetch: a busy answer can take seconds to come back
		j.running = false
		r.inflight--
		r.callInflight(r.inflight)

		// Shutting down: a failure now is most likely the cancelled context. Store nothing; the
		// job stays live in the store for the restart reaper.
		if r.closed && err != nil {
			r.cond.Broadcast()
			r.mu.Unlock()
			continue
		}

		// A cancelled job whose in-flight target was only busy is finished: nothing more runs.
		if errors.Is(err, ErrRetryLater) && j.cancelled {
			r.removeLocked(j)
			r.callDepth(len(r.live))
			r.cond.Broadcast()
			r.mu.Unlock()
			continue
		}

		// A target that was only busy stays pending and is retried after RetryWait, unless it
		// has been busy for longer than RetryFor, in which case it is stored failed below.
		if errors.Is(err, ErrRetryLater) {
			if j.busySince.IsZero() {
				j.busySince = now
			}
			if now.Sub(j.busySince) <= r.cfg.RetryFor {
				j.retryAt = now.Add(r.cfg.RetryWait)
				r.cond.Broadcast()
				r.mu.Unlock()
				continue
			}
		}

		resultBytes := r.storeTarget(j, idx, res, err)
		j.next++
		j.busySince, j.retryAt = time.Time{}, time.Time{}
		if err == nil {
			j.done++
		} else {
			j.failed++
		}
		r.callTarget(err, time.Since(startedAt), resultBytes)
		r.db.EvictBatchOldest(r.cfg.MaxBytes)

		if j.cancelled || j.next >= len(j.targets) {
			if !j.cancelled {
				r.setTerminal(j)
			}
			r.removeLocked(j)
			r.callDepth(len(r.live))
		}
		r.cond.Broadcast()
		r.mu.Unlock()
	}
}

// pick returns the first live job, scanning from cursor and wrapping once, that has a pending
// target that is not running, not cancelled, and ready (now past any retry). It advances cursor to
// the index after the job, so the rotation is round-robin across jobs in submit order.
func (r *Runner) pick(now time.Time) *liveJob {
	n := len(r.live)
	if n == 0 {
		return nil
	}
	for i := 0; i < n; i++ {
		j := r.live[(r.cursor+i)%n]
		if j.running || j.cancelled || j.next >= len(j.targets) {
			continue
		}
		if !now.Before(j.retryAt) { // now >= retryAt; a zero retryAt is always ready
			// The job after this one; n (one past the end) means "whatever is appended next,
			// else wrap to 0" — reads take cursor % n.
			r.cursor = (r.cursor+i)%n + 1
			return j
		}
	}
	return nil
}

// scheduleRetryLocked arms a timer that wakes the workers at the earliest future retryAt, so they
// do not sleep past a busy target that is ready to be tried again.
func (r *Runner) scheduleRetryLocked(now time.Time) {
	var earliest time.Time
	for _, j := range r.live {
		if j.cancelled || j.next >= len(j.targets) || !j.retryAt.After(now) {
			continue
		}
		if earliest.IsZero() || j.retryAt.Before(earliest) {
			earliest = j.retryAt
		}
	}
	if !earliest.IsZero() {
		d := time.Until(earliest)
		if d < 0 {
			d = 0
		}
		time.AfterFunc(d, r.cond.Broadcast)
	}
}

// removeLocked drops j from live and fixes cursor so the rotation still points at the next job.
func (r *Runner) removeLocked(j *liveJob) {
	idx := -1
	for i := range r.live {
		if r.live[i] == j {
			idx = i
			break
		}
	}
	if idx < 0 {
		return
	}
	r.live = append(r.live[:idx], r.live[idx+1:]...)
	if r.cursor > idx {
		r.cursor--
	}
	if r.cursor > len(r.live) {
		r.cursor = len(r.live)
	}
}

// setTerminal writes the job's finished status once: all done → done, all failed → failed, a mix
// → partial.
func (r *Runner) setTerminal(j *liveJob) {
	var status string
	switch {
	case j.failed == 0:
		status = "done"
	case j.done == 0:
		status = "failed"
	default:
		status = "partial"
	}
	if err := r.db.SetBatchJobStatus(j.id, status, "", time.Now().Unix()); err != nil {
		r.log.Warn("batch: set terminal status", "job", j.id, "err", err)
	}
	r.callJob(status)
	r.callJobSeconds(time.Since(j.created))
}

// runTarget fetches one target, recovering a panic into ErrInternal (never echoing the value).
func (r *Runner) runTarget(ctx context.Context, j *liveJob, idx int, t Target) (res Result, err error) {
	defer func() {
		if rec := recover(); rec != nil {
			r.log.Warn("batch: recover fetch panic", "job", j.id, "idx", idx)
			res = Result{}
			err = ErrInternal
		}
	}()
	return r.fetch(ctx, j.caller, t)
}

// storeTarget stores one target: success → done with the result JSON, error → failed with the
// TargetError's Msg/Hint or the error's text. It returns the number of result bytes stored.
func (r *Runner) storeTarget(j *liveJob, idx int, res Result, err error) int {
	now := time.Now().Unix()
	status := "failed"
	var resultStr, errMsg, hint *string
	if err == nil {
		status = "done"
		b, merr := json.Marshal(res)
		if merr != nil {
			r.log.Warn("batch: marshal result", "job", j.id, "idx", idx, "err", merr)
		} else {
			s := string(b)
			resultStr = &s
		}
	} else {
		var terr *TargetError
		if errors.As(err, &terr) {
			m := terr.Msg
			errMsg = &m
			if terr.Hint != "" {
				h := terr.Hint
				hint = &h
			}
		} else {
			e := err.Error()
			errMsg = &e
		}
	}
	if serr := r.db.UpdateBatchTarget(j.id, idx, now, status, resultStr, errMsg, hint); serr != nil {
		r.log.Warn("batch: store target", "job", j.id, "idx", idx, "err", serr)
	}
	if resultStr != nil {
		return len(*resultStr)
	}
	return 0
}

func (r *Runner) callDepth(live int) {
	if r.cfg.Hooks.Depth != nil {
		r.cfg.Hooks.Depth(live)
	}
}

func (r *Runner) callInflight(n int) {
	if r.cfg.Hooks.Inflight != nil {
		r.cfg.Hooks.Inflight(n)
	}
}

func (r *Runner) callTarget(err error, d time.Duration, resultBytes int) {
	if r.cfg.Hooks.Target != nil {
		r.cfg.Hooks.Target(err, d, resultBytes)
	}
}

func (r *Runner) callJobSeconds(d time.Duration) {
	if r.cfg.Hooks.JobSeconds != nil {
		r.cfg.Hooks.JobSeconds(d)
	}
}