sift/internal/api/batch.go

package api

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

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

// batchCtxKey marks a context as a batch fetch so extractPDF charges the PDF pool to the
// caller's "/batch" slot instead of its interactive one.
type batchCtxKey struct{}

func isBatchFetch(ctx context.Context) bool {
	ok, _ := ctx.Value(batchCtxKey{}).(bool)
	return ok
}

// NewBatchRunner builds the runner that serve (task 07) and the tests use. Its FetchFunc is
// s.batchFetch and its Hooks feed s.Prom; any cfg.Hooks from the caller is overwritten.
func (s *Server) NewBatchRunner(db *dcache.DB, cfg batch.Config) *batch.Runner {
	cfg.Hooks = batch.Hooks{
		Job: func(outcome string) {
			s.Prom.ObserveBatchJob(outcome)
		},
		Target: func(err error, d time.Duration, n int) {
			s.Prom.ObserveBatchTarget(batchOutcome(err))
			s.Prom.ObserveBatchTargetSeconds(d)
			if err == nil {
				s.Prom.ObserveBatchResultBytes(n)
			}
		},
		JobSeconds: func(d time.Duration) {
			s.Prom.ObserveBatchJobSeconds(d)
		},
		Depth: func(live int) {
			s.Prom.SetBatchQueueDepth(live)
		},
		Inflight: func(n int) {
			s.Prom.SetBatchTargetsInflight(n)
		},
	}
	return batch.New(db, s.batchFetch, cfg, s.Log)
}

// batchFetch is the runner's FetchFunc: one target through s.fetchTarget under the request
// deadline, tagged as a batch fetch. Its errors map to what the runner and the metrics need.
func (s *Server) batchFetch(ctx context.Context, c *auth.Caller, t batch.Target) (batch.Result, error) {
	ctx, cancel := context.WithTimeout(ctx, requestDeadline)
	defer cancel()
	ctx = context.WithValue(ctx, batchCtxKey{}, true)
	res, err := s.fetchTarget(ctx, c, t)
	if err != nil {
		var fe *fetch.Error
		if errors.As(err, &fe) {
			switch fe.Status {
			case 429, 503:
				return batch.Result{}, fmt.Errorf("%s: %w", fe.Msg, batch.ErrRetryLater)
			}
			// Copy, never mutate: the PDF pool's single-flight may share this error with an
			// interactive request.
			status, msg := fe.Status, fe.Msg
			if status == 502 {
				status, msg = 500, "internal error"
			}
			if status == 403 {
				s.blockN.Add(1)
			}
			return batch.Result{}, &batch.TargetError{Status: status, Msg: msg, Hint: fe.Hint}
		}
		s.Log.Error("internal", "caller", c.Name, "host", "-", "err", s.logErr(c, err))
		return batch.Result{}, &batch.TargetError{Status: 500, Msg: "internal error"}
	}
	return res, nil
}

// batchOutcome classifies a finished target's error into the sift_batch_targets_total label:
// a nil error is ok, a retry-later the runner gave up on is busy, a TargetError by its status,
// and anything else upstream_5xx so a raw string never becomes a label.
func batchOutcome(err error) string {
	switch {
	case err == nil:
		return "ok"
	case errors.Is(err, batch.ErrRetryLater):
		return "busy"
	}
	var te *batch.TargetError
	if errors.As(err, &te) {
		switch te.Status {
		case 403:
			return "blocked"
		case 429:
			return "busy"
		case 504:
			return "timeout"
		case 415:
			return "unsupported"
		case 422:
			if strings.HasPrefix(te.Msg, "document too long") {
				return "too_long"
			}
			return "upstream_4xx"
		}
		if te.Status >= 400 && te.Status < 500 {
			return "upstream_4xx"
		}
	}
	return "upstream_5xx"
}

// The four /v1/batch handlers (task 07). A job that is unknown or owned by another caller is a
// 404 for every GET and DELETE, never a 403, so a caller cannot probe for another caller's jobs.

