sift/internal/dcache/batch_test.go

package dcache

import (
	"io"
	"log/slog"
	"strings"
	"testing"
	"time"
)

func newBatchDB(t *testing.T) *DB {
	t.Helper()
	db, err := Open("", 1<<20, slog.New(slog.NewTextHandler(io.Discard, nil)))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Close() })
	return db
}

func strp(s string) *string { return &s }

func TestBatchJobRoundTrip(t *testing.T) {
	db := newBatchDB(t)
	j := BatchJob{
		ID: "j1", Caller: "alice", CallerClass: ClassFleet, IDemKey: "k1",
		Status: "queued", TargetsTotal: 3, CreatedAt: 100, Lease: 100,
	}
	if err := db.CreateBatchJob(j); err != nil {
		t.Fatal(err)
	}
	got, ok := db.BatchJob("j1")
	if !ok {
		t.Fatal("missing job")
	}
	if got.ID != j.ID || got.Caller != "alice" || got.CallerClass != ClassFleet ||
		got.Status != "queued" || got.TargetsTotal != 3 || got.Lease != 100 || got.CreatedAt != 100 {
		t.Fatalf("round trip mismatch: %+v", got)
	}
	if _, ok := db.BatchJob("nope"); ok {
		t.Fatal("unknown id must be ok=false")
	}
}

func TestBatchIDemLookup(t *testing.T) {
	db := newBatchDB(t)
	if err := db.CreateBatchJob(BatchJob{ID: "j1", Caller: "a", CallerClass: ClassFleet, IDemKey: "k", Status: "running", TargetsTotal: 1, CreatedAt: 100, Lease: 1000}); err != nil {
		t.Fatal(err)
	}
	if id, ok := db.BatchJobByIDem("a", "k", 2000, time.Hour); !ok || id != "j1" {
		t.Fatalf("live idem: %q %v", id, ok)
	}
	if _, ok := db.BatchJobByIDem("a", "k", 5000, time.Hour); ok {
		t.Fatal("lease-expired job must not match")
	}
	if _, ok := db.BatchJobByIDem("b", "k", 2000, time.Hour); ok {
		t.Fatal("other caller must not match")
	}
	// Only a job that is live or has results is reused. A failed (e.g. killed by a restart)
	// or cancelled job must not be handed back to a resubmit for the rest of its lease.
	for _, c := range []struct {
		status string
		reuse  bool
	}{{"queued", true}, {"done", true}, {"partial", true}, {"failed", false}, {"cancelled", false}} {
		key := "k-" + c.status
		if err := db.CreateBatchJob(BatchJob{ID: "j-" + c.status, Caller: "a", CallerClass: ClassFleet, IDemKey: key, Status: c.status, TargetsTotal: 1, CreatedAt: 100, Lease: 1000}); err != nil {
			t.Fatal(err)
		}
		if _, ok := db.BatchJobByIDem("a", key, 2000, time.Hour); ok != c.reuse {
			t.Errorf("status %s: reused=%v, want %v", c.status, ok, c.reuse)
		}
	}
}

func TestBatchLiveCounts(t *testing.T) {
	db := newBatchDB(t)
	mk := func(id, caller, status string) {
		if err := db.CreateBatchJob(BatchJob{ID: id, Caller: caller, CallerClass: ClassFleet, IDemKey: id, Status: status, TargetsTotal: 1, CreatedAt: 1, Lease: 1}); err != nil {
			t.Fatal(err)
		}
	}
	mk("j1", "a", "queued")
	mk("j2", "a", "running")
	mk("j3", "a", "done")
	mk("j4", "b", "running")
	if n := db.LiveBatchJobs("a"); n != 2 {
		t.Fatalf("caller a live=%d want 2", n)
	}
	if n := db.LiveBatchJobsTotal(); n != 3 {
		t.Fatalf("total live=%d want 3", n)
	}
}

