internal/depot/sweep.go

@@ -0,0 +1,156 @@
+package depot
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"sort"
+	"sync"
+	"time"
+)
+
+// confirmSleep is the backoff sleeper for the serial confirm phase;
+// overridden in tests.
+var confirmSleep = time.Sleep
+
+// SweepResult partitions the swept shas by confirmed state. Unconfirmed is
+// distinct from Missing and must never be collapsed into it: it means the
+// probe budget was exhausted, not that the blob is known absent.
+type SweepResult struct {
+	Present     []string
+	Missing     []string
+	Unconfirmed []string
+}
+
+type probeOutcome struct {
+	sha       string
+	state     ProbeState
+	transient bool
+}
+
+// Sweep probes every sha against b: one parallel pass (cfg.ProbeJobs workers)
+// followed by a serial confirm phase over every non-Present result (unless
+// the non-Present count exceeds cfg.ConfirmMax, in which case all of them
+// are reported Unconfirmed with zero re-probes). Progress is reported to
+// progress roughly every 1000 blobs.
+func Sweep(ctx context.Context, b Backend, shas []string, cfg Config, progress io.Writer) (SweepResult, error) {
+	sorted := append([]string(nil), shas...)
+	sort.Strings(sorted)
+
+	outcomes := make([]probeOutcome, len(sorted))
+	jobs := cfg.ProbeJobs
+	if jobs < 1 {
+		jobs = 1
+	}
+
+	var wg sync.WaitGroup
+	work := make(chan int)
+	var probed int64
+	var mu sync.Mutex
+	var firstErr error
+
+	worker := func() {
+		defer wg.Done()
+		for i := range work {
+			state, transient, err := b.Probe(ctx, sorted[i])
+			mu.Lock()
+			if err != nil && !transient {
+				if firstErr == nil {
+					firstErr = fmt.Errorf("probe %s: %w", sorted[i], err)
+				}
+			} else {
+				outcomes[i] = probeOutcome{sha: sorted[i], state: state, transient: transient}
+			}
+			probed++
+			n := probed
+			mu.Unlock()
+			if n%1000 == 0 {
+				fmt.Fprintf(progress, "probed %d/%d\n", n, len(sorted))
+			}
+		}
+	}
+
+	for w := 0; w < jobs; w++ {
+		wg.Add(1)
+		go worker()
+	}
+	for i := range sorted {
+		work <- i
+	}
+	close(work)
+	wg.Wait()
+
+	if firstErr != nil {
+		return SweepResult{}, firstErr
+	}
+
+	var res SweepResult
+	var pending []probeOutcome
+	for _, o := range outcomes {
+		switch o.state {
+		case Present:
+			res.Present = append(res.Present, o.sha)
+		case Absent:
+			res.Missing = append(res.Missing, o.sha)
+		default:
+			pending = append(pending, o)
+		}
+	}
+
+	if cfg.ConfirmMax > 0 && len(pending) > cfg.ConfirmMax {
+		for _, o := range pending {
+			res.Unconfirmed = append(res.Unconfirmed, o.sha)
+		}
+		return res, nil
+	}
+
+	for _, o := range pending {
+		state, err := confirm(ctx, b, o.sha, cfg.ConfirmRetries)
+		if err != nil {
+			return SweepResult{}, err
+		}
+		switch state {
+		case Present:
+			res.Present = append(res.Present, o.sha)
+		case Absent:
+			res.Missing = append(res.Missing, o.sha)
+		default:
+			res.Unconfirmed = append(res.Unconfirmed, o.sha)
+		}
+	}
+
+	sort.Strings(res.Present)
+	sort.Strings(res.Missing)
+	sort.Strings(res.Unconfirmed)
+
+	if len(sorted)%1000 != 0 {
+		fmt.Fprintf(progress, "probed %d/%d\n", len(sorted), len(sorted))
+	}
+
+	return res, nil
+}
+
+// confirm re-probes sha serially up to retries attempts with linear backoff
+// (1s, 2s, ... between attempts). A 2xx result confirms Present, a 404/410
+// confirms Missing (Absent) immediately. If it is still transient after all
+// attempts, the state is left Unconfirmed rather than guessed as Missing.
+func confirm(ctx context.Context, b Backend, sha string, retries int) (ProbeState, error) {
+	if retries < 1 {
+		retries = 1
+	}
+	var last ProbeState = Unconfirmed
+	for attempt := 1; attempt <= retries; attempt++ {
+		state, transient, err := b.Probe(ctx, sha)
+		if err != nil && !transient {
+			return Unconfirmed, fmt.Errorf("confirm %s: %w", sha, err)
+		}
+		if state == Present || state == Absent {
+			return state, nil
+		}
+		last = Unconfirmed
+		if attempt < retries {
+			confirmSleep(time.Duration(attempt) * time.Second)
+		}
+	}
+	return last, nil
+}

internal/depot/sweep_test.go

