sow-tools/internal/depot/run.go

package depot

import (
	"context"
	"crypto/sha256"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"
)

const (
	exitOK          = 0
	exitDrift       = 1
	exitUnconfirmed = 2
	exitUsage       = 64
	exitInternal    = 70
)

// Run executes a depot subcommand. args[0] is the subcommand
// (status|push|verify|get|pull); returns the process exit code.
func Run(args []string, stdout, stderr io.Writer, getenv func(string) string) int {
	if len(args) == 0 {
		printRunUsage(stderr)
		return exitUsage
	}
	rest := args[1:]
	switch args[0] {
	case "status":
		return runStatus(rest, stdout, stderr, getenv)
	case "push":
		return runPush(rest, stdout, stderr, getenv)
	case "verify":
		return runVerify(rest, stdout, stderr, getenv)
	case "get":
		return runGet(rest, stdout, stderr, getenv)
	case "pull":
		return runPull(rest, stdout, stderr, getenv)
	default:
		fmt.Fprintf(stderr, "depot: unknown subcommand %q\n\n", args[0])
		printRunUsage(stderr)
		return exitUsage
	}
}

func printRunUsage(w io.Writer) {
	fmt.Fprint(w, `usage:
  depot status [--manifests DIR] [--source LOCAL_DEPOT] --target bunny|cdn|local
  depot push   [--manifests DIR] --source LOCAL_DEPOT   --target bunny
  depot verify [--manifests DIR] --target bunny|cdn     [--sample N]
  depot get    <sha> <dest>       --target cdn|bunny|local
  depot pull   [--manifests DIR] --dest DIR             --target cdn|bunny

--target local uses the DEPOT_DIR environment variable as the local root.
`)
}

// resolveBackend builds the named backend, resolving "local" against DEPOT_DIR.
func resolveBackend(target string, getenv func(string) string, cfg Config) (Backend, error) {
	root := ""
	if target == "local" {
		root = getenv("DEPOT_DIR")
		if root == "" {
			return nil, fmt.Errorf("--target local requires DEPOT_DIR to be set")
		}
	}
	return NewBackend(target, root, cfg)
}

func printSweepStatus(stdout io.Writer, referenced int, res SweepResult) {
	fmt.Fprintf(stdout, "referenced=%d present=%d missing=%d unconfirmed=%d\n",
		referenced, len(res.Present), len(res.Missing), len(res.Unconfirmed))
	for _, sha := range res.Missing {
		fmt.Fprintf(stdout, "missing %s\n", sha)
	}
	for _, sha := range res.Unconfirmed {
		fmt.Fprintf(stdout, "unconfirmed %s\n", sha)
	}
}

func sweepExitCode(res SweepResult) int {
	if len(res.Missing) > 0 {
		return exitDrift
	}
	if len(res.Unconfirmed) > 0 {
		return exitUnconfirmed
	}
	return exitOK
}

func shaKeys(m map[string]int64) []string {
	out := make([]string, 0, len(m))
	for sha := range m {
		out = append(out, sha)
	}
	sort.Strings(out)
	return out
}

func runStatus(args []string, stdout, stderr io.Writer, getenv func(string) string) int {
	fs := flag.NewFlagSet("status", flag.ContinueOnError)
	fs.SetOutput(stderr)
	manifests := fs.String("manifests", "assets", "directory of *.yml manifests")
	_ = fs.String("source", "", "local depot root (unused for status target=local; see DEPOT_DIR)")
	target := fs.String("target", "", "bunny|cdn|local")
	if err := fs.Parse(args); err != nil {
		return exitUsage
	}
	if *target == "" {
		fmt.Fprintln(stderr, "depot status: --target is required")
		return exitUsage
	}

	cfg := LoadConfig(getenv)
	backend, err := resolveBackend(*target, getenv, cfg)
	if err != nil {
		fmt.Fprintln(stderr, "depot status:", err)
		return exitInternal
	}

	shaSizes, err := ReferencedSHAs(*manifests)
	if err != nil {
		fmt.Fprintln(stderr, "depot status:", err)
		return exitInternal
	}
	shas := shaKeys(shaSizes)

	res, err := Sweep(context.Background(), backend, shas, cfg, io.Discard)
	if err != nil {
		fmt.Fprintln(stderr, "depot status:", err)
		return exitInternal
	}

	printSweepStatus(stdout, len(shas), res)
	return sweepExitCode(res)
}