func TestBatchTargetsOrderedAndUpdate(t *testing.T) {
	db := newBatchDB(t)
	must := func(err error) {
		t.Helper()
		if err != nil {
			t.Fatal(err)
		}
	}
	must(db.CreateBatchJob(BatchJob{ID: "j1", Caller: "a", CallerClass: ClassFleet, IDemKey: "k", Status: "running", TargetsTotal: 3, CreatedAt: 100, Lease: 100}))
	must(db.PutBatchTargets("j1", []BatchTarget{
		{JobID: "j1", Idx: 0, URL: "https://a.example/", Options: `{"mode":"clean"}`, Status: "pending"},
		{JobID: "j1", Idx: 1, URL: "https://b.example/", Status: "pending"},
		{JobID: "j1", Idx: 2, URL: "https://c.example/", Status: "pending"},
	}))
	got, ok := db.BatchTargets("j1")
	if !ok || len(got) != 3 {
		t.Fatalf("targets len=%d ok=%v", len(got), ok)
	}
	for i := range got {
		if got[i].Idx != i {
			t.Fatalf("order broken: got[%d].Idx=%d", i, got[i].Idx)
		}
	}

	// A done write bumps the lease to now and counts the target once.
	res := strp(`{"content":"ok"}`)
	must(db.UpdateBatchTarget("j1", 0, 200, "done", res, nil, nil))
	j, _ := db.BatchJob("j1")
	if j.Lease != 200 || j.TargetsDone != 1 {
		t.Fatalf("after first write lease=%d done=%d", j.Lease, j.TargetsDone)
	}
	// Re-writing the same target (still a done transition) must not double-count.
	must(db.UpdateBatchTarget("j1", 0, 201, "done", res, nil, nil))
	j, _ = db.BatchJob("j1")
	if j.TargetsDone != 1 {
		t.Fatalf("double-counted: done=%d", j.TargetsDone)
	}
	// A failed write counts too.
	must(db.UpdateBatchTarget("j1", 1, 202, "failed", nil, strp("upstream 500"), strp("browser")))
	j, _ = db.BatchJob("j1")
	if j.TargetsDone != 2 || j.Lease != 202 {
		t.Fatalf("after failed write lease=%d done=%d", j.Lease, j.TargetsDone)
	}
}

func TestBatchTouchAndStatus(t *testing.T) {
	db := newBatchDB(t)
	if err := db.CreateBatchJob(BatchJob{ID: "j1", Caller: "a", CallerClass: ClassFleet, IDemKey: "k", Status: "running", TargetsTotal: 1, CreatedAt: 100, Lease: 100}); err != nil {
		t.Fatal(err)
	}
	if err := db.TouchBatchJob("j1", 500); err != nil {
		t.Fatal(err)
	}
	if j, _ := db.BatchJob("j1"); j.Lease != 500 {
		t.Fatalf("touch lease=%d want 500", j.Lease)
	}
	if err := db.SetBatchJobStatus("j1", "cancelled", "", 600); err != nil {
		t.Fatal(err)
	}
	if j, _ := db.BatchJob("j1"); j.Status != "cancelled" || j.Lease != 600 {
		t.Fatalf("status=%q lease=%d", j.Status, j.Lease)
	}
}

func TestSweepBatch(t *testing.T) {
	db := newBatchDB(t)
	mk := func(id string, lease int64) {
		if err := db.CreateBatchJob(BatchJob{ID: id, Caller: "a", CallerClass: ClassFleet, IDemKey: id, Status: "done", TargetsTotal: 1, CreatedAt: 1, Lease: lease}); err != nil {
			t.Fatal(err)
		}
		if err := db.PutBatchTargets(id, []BatchTarget{{JobID: id, Idx: 0, URL: "https://x.example/", Status: "done", Result: strp("x")}}); err != nil {
			t.Fatal(err)
		}
	}
	mk("old", 100)
	mk("new", 1000)
	n := db.SweepBatch(300, 100*time.Second) // old: 100+100=200<=300 expired; new: 1100>300 kept
	if n != 1 {
		t.Fatalf("swept %d want 1", n)
	}
	if _, ok := db.BatchJob("old"); ok {
		t.Fatal("old job must be gone")
	}
	if _, ok := db.BatchJob("new"); !ok {
		t.Fatal("new job must remain")
	}
	if _, ok := db.BatchTargets("old"); ok {
		t.Fatal("old targets must be gone")
	}
}

