at-container-registry/internal/reqcount/reqcount.go

// Package reqcount counts HTTP requests per backend and can inject a fixed
// latency in front of each, so an in-process benchmark of the ATCR stack can
// report round trips per operation. Round-trip counts are the reliable signal
// in-process: with no latency injected, wall time is dominated by gofakes3's
// in-memory copies and crane's client-side gzip, not by ATCR code.
package reqcount

import (
	"net/http"
	"sort"
	"strings"
	"sync"
	"time"
)

// Tap wraps handlers and counts what passes through them.
type Tap struct {
	mu      sync.Mutex
	counts  map[string]map[string]int
	latency map[string]time.Duration
}

// New returns an empty Tap.
func New() *Tap {
	return &Tap{
		counts:  make(map[string]map[string]int),
		latency: make(map[string]time.Duration),
	}
}

// SetLatency makes every request to backend sleep for d before it is served.
func (t *Tap) SetLatency(backend string, d time.Duration) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.latency[backend] = d
}

// Wrap returns next wrapped with counting and latency for backend.
func (t *Tap) Wrap(backend string, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		key := classify(backend, r)
		t.mu.Lock()
		m := t.counts[backend]
		if m == nil {
			m = make(map[string]int)
			t.counts[backend] = m
		}
		m[key]++
		d := t.latency[backend]
		t.mu.Unlock()
		if d > 0 {
			time.Sleep(d)
		}
		next.ServeHTTP(w, r)
	})
}

// Reset clears every count.
func (t *Tap) Reset() {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.counts = make(map[string]map[string]int)
}

// Snapshot returns a copy of the counts: backend -> key -> requests.
func (t *Tap) Snapshot() map[string]map[string]int {
	t.mu.Lock()
	defer t.mu.Unlock()
	out := make(map[string]map[string]int, len(t.counts))
	for b, m := range t.counts {
		c := make(map[string]int, len(m))
		for k, n := range m {
			c[k] = n
		}
		out[b] = c
	}
	return out
}

// Total sums one backend's counts.
func Total(m map[string]int) int {
	n := 0
	for _, v := range m {
		n += v
	}
	return n
}

// Keys returns a backend's keys sorted by count, largest first.
func Keys(m map[string]int) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Slice(keys, func(i, j int) bool {
		if m[keys[i]] != m[keys[j]] {
			return m[keys[i]] > m[keys[j]]
		}
		return keys[i] < keys[j]
	})
	return keys
}

// classify turns a request into a short, stable key for its backend.
func classify(backend string, r *http.Request) string {
	q := r.URL.Query()
	switch backend {
	case "s3":
		switch {
		case q.Has("uploads"):
			return "POST initiate"
		case q.Has("uploadId") && q.Has("partNumber"):
			return "PUT part"
		case q.Has("uploadId") && r.Method == http.MethodPost:
			return "POST complete"
		case q.Has("uploadId") && r.Method == http.MethodDelete:
			return "DELETE abort"
		case r.Header.Get("X-Amz-Copy-Source") != "":
			return "PUT copy"
		}
		return r.Method + " object"
	case "hold":
		p := strings.TrimPrefix(r.URL.Path, "/xrpc/")
		if strings.HasSuffix(p, "sync.getBlob") && q.Get("method") != "" {
			p += "?method=" + q.Get("method")
		}
		return r.Method + " " + p
	default:
		return r.Method + " " + strings.TrimPrefix(r.URL.Path, "/xrpc/")
	}
}

// Transport wraps an http.RoundTripper so client-side requests are counted
// under backend, keyed by method and a coarse path class. It is meant for the
// registry itself, which sits in front of the tapped backends.
func (t *Tap) Transport(backend string, base http.RoundTripper) http.RoundTripper {
	if base == nil {
		base = http.DefaultTransport
	}
	return roundTripFunc(func(r *http.Request) (*http.Response, error) {
		t.mu.Lock()
		m := t.counts[backend]
		if m == nil {
			m = make(map[string]int)
			t.counts[backend] = m
		}
		m[r.Method+" "+pathClass(r.URL.Path)]++
		t.mu.Unlock()
		return base.RoundTrip(r)
	})
}

type roundTripFunc func(*http.Request) (*http.Response, error)

func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }

func pathClass(p string) string {
	switch {
	case !strings.HasPrefix(p, "/v2/") && !strings.HasPrefix(p, "/auth/"):
		// A redirect the client followed, to S3 or elsewhere.
		return "redirect"
	case strings.Contains(p, "/blobs/uploads"):
		return "blobs/uploads"
	case strings.Contains(p, "/blobs/"):
		return "blobs"
	case strings.Contains(p, "/manifests/"):
		return "manifests"
	case strings.HasPrefix(p, "/auth/"):
		return "auth"
	case p == "/v2/" || p == "/v2":
		return "ping"
	}
	return "other"
}