func runPush(args []string, stdout, stderr io.Writer, getenv func(string) string) int {
	fs := flag.NewFlagSet("push", flag.ContinueOnError)
	fs.SetOutput(stderr)
	manifests := fs.String("manifests", "assets", "directory of *.yml manifests")
	source := fs.String("source", "", "local depot root to read blobs from")
	target := fs.String("target", "", "must be bunny")
	if err := fs.Parse(args); err != nil {
		return exitUsage
	}
	if *target != "bunny" {
		fmt.Fprintln(stderr, "depot push: --target must be bunny")
		return exitUsage
	}
	if *source == "" {
		fmt.Fprintln(stderr, "depot push: --source is required")
		return exitUsage
	}

	cfg := LoadConfig(getenv)
	backend, err := NewBackend("bunny", "", cfg)
	if err != nil {
		fmt.Fprintln(stderr, "depot push:", err)
		return exitInternal
	}

	shaSizes, err := ReferencedSHAs(*manifests)
	if err != nil {
		fmt.Fprintln(stderr, "depot push:", err)
		return exitInternal
	}
	shas := shaKeys(shaSizes)

	res, err := Sweep(context.Background(), backend, shas, cfg, io.Discard)
	if err != nil {
		fmt.Fprintln(stderr, "depot push:", err)
		return exitInternal
	}

	toUpload := append(append([]string(nil), res.Missing...), res.Unconfirmed...)
	sort.Strings(toUpload)

	var (
		mu           sync.Mutex
		uploaded     int
		failed       int
		transportErr error
	)

	work := func(sha string) {
		srcPath := filepath.Join(*source, BlobKey(sha))
		if _, statErr := os.Stat(srcPath); statErr != nil {
			mu.Lock()
			failed++
			mu.Unlock()
			fmt.Fprintf(stderr, "push: missing source blob %s\n", sha)
			return
		}
		if err := backend.Put(context.Background(), sha, srcPath); err != nil {
			mu.Lock()
			if transportErr == nil {
				transportErr = err
			}
			mu.Unlock()
			fmt.Fprintf(stderr, "push: upload %s: %v\n", sha, err)
			return
		}
		mu.Lock()
		uploaded++
		mu.Unlock()
	}

	runWorkers(cfg.Jobs, toUpload, work)

	fmt.Fprintf(stdout, "uploaded=%d failed=%d\n", uploaded, failed)

	if transportErr != nil {
		return exitInternal
	}
	if failed > 0 {
		return exitDrift
	}
	return exitOK
}

func runVerify(args []string, stdout, stderr io.Writer, getenv func(string) string) int {
	fs := flag.NewFlagSet("verify", flag.ContinueOnError)
	fs.SetOutput(stderr)
	manifests := fs.String("manifests", "assets", "directory of *.yml manifests")
	target := fs.String("target", "", "bunny|cdn")
	sample := fs.Int("sample", 3, "number of present blobs to spot-check by download (0=skip)")
	if err := fs.Parse(args); err != nil {
		return exitUsage
	}
	if *target != "bunny" && *target != "cdn" {
		fmt.Fprintln(stderr, "depot verify: --target must be bunny or cdn")
		return exitUsage
	}

	cfg := LoadConfig(getenv)
	backend, err := NewBackend(*target, "", cfg)
	if err != nil {
		fmt.Fprintln(stderr, "depot verify:", err)
		return exitInternal
	}

	shaSizes, err := ReferencedSHAs(*manifests)
	if err != nil {
		fmt.Fprintln(stderr, "depot verify:", err)
		return exitInternal
	}
	shas := shaKeys(shaSizes)

	res, err := Sweep(context.Background(), backend, shas, cfg, io.Discard)
	if err != nil {
		fmt.Fprintln(stderr, "depot verify:", err)
		return exitInternal
	}

	printSweepStatus(stdout, len(shas), res)
	if code := sweepExitCode(res); code != exitOK {
		return code
	}

	if *sample <= 0 || len(res.Present) == 0 {
		return exitOK
	}

	n := *sample
	if n > len(res.Present) {
		n = len(res.Present)
	}
	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
	picks := rng.Perm(len(res.Present))[:n]

	tmpDir, err := os.MkdirTemp("", "depot-verify-")
	if err != nil {
		fmt.Fprintln(stderr, "depot verify:", err)
		return exitInternal
	}
	defer os.RemoveAll(tmpDir)

	for _, idx := range picks {
		sha := res.Present[idx]
		dest := filepath.Join(tmpDir, sha)
		if err := backend.Get(context.Background(), sha, dest); err != nil {
			fmt.Fprintf(stderr, "depot verify: sample %s: %v\n", sha, err)
			return exitDrift
		}
	}
	return exitOK
}

