at-container-registry/test/integration/bench_real_test.go

//go:build integration

package integration

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/google/go-containerregistry/pkg/crane"
	"github.com/google/go-containerregistry/pkg/name"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/empty"
	"github.com/google/go-containerregistry/pkg/v1/mutate"
	"github.com/google/go-containerregistry/pkg/v1/random"
	"github.com/google/go-containerregistry/pkg/v1/types"

	"atcr.io/internal/reqcount"
	"atcr.io/internal/testharness"
)

// Layer sizes in MB copied from three real manifests in the production
// appview database (2026-09-10): the median image, the p75 image, and the
// p90 image by layer count.
var benchProfiles = map[string][]float64{
	"median": {3.63, 1.29, 0.01, 36.32, 1.64, 0.34},
	"p75":    {3.95, 0.44, 13.09, 0.01, 17.34, 48.04, 0.15, 0.02, 0.01, 0.15, 0.01},
	"p90":    {3.47, 47.69, 1.20, 0.01, 1.82, 0.01, 0.03, 0.01, 0.01, 0.01, 0.01, 27.52, 15.71, 0.01, 0.07, 0.01, 0.01, 0.02, 0.01, 0.01, 0.01, 0.01},
}

// TestBenchRealImages pushes and pulls production-shaped images through the
// in-process stack and reports round trips per backend and wall time.
//
// Skipped unless BENCH_PROFILES is set (a comma list of profile names, or
// "all"). BENCH_LAT_PDS, BENCH_LAT_HOLD and BENCH_LAT_S3 inject per-request
// latency in milliseconds. BENCH_OUT appends one JSON line per operation.
//
//	BENCH_PROFILES=all BENCH_LAT_PDS=50 BENCH_LAT_HOLD=5 BENCH_LAT_S3=20 \
//	  go test -tags=integration,testmode -run TestBenchRealImages -v -count=1 ./test/integration/
func TestBenchRealImages(t *testing.T) {
	want := os.Getenv("BENCH_PROFILES")
	if want == "" {
		t.Skip("set BENCH_PROFILES to run")
	}
	var names []string
	if want == "all" {
		names = []string{"median", "p75", "p90"}
	} else {
		names = strings.Split(want, ",")
	}

	tap := reqcount.New()
	for _, b := range []string{"pds", "hold", "s3"} {
		if ms, _ := strconv.Atoi(os.Getenv("BENCH_LAT_" + strings.ToUpper(b))); ms > 0 {
			tap.SetLatency(b, time.Duration(ms)*time.Millisecond)
		}
	}

	h := testharness.New(t, testharness.WithBackendTap(func(backend string, next http.Handler) http.Handler {
		return tap.Wrap(backend, next)
	}))
	// crane, with a counting transport so client-to-registry requests are
	// reported under "registry" next to the backends.
	craneOpts := func(creds testharness.Auth) []crane.Option {
		return []crane.Option{
			crane.WithAuth(toAuthn(creds)),
			crane.Insecure,
			crane.WithTransport(tap.Transport("registry", nil)),
		}
	}

	var out *os.File
	if p := os.Getenv("BENCH_OUT"); p != "" {
		f, err := os.OpenFile(p, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
		if err != nil {
			t.Fatalf("open BENCH_OUT: %v", err)
		}
		defer f.Close()
		out = f
	}

	for _, profile := range names {
		sizes, ok := benchProfiles[profile]
		if !ok {
			t.Fatalf("unknown profile %q", profile)
		}
		sailor := h.AddSailor(profile + ".test")
		creds := h.RegistryCreds(sailor)

		// Warm the per-user caches (hold prefs row, service token) with a
		// tiny push so the measured ops see steady state, like a returning
		// user's second push.
		warm, err := random.Image(1<<10, 1)
		if err != nil {
			t.Fatal(err)
		}
		warmRef, _ := name.ParseReference(fmt.Sprintf("%s/%s/warm:latest", h.AppViewHostPort(), sailor.Handle()), name.Insecure)
		if err := crane.Push(warm, warmRef.String(), craneOpts(creds)...); err != nil {
			t.Fatalf("warm push: %v", err)
		}

		img := buildImage(t, sizes)
		ref, _ := name.ParseReference(fmt.Sprintf("%s/%s/bench:%s", h.AppViewHostPort(), sailor.Handle(), profile), name.Insecure)

		tap.Reset()
		start := time.Now()
		if err := crane.Push(img, ref.String(), craneOpts(creds)...); err != nil {
			t.Fatalf("%s push: %v", profile, err)
		}
		report(t, out, profile, "push", time.Since(start), tap.Snapshot())

		tap.Reset()
		start = time.Now()
		if err := pullAll(ref.String(), craneOpts(creds)); err != nil {
			t.Fatalf("%s pull: %v", profile, err)
		}
		report(t, out, profile, "pull", time.Since(start), tap.Snapshot())
	}
}

// pullAll fetches the manifest and reads every layer's bytes, the same shape
// as the crane client in clients.go.
func pullAll(ref string, opts []crane.Option) error {
	img, err := crane.Pull(ref, opts...)
	if err != nil {
		return err
	}
	layers, err := img.Layers()
	if err != nil {
		return err
	}
	for _, l := range layers {
		rc, err := l.Compressed()
		if err != nil {
			return err
		}
		_, cerr := io.Copy(io.Discard, rc)
		rc.Close()
		if cerr != nil {
			return cerr
		}
	}
	return nil
}

func buildImage(t *testing.T, sizesMB []float64) v1.Image {
	t.Helper()
	layers := make([]v1.Layer, 0, len(sizesMB))
	for _, mb := range sizesMB {
		l, err := random.Layer(int64(mb*1024*1024), types.DockerLayer)
		if err != nil {
			t.Fatalf("random layer: %v", err)
		}
		layers = append(layers, l)
	}
	img, err := mutate.AppendLayers(empty.Image, layers...)
	if err != nil {
		t.Fatalf("append layers: %v", err)
	}
	return img
}

func report(t *testing.T, out *os.File, profile, op string, wall time.Duration, snap map[string]map[string]int) {
	t.Helper()
	pds, hold, s3 := reqcount.Total(snap["pds"]), reqcount.Total(snap["hold"]), reqcount.Total(snap["s3"])
	reg := reqcount.Total(snap["registry"])
	t.Logf("%-6s %-4s %6dms  registry=%-3d pds=%-3d hold=%-3d s3=%-3d", profile, op, wall.Milliseconds(), reg, pds, hold, s3)
	for _, b := range []string{"registry", "hold", "s3", "pds"} {
		for _, k := range reqcount.Keys(snap[b]) {
			t.Logf("    %-4s %4d  %s", b, snap[b][k], k)
		}
	}
	if out != nil {
		_ = json.NewEncoder(out).Encode(map[string]any{
			"profile": profile, "op": op, "wall_ms": wall.Milliseconds(),
			"registry": reg, "pds": pds, "hold": hold, "s3": s3, "detail": snap,
		})
	}
}