batch: workers — round-robin, serial per job, busy retry, cancel in flight, hooks

Start cfg.Workers goroutines in New that drain one rotation of live jobs in submit order, each taking the next job with a pending target and none in flight, so targets alternate across jobs and each job stays serial. One target per fetch under the Close-cancelled context: panic recovers into ErrInternal (never echoed), ErrRetryLater retries after RetryWait (bounded by RetryFor), otherwise the target is stored done/failed and the hooks fire. Terminal status is set once per job; cancelling a running job stores the in-flight target without overwriting cancelled and leaves later targets pending. Co-Authored-By: ornith-1.5-35b-a3b <ornith-1.5-35b-a3b@llama.cpp.invalid>
2026-09-26 22:18:35 -07:00
parent 0e691a6b4f
commit 816a76fe5b
4 changed files with 715 additions and 6 deletions
@@ -0,0 +1,271 @@
+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()
+		now := time.Now()
+		j := r.pick(now)
+		if j == nil {
+			if r.closed {
+				r.mu.Unlock()
+				return
+			}
+			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()
+		j.running = false
+		r.inflight--
+		r.callInflight(r.inflight)
+
+		// 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.forgetCaller(j.id)
+			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
+			r.cursor += i + 1 // absolute; wraps at access, so the rotation survives jobs being added
+			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--
+	}
+}
+
+// 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
+}
+
+// forgetCaller drops the in-memory caller for a finished job.
+func (r *Runner) forgetCaller(id string) {
+	r.callMu.Lock()
+	delete(r.callers, id)
+	r.callMu.Unlock()
+}
+
+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)
+	}
+}