diff --git a/internal/depot/sweep.go b/internal/depot/sweep.go new file mode 100644 index 0000000..81a3d20 --- /dev/null +++ b/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 +} diff --git a/internal/depot/sweep_test.go b/internal/depot/sweep_test.go new file mode 100644 index 0000000..78fdc1e --- /dev/null +++ b/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) + } +}