sow-tools/internal/depot/sweep.go

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
}