api: NewBatchRunner — the batch fetch adapter, retry on busy, outcome metrics
Wire the batch runner to fetchTarget and feed the sift_batch_* metrics. New internal/api/batch.go: NewBatchRunner overwrites cfg.Hooks (Job, Target, JobSeconds, Depth, Inflight → s.Prom) and sets the runner's FetchFunc to batchFetch, which tags the context as a batch fetch, runs under requestDeadline, and maps fetchTarget's errors to what the runner retries (429/503 → ErrRetryLater) and what the metrics classify (batchOutcome). Added the Batch field to Server (nil → task 07 answers 503 "batch disabled") and, in extractPDF, charged the PDF pool to c.Name+"/batch" under the batch mark so a batch never takes a caller's interactive slot while the single-flight key is unchanged. No HTTP endpoints. Co-Authored-By: ornith-1.5-35b-a3b <ornith-1.5-35b-a3b@llama.cpp.invalid>
This commit is contained in:
ornith-1.5-35b-a3b
@@ -73,6 +73,10 @@ type Server struct {
|
||||
|
||||
pdfPoolOnce sync.Once
|
||||
|
||||
// Batch is the batch runner; nil means the /v1/batch endpoints (task 07) answer 503
|
||||
// "batch disabled".
|
||||
Batch *batch.Runner
|
||||
|
||||
// uploads bounds how many /v1/extract bodies (up to PDFMaxUpload each) are read into
|
||||
// memory at once, ahead of the PDF pool's own limits.
|
||||
uploads chan struct{}
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
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"
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package api
|
||||
|
||||
// P6 task 06: the runner wired to fetchTarget and the metrics, with no HTTP yet. The helpers
|
||||
// here (newBatchServer, pageServer, fakeDNS) are shared with task 07's batch_test.go.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/netip"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"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"
|
||||
"git.wntrmute.dev/kyle/sift/internal/guard"
|
||||
)
|
||||
|
||||
var loopback = []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}
|
||||
|
||||
// fakeDNS resolves internal.test to a private address; every other name fails to resolve.
|
||||
func fakeDNS(_ context.Context, host string) ([]netip.Addr, error) {
|
||||
if host == "internal.test" {
|
||||
return []netip.Addr{netip.MustParseAddr("10.0.0.5")}, nil
|
||||
}
|
||||
return nil, fmt.Errorf("no such host %s", host)
|
||||
}
|
||||
|
||||
// newBatchServer builds a Server with the batch runner wired over an in-memory dcache through
|
||||
// s.NewBatchRunner, exactly as serve does. callers are the names to provision (each gets a key).
|
||||
func newBatchServer(t *testing.T, callers ...string) (*Server, map[string]string) {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
keysPath := filepath.Join(dir, "keys.toml")
|
||||
keys := map[string]string{}
|
||||
for _, name := range callers {
|
||||
out := filepath.Join(dir, name)
|
||||
if err := auth.Add(keysPath, name, "hermes-fleet", out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b, err := os.ReadFile(out)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keys[name] = strings.TrimSpace(string(b))
|
||||
}
|
||||
store, err := auth.Load(keysPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
g := &guard.Guard{AllowPrefixes: loopback, LookupIP: fakeDNS}
|
||||
f := &fetch.Fetcher{Guard: g, UA: "sift-test"}
|
||||
log := slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
s := New("test", g, f, store, log)
|
||||
db, err := dcache.Open("", 1<<20, log)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
s.Batch = s.NewBatchRunner(db, batch.Config{})
|
||||
t.Cleanup(s.Batch.Close) // before db.Close
|
||||
return s, keys
|
||||
}
|
||||
|
||||
// pageServer serves a small HTML page at every path; hold, when non-nil, blocks each request
|
||||
// until it is closed.
|
||||
func pageServer(t *testing.T, hold chan struct{}) *httptest.Server {
|
||||
t.Helper()
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if hold != nil {
|
||||
select {
|
||||
case <-hold:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
fmt.Fprintf(w, `<html><head><title>T</title></head><body><article><p>hello batch %s</p></article></body></html>`, r.URL.Path)
|
||||
}))
|
||||
t.Cleanup(ts.Close)
|
||||
return ts
|
||||
}
|
||||
|
||||
func waitJob(t *testing.T, r *batch.Runner, id string) batch.Job {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(10 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
j, _, ok := r.Status(id)
|
||||
if ok && (j.Status == "done" || j.Status == "partial" || j.Status == "failed" || j.Status == "cancelled") {
|
||||
return j
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("job never finished")
|
||||
return batch.Job{}
|
||||
}
|
||||
|
||||
func promText(t *testing.T, s *Server) string {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodGet, "/metrics", nil)
|
||||
req.RemoteAddr, req.Host = "127.0.0.1:9", "127.0.0.1:8890"
|
||||
rec := httptest.NewRecorder()
|
||||
s.ServeHTTP(rec, req)
|
||||
return rec.Body.String()
|
||||
}
|
||||
|
||||
// The runner fetches through fetchTarget as the submitting caller; a name that resolves to a
|
||||
// private address is a per-target "blocked" failure; the hooks reach the metrics.
|
||||
func TestBatchRunnerWired(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, _ := newBatchServer(t, "test")
|
||||
c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
|
||||
id, reused, err := s.Batch.Submit(context.Background(), c, dcache.ClassFleet, []batch.Target{
|
||||
{URL: ts.URL + "/one"}, {URL: "http://internal.test/"},
|
||||
})
|
||||
if err != nil || reused {
|
||||
t.Fatalf("submit: %v %v", err, reused)
|
||||
}
|
||||
if j := waitJob(t, s.Batch, id); j.Status != "partial" {
|
||||
t.Fatalf("job %+v", j)
|
||||
}
|
||||
_, ts2, _ := s.Batch.Status(id)
|
||||
if ts2[0].Status != "done" || ts2[0].Result == nil || !strings.Contains(ts2[0].Result.Content, "hello batch /one") {
|
||||
t.Fatalf("target 0: %+v", ts2[0])
|
||||
}
|
||||
if ts2[0].Result.Meta["mode"] != "clean" || ts2[0].Result.Meta["untrusted"] != true {
|
||||
t.Fatalf("target 0 meta %v: not the fetchTarget envelope", ts2[0].Result.Meta)
|
||||
}
|
||||
if ts2[1].Status != "failed" || !strings.Contains(ts2[1].Error, "blocked") {
|
||||
t.Fatalf("target 1: %+v", ts2[1])
|
||||
}
|
||||
out := promText(t, s)
|
||||
for _, want := range []string{
|
||||
`sift_batch_jobs_total{outcome="created"} 1`,
|
||||
`sift_batch_jobs_total{outcome="partial"} 1`,
|
||||
`sift_batch_targets_total{outcome="ok"} 1`,
|
||||
`sift_batch_targets_total{outcome="blocked"} 1`,
|
||||
`sift_batch_result_bytes_count 1`,
|
||||
`sift_batch_queue_depth 0`,
|
||||
} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Errorf("metrics lack %q", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The adapter turns fetch errors into what the runner and the metrics need.
|
||||
func TestBatchOutcomeClasses(t *testing.T) {
|
||||
cases := []struct {
|
||||
err error
|
||||
want string
|
||||
}{
|
||||
{nil, "ok"},
|
||||
{fmt.Errorf("busy: %w", batch.ErrRetryLater), "busy"},
|
||||
{&batch.TargetError{Status: 403, Msg: "blocked: private address"}, "blocked"},
|
||||
{&batch.TargetError{Status: 429, Msg: "busy"}, "busy"},
|
||||
{&batch.TargetError{Status: 504, Msg: "timeout"}, "timeout"},
|
||||
{&batch.TargetError{Status: 415, Msg: "unsupported content type"}, "unsupported"},
|
||||
{&batch.TargetError{Status: 422, Msg: "document too long for question (6001 chunks); use grep or pages"}, "too_long"},
|
||||
{&batch.TargetError{Status: 422, Msg: "unreadable PDF"}, "upstream_4xx"},
|
||||
{&batch.TargetError{Status: 404, Msg: "not found"}, "upstream_4xx"},
|
||||
{&batch.TargetError{Status: 500, Msg: "internal error"}, "upstream_5xx"},
|
||||
{errors.New("something else"), "upstream_5xx"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := batchOutcome(c.err); got != c.want {
|
||||
t.Errorf("batchOutcome(%v) = %q, want %q", c.err, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// batchFetch maps a busy fetch to ErrRetryLater (the runner retries it) and every other
|
||||
// failure to a TargetError with the /v1/fetch status.
|
||||
func TestBatchFetchErrors(t *testing.T) {
|
||||
s, _ := newBatchServer(t, "test")
|
||||
c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
|
||||
_, err := s.batchFetch(context.Background(), c, batch.Target{URL: "http://internal.test/"})
|
||||
var te *batch.TargetError
|
||||
if !errors.As(err, &te) || te.Status != 403 || !strings.Contains(te.Msg, "blocked") {
|
||||
t.Fatalf("private name: %v, want a TargetError 403 blocked", err)
|
||||
}
|
||||
// Fill the global fetch gate so the next fetch is refused busy.
|
||||
for i := 0; i < 8; i++ {
|
||||
if err := s.acquire(fmt.Sprintf("filler%d", i), "http://filler.test/"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ts := pageServer(t, nil)
|
||||
_, err = s.batchFetch(context.Background(), c, batch.Target{URL: ts.URL})
|
||||
if !errors.Is(err, batch.ErrRetryLater) {
|
||||
t.Fatalf("busy gate: %v, want an error wrapping batch.ErrRetryLater", err)
|
||||
}
|
||||
}
|
||||
+5
-1
@@ -58,7 +58,11 @@ func (s *Server) extractPDF(ctx context.Context, c *auth.Caller, res *fetch.Reso
|
||||
return v, err
|
||||
}
|
||||
start := time.Now()
|
||||
v, err := s.pool().Do(ctx, s.callerClass(c)+"|"+res.SHA256(), c.Name, fn)
|
||||
caller := c.Name
|
||||
if isBatchFetch(ctx) {
|
||||
caller = c.Name + "/batch"
|
||||
}
|
||||
v, err := s.pool().Do(ctx, s.callerClass(c)+"|"+res.SHA256(), caller, fn)
|
||||
total := time.Since(start)
|
||||
// The job's own extraction time is what this call measured it to be; the
|
||||
// rest of Do is time spent waiting for a slot.
|
||||
|
||||
Reference in New Issue
Block a user