// ownership answers the caller for id: a disabled batch is 503, an unknown id or one not owned by
// c is a 404. The empty job returned on either path signals "already answered" to the caller.
func (s *Server) ownership(w http.ResponseWriter, c *auth.Caller, id string) (batch.Job, []batch.TargetStatus) {
	if s.Batch == nil {
		writeErr(w, 503, "batch disabled", "")
		return batch.Job{}, nil
	}
	j, ts, ok := s.Batch.Status(id)
	if !ok || j.Caller != c.Name {
		writeErr(w, 404, "not found", "")
		return batch.Job{}, nil
	}
	return j, ts
}

// logBatch writes the one batch line per handler: caller, job_id, targets, then either status or
// reused. It never carries a target URL, hostname or the idem_key.
func (s *Server) logBatch(c *auth.Caller, jobID string, targets int, extra ...any) {
	s.Log.Info("batch",
		append([]any{
			"caller", c.Name,
			"job_id", jobID,
			"targets", targets,
		}, extra...)...,
	)
}

func validBatchMode(m string) bool {
	switch m {
	case "", "clean", "raw", "digest":
		return true
	}
	return false
}

func (s *Server) batchSubmit(w http.ResponseWriter, r *http.Request) {
	c := s.caller(w, r)
	if c == nil {
		return
	}
	if s.Batch == nil { // [batch] disabled, or the store did not open
		writeErr(w, 503, "batch disabled", "")
		return
	}
	var req struct {
		Targets []v1Req `json:"targets"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeErr(w, 400, "invalid batch: bad json", "")
		return
	}
	if len(req.Targets) == 0 {
		writeErr(w, 400, "invalid batch: no targets", "")
		return
	}
	targets := make([]batch.Target, 0, len(req.Targets))
	for _, t := range req.Targets {
		if t.URL == "" {
			writeErr(w, 400, "invalid batch: empty url", "")
			return
		}
		if !validBatchMode(t.Mode) {
			writeErr(w, 400, "invalid batch: bad mode", "")
			return
		}
		targets = append(targets, batch.Target{
			URL: t.URL,
			Options: batch.Options{
				Mode:     t.Mode,
				Question: t.Question,
				Pages:    t.Pages,
				Grep:     t.Grep,
				JQ:       t.JQ,
				Context:  t.Context,
				Recover:  t.Recover,
			},
		})
	}
	for i, t := range targets {
		if err := s.Guard.PreCheck(t.URL); err != nil {
			s.logBatch(c, "", len(targets), "status", "blocked")
			writeErr(w, 422, fmt.Sprintf("blocked: target %d: %v", i, err), "")
			return
		}
	}
	id, reused, err := s.Batch.Submit(r.Context(), c, s.callerClass(c), targets)
	if err != nil {
		s.logBatch(c, "", len(targets), "status", "failed")
		switch {
		case errors.Is(err, batch.ErrTooManyTargets):
			writeErr(w, 422, "too many targets (max 32)", "")
		case errors.Is(err, batch.ErrCallerBusy):
			w.Header().Set("Retry-After", "30")
			writeErr(w, 429, "busy: a batch of yours is still running", "")
		case errors.Is(err, batch.ErrBusy):
			w.Header().Set("Retry-After", "60")
			writeErr(w, 429, "batch jobs busy", "")
		case errors.Is(err, batch.ErrNoTargets):
			writeErr(w, 400, "invalid batch: no targets", "")
		default:
			writeErr(w, 500, "internal error", "")
		}
		return
	}
	status := "queued"
	if reused {
		if j, _, ok := s.Batch.Status(id); ok {
			status = j.Status
		}
	}
	if reused {
		s.logBatch(c, id, len(targets), "reused", true)
	} else {
		s.logBatch(c, id, len(targets), "status", status)
	}
	writeJSON(w, 202, map[string]any{
		"success": true,
		"job_id":  id,
		"status":  status,
		"targets": len(targets),
		"reused":  reused,
	})
}

func (s *Server) batchStatus(w http.ResponseWriter, r *http.Request) {
	c := s.caller(w, r)
	if c == nil {
		return
	}
	id := r.PathValue("id")
	j, ts := s.ownership(w, c, id)
	if j.ID == "" {
		return
	}
	s.Batch.Touch(id)
	s.logBatch(c, id, j.TargetsTotal, "status", j.Status)
	writeJSON(w, 200, statusEnvelope(j, ts, r))
}

func (s *Server) batchResult(w http.ResponseWriter, r *http.Request) {
	c := s.caller(w, r)
	if c == nil {
		return
	}
	id := r.PathValue("id")
	j, ts := s.ownership(w, c, id)
	if j.ID == "" {
		return
	}
	idx, err := strconv.Atoi(r.PathValue("idx"))
	if err != nil || idx < 0 || idx >= j.TargetsTotal || idx >= len(ts) {
		writeErr(w, 404, "not found", "")
		return
	}
	s.Batch.Touch(id)
	t := ts[idx]
	out := map[string]any{
		"job_id": j.ID,
		"idx":    t.Idx,
		"url":    t.URL,
		"status": t.Status,
	}
	if t.Result != nil {
		out["result"] = resultEnvelope(t.Result)
	}
	if t.Error != "" {
		out["error"] = t.Error
	}
	if t.Hint != "" {
		out["hint"] = t.Hint
	}
	s.logBatch(c, id, j.TargetsTotal, "status", t.Status)
	writeJSON(w, 200, out)
}

func (s *Server) batchCancel(w http.ResponseWriter, r *http.Request) {
	c := s.caller(w, r)
	if c == nil {
		return
	}
	id := r.PathValue("id")
	j, _ := s.ownership(w, c, id)
	if j.ID == "" {
		return
	}
	switch j.Status {
	case "done", "partial", "failed", "cancelled":
		s.logBatch(c, id, j.TargetsTotal, "status", "failed")
		writeErr(w, 409, "job already finished", "")
		return
	}
	if s.Batch.Cancel(id) {
		s.logBatch(c, id, j.TargetsTotal, "status", "cancelled")
		writeJSON(w, 200, map[string]any{"success": true, "status": "cancelled"})
		return
	}
	s.logBatch(c, id, j.TargetsTotal, "status", "failed")
	writeErr(w, 409, "job already finished", "")
}

// statusEnvelope is the summary GET /v1/batch/{id} returns: statuses, never inline results,
// because 32 results can be several megabytes. With ?results=1 each done target carries the
// /v1/fetch envelope (Meta plus content).
func statusEnvelope(j batch.Job, ts []batch.TargetStatus, r *http.Request) map[string]any {
	out := map[string]any{
		"job_id":        j.ID,
		"status":        j.Status,
		"targets_total": j.TargetsTotal,
		"targets_done":  j.TargetsDone,
		"created_at":    time.Unix(j.CreatedAt, 0).UTC().Format(time.RFC3339),
		"targets":       targetSummaries(ts, r),
	}
	if j.Error != "" {
		out["error"] = j.Error
	}
	return out
}

func targetSummaries(ts []batch.TargetStatus, r *http.Request) []map[string]any {
	results := r.URL.Query().Get("results") == "1"
	out := make([]map[string]any, 0, len(ts))
	for _, t := range ts {
		m := map[string]any{
			"idx":    t.Idx,
			"url":    t.URL,
			"status": t.Status,
		}
		if t.Error != "" {
			m["error"] = t.Error
		}
		if t.Hint != "" {
			m["hint"] = t.Hint
		}
		if results && t.Status == "done" && t.Result != nil {
			m["result"] = resultEnvelope(t.Result)
		}
		out = append(out, m)
	}
	return out
}

// resultEnvelope is the /v1/fetch shape: every Meta field plus content.
func resultEnvelope(res *batch.Result) map[string]any {
	env := make(map[string]any, len(res.Meta)+1)
	for k, v := range res.Meta {
		env[k] = v
	}
	env["content"] = res.Content
	return env
}