@@ -0,0 +1,169 @@
+package depot
+
+import (
+	"context"
+	"io"
+	"sort"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+// stubBackend lets tests script Probe behavior per call.
+type stubBackend struct {
+	probe func(ctx context.Context, sha string) (ProbeState, bool, error)
+}
+
+func (s *stubBackend) Name() string { return "stub" }
+func (s *stubBackend) Probe(ctx context.Context, sha string) (ProbeState, bool, error) {
+	return s.probe(ctx, sha)
+}
+func (s *stubBackend) Get(ctx context.Context, sha, dest string) error { return nil }
+func (s *stubBackend) Put(ctx context.Context, sha, src string) error  { return nil }
+
+func sweepTestCfg() Config {
+	return Config{ProbeJobs: 4, ConfirmRetries: 3, ConfirmMax: 200}
+}
+
+func TestSweepAllPresent(t *testing.T) {
+	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
+		return Present, false, nil
+	}}
+	shas := []string{"bbb", "aaa", "ccc"}
+	res, err := Sweep(context.Background(), b, shas, sweepTestCfg(), io.Discard)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	want := []string{"aaa", "bbb", "ccc"}
+	sort.Strings(want)
+	if len(res.Present) != 3 || len(res.Missing) != 0 || len(res.Unconfirmed) != 0 {
+		t.Fatalf("got %+v", res)
+	}
+	for i, sha := range res.Present {
+		if sha != want[i] {
+			t.Fatalf("expected sorted output, got %v", res.Present)
+		}
+	}
+}
+
+func TestSweepMissing(t *testing.T) {
+	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
+		return Absent, false, nil
+	}}
+	shas := []string{"aaa"}
+	res, err := Sweep(context.Background(), b, shas, sweepTestCfg(), io.Discard)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(res.Missing) != 1 || res.Missing[0] != "aaa" {
+		t.Fatalf("got %+v", res)
+	}
+	if len(res.Present) != 0 || len(res.Unconfirmed) != 0 {
+		t.Fatalf("got %+v", res)
+	}
+}
+
+func TestSweepThrottledUnconfirmed(t *testing.T) {
+	var sleeps []time.Duration
+	orig := confirmSleep
+	confirmSleep = func(d time.Duration) { sleeps = append(sleeps, d) }
+	defer func() { confirmSleep = orig }()
+
+	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
+		return Unconfirmed, true, nil
+	}}
+	shas := []string{"aaa"}
+	cfg := sweepTestCfg()
+	res, err := Sweep(context.Background(), b, shas, cfg, io.Discard)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(res.Unconfirmed) != 1 || res.Unconfirmed[0] != "aaa" {
+		t.Fatalf("got %+v", res)
+	}
+	if len(res.Missing) != 0 {
+		t.Fatalf("must never collapse into missing: %+v", res)
+	}
+	if len(sleeps) != 2 || sleeps[0] != 1*time.Second || sleeps[1] != 2*time.Second {
+		t.Fatalf("expected sleeps [1s 2s], got %v", sleeps)
+	}
+}
+
+func TestSweepTransientRecovers(t *testing.T) {
+	orig := confirmSleep
+	confirmSleep = func(d time.Duration) {}
+	defer func() { confirmSleep = orig }()
+
+	var calls int32
+	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
+		n := atomic.AddInt32(&calls, 1)
+		if n == 1 {
+			return Unconfirmed, true, nil // initial parallel probe: 428
+		}
+		return Present, false, nil // serial re-probe: 206
+	}}
+	res, err := Sweep(context.Background(), b, []string{"aaa"}, sweepTestCfg(), io.Discard)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(res.Present) != 1 || res.Present[0] != "aaa" {
+		t.Fatalf("got %+v", res)
+	}
+}
+
+func TestSweepConfirmMaxCap(t *testing.T) {
+	orig := confirmSleep
+	var sleepCalls int32
+	confirmSleep = func(d time.Duration) { atomic.AddInt32(&sleepCalls, 1) }
+	defer func() { confirmSleep = orig }()
+
+	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
+		return Unconfirmed, true, nil
+	}}
+	shas := []string{"a", "b", "c", "d", "e"}
+	cfg := sweepTestCfg()
+	cfg.ConfirmMax = 2
+	res, err := Sweep(context.Background(), b, shas, cfg, io.Discard)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(res.Unconfirmed) != 5 {
+		t.Fatalf("expected all 5 unconfirmed, got %+v", res)
+	}
+	if len(res.Missing) != 0 || len(res.Present) != 0 {
+		t.Fatalf("got %+v", res)
+	}
+	if sleepCalls != 0 {
+		t.Fatalf("expected zero serial re-probes/sleeps when over ConfirmMax, got %d calls", sleepCalls)
+	}
+}
+
+func TestSweepParallelBounded(t *testing.T) {
+	var inFlight int32
+	var maxInFlight int32
+	shas := make([]string, 100)
+	for i := range shas {
+		shas[i] = string(rune('a'+i%26)) + string(rune('A'+i/26))
+	}
+	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
+		n := atomic.AddInt32(&inFlight, 1)
+		for {
+			cur := atomic.LoadInt32(&maxInFlight)
+			if n <= cur || atomic.CompareAndSwapInt32(&maxInFlight, cur, n) {
+				break
+			}
+		}
+		time.Sleep(time.Millisecond)
+		atomic.AddInt32(&inFlight, -1)
+		return Present, false, nil
+	}}
+	cfg := sweepTestCfg()
+	cfg.ProbeJobs = 4
+	_, err := Sweep(context.Background(), b, shas, cfg, io.Discard)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if maxInFlight > 4 {
+		t.Fatalf("expected max concurrency <= 4, got %d", maxInFlight)
+	}
+}