func TestEvictBatchOldest(t *testing.T) {
	db := newBatchDB(t)
	done := func(id string, lease int64, size int) {
		if err := db.CreateBatchJob(BatchJob{ID: id, Caller: "a", CallerClass: ClassFleet, IDemKey: id, Status: "done", TargetsTotal: 1, CreatedAt: 1, Lease: lease}); err != nil {
			t.Fatal(err)
		}
		if err := db.PutBatchTargets(id, []BatchTarget{{JobID: id, Idx: 0, URL: "https://x.example/", Status: "pending"}}); err != nil {
			t.Fatal(err)
		}
		if err := db.UpdateBatchTarget(id, 0, lease, "done", strp(strings.Repeat("x", size)), nil, nil); err != nil {
			t.Fatal(err)
		}
	}
	done("old1", 100, 100)
	done("old2", 200, 200)
	done("new", 300, 300)

	// A running job with one completed target must never be byte-evicted.
	if err := db.CreateBatchJob(BatchJob{ID: "run", Caller: "a", CallerClass: ClassFleet, IDemKey: "run", Status: "running", TargetsTotal: 2, TargetsDone: 1, CreatedAt: 1, Lease: 400}); err != nil {
		t.Fatal(err)
	}
	if err := db.PutBatchTargets("run", []BatchTarget{
		{JobID: "run", Idx: 0, URL: "https://x.example/", Status: "pending"},
		{JobID: "run", Idx: 1, URL: "https://y.example/", Status: "pending"},
	}); err != nil {
		t.Fatal(err)
	}
	if err := db.UpdateBatchTarget("run", 0, 400, "done", strp(strings.Repeat("x", 500)), nil, nil); err != nil {
		t.Fatal(err)
	}

	// total nbytes = 100+200+300+500 = 1100. Evict oldest done until <= 250:
	// old1(100), old2(200), new(300) all evicted (oldest first); run stays (live).
	db.EvictBatchOldest(250)
	for _, id := range []string{"old1", "old2", "new"} {
		if _, ok := db.BatchJob(id); ok {
			t.Fatalf("%s should be byte-evicted", id)
		}
	}
	if _, ok := db.BatchJob("run"); !ok {
		t.Fatal("running job must not be byte-evicted")
	}
	if n := db.BatchNbytes(); n != 500 {
		t.Fatalf("BatchNbytes=%d want 500", n)
	}
}

func TestMarkInFlightFailed(t *testing.T) {
	db := newBatchDB(t)
	mk := func(id, status string) {
		if err := db.CreateBatchJob(BatchJob{ID: id, Caller: "a", CallerClass: ClassFleet, IDemKey: id, Status: status, TargetsTotal: 1, CreatedAt: 1, Lease: 1}); err != nil {
			t.Fatal(err)
		}
	}
	mk("q", "queued")
	mk("r", "running")
	mk("d", "done")
	if n := db.MarkInFlightFailed("restart", 999); n != 2 {
		t.Fatalf("marked %d want 2", n)
	}
	if j, _ := db.BatchJob("q"); j.Status != "failed" || j.Error != "restart" || j.Lease != 999 {
		t.Fatalf("q %+v", j)
	}
	if j, _ := db.BatchJob("r"); j.Status != "failed" || j.Error != "restart" {
		t.Fatalf("r %+v", j)
	}
	if j, _ := db.BatchJob("d"); j.Status != "done" {
		t.Fatalf("done job changed: %+v", j)
	}
}