func runGet(args []string, stdout, stderr io.Writer, getenv func(string) string) int {
	fs := flag.NewFlagSet("get", flag.ContinueOnError)
	fs.SetOutput(stderr)
	target := fs.String("target", "", "cdn|bunny|local")
	if err := fs.Parse(args); err != nil {
		return exitUsage
	}
	positional := fs.Args()
	if len(positional) != 2 {
		fmt.Fprintln(stderr, "depot get: usage: depot get <sha> <dest> --target cdn|bunny|local")
		return exitUsage
	}
	sha, dest := positional[0], positional[1]
	if !ValidSHA(sha) {
		fmt.Fprintf(stderr, "depot get: invalid sha %q\n", sha)
		return exitUsage
	}
	if *target == "" {
		fmt.Fprintln(stderr, "depot get: --target is required")
		return exitUsage
	}

	cfg := LoadConfig(getenv)
	backend, err := resolveBackend(*target, getenv, cfg)
	if err != nil {
		fmt.Fprintln(stderr, "depot get:", err)
		return exitInternal
	}

	if err := backend.Get(context.Background(), sha, dest); err != nil {
		fmt.Fprintln(stderr, "depot get:", err)
		return exitInternal
	}
	return exitOK
}

func runPull(args []string, stdout, stderr io.Writer, getenv func(string) string) int {
	fs := flag.NewFlagSet("pull", flag.ContinueOnError)
	fs.SetOutput(stderr)
	manifests := fs.String("manifests", "assets", "directory of *.yml manifests")
	dest := fs.String("dest", "", "local directory to pull blobs into")
	target := fs.String("target", "", "cdn|bunny")
	if err := fs.Parse(args); err != nil {
		return exitUsage
	}
	if *target != "bunny" && *target != "cdn" {
		fmt.Fprintln(stderr, "depot pull: --target must be bunny or cdn")
		return exitUsage
	}
	if *dest == "" {
		fmt.Fprintln(stderr, "depot pull: --dest is required")
		return exitUsage
	}

	cfg := LoadConfig(getenv)
	backend, err := NewBackend(*target, "", cfg)
	if err != nil {
		fmt.Fprintln(stderr, "depot pull:", err)
		return exitInternal
	}

	shaSizes, err := ReferencedSHAs(*manifests)
	if err != nil {
		fmt.Fprintln(stderr, "depot pull:", err)
		return exitInternal
	}
	shas := shaKeys(shaSizes)

	// Split into "already present and correct locally" (skip) vs "needs a
	// probe/download decision".
	var present int
	var toCheck []string
	for _, sha := range shas {
		destPath := filepath.Join(*dest, BlobKey(sha))
		if localFileMatchesSHA(destPath, sha) {
			present++
			continue
		}
		toCheck = append(toCheck, sha)
	}

	// Sweep the source to distinguish "not there" (exit 1, listed) from
	// "there, download it".
	res, err := Sweep(context.Background(), backend, toCheck, cfg, io.Discard)
	if err != nil {
		fmt.Fprintln(stderr, "depot pull:", err)
		return exitInternal
	}

	var (
		mu      sync.Mutex
		pulled  int
		anyFail bool
	)
	work := func(sha string) {
		destPath := filepath.Join(*dest, BlobKey(sha))
		if err := backend.Get(context.Background(), sha, destPath); err != nil {
			mu.Lock()
			anyFail = true
			mu.Unlock()
			fmt.Fprintf(stderr, "pull: %s: %v\n", sha, err)
			return
		}
		mu.Lock()
		pulled++
		mu.Unlock()
	}
	runWorkers(cfg.Jobs, res.Present, work)

	// ponytail: unconfirmed source state (probe budget exhausted / flaky
	// responses) is treated as a download failure rather than a third exit
	// path; if that proves too coarse in practice, give pull its own
	// unconfirmed accounting like status.
	if len(res.Unconfirmed) > 0 {
		anyFail = true
		for _, sha := range res.Unconfirmed {
			fmt.Fprintf(stderr, "pull: %s: unconfirmed at source\n", sha)
		}
	}

	for _, sha := range res.Missing {
		fmt.Fprintf(stdout, "missing %s\n", sha)
	}

	fmt.Fprintf(stdout, "pulled=%d present=%d\n", pulled, present)

	if anyFail {
		return exitInternal
	}
	if len(res.Missing) > 0 {
		return exitDrift
	}
	return exitOK
}

// localFileMatchesSHA reports whether path exists and hashes to sha.
func localFileMatchesSHA(path, sha string) bool {
	f, err := os.Open(path)
	if err != nil {
		return false
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return false
	}
	return fmt.Sprintf("%x", h.Sum(nil)) == sha
}

// runWorkers runs work(sha) for each sha in items using up to jobs goroutines.
func runWorkers(jobs int, items []string, work func(sha string)) {
	if jobs < 1 {
		jobs = 1
	}
	ch := make(chan string)
	var wg sync.WaitGroup
	for i := 0; i < jobs; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for sha := range ch {
				work(sha)
			}
		}()
	}
	for _, sha := range items {
		ch <- sha
	}
	close(ch)
	wg.Wait()
}