sift/internal/api/batch.go

package api

import (
	"context"
	"errors"
	"fmt"
	"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)
			}
			if fe.Status == 502 {
				fe.Status = 500
				fe.Msg = "internal error"
			}
			if fe.Status == 403 {
				s.blockN.Add(1)
			}
			return batch.Result{}, &batch.TargetError{Status: fe.Status, Msg: fe.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"
}