tiny-proxy/internal/proxycache/proxycache_test.go

package proxycache

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"tiny-proxy/internal/metrics"
	"tiny-proxy/internal/storage"
)

// markVRepo binds req to a vrepo's stats, standing in for the
// metrics.MarkVRepo a real handler makes once it has resolved the repo.
func markVRepo(req *http.Request, stats []*metrics.VRepoStats) *http.Request {
	if len(stats) == 0 {
		return req
	}
	return req.WithContext(metrics.WithVRepo(req.Context(), stats[0]))
}

// testVRepo registers one vrepo and returns its stats.
func testVRepo() *metrics.VRepoStats {
	return metrics.NewRegistry().RegisterVRepo("v", "maven")
}

func testCache(t *testing.T) *Cache {
	t.Helper()
	store, err := storage.NewLocalFS(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	return New(store)
}

// get serves key through the cache. The optional stats binds the request to a
// vrepo the way a handler's MarkVRepo does, so the cache counters have
// somewhere to land; tests that do not look at metrics pass nothing.
func get(t *testing.T, c *Cache, upstream http.Handler, key string, ttl time.Duration, stats ...*metrics.VRepoStats) (int, string) {
	t.Helper()
	return getPolicy(t, c, upstream, key, ttl, false, stats...)
}

// getPolicy is get with an explicit NoStore policy on the fetch result.
func getPolicy(t *testing.T, c *Cache, upstream http.Handler, key string, ttl time.Duration, noStore bool, stats ...*metrics.VRepoStats) (int, string) {
	t.Helper()
	rec := httptest.NewRecorder()
	req := markVRepo(httptest.NewRequest(http.MethodGet, "/", nil), stats)
	c.Serve(rec, req, key, ttl, func(ctx context.Context) (Result, error) {
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "", nil)
		rec2 := httptest.NewRecorder()
		upstream.ServeHTTP(rec2, req)
		return Result{Resp: rec2.Result(), NoStore: noStore}, nil
	})
	res := rec.Result()
	body, _ := io.ReadAll(res.Body)
	return res.StatusCode, string(body)
}

func TestMissThenHit(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		w.Header().Set("Content-Type", "text/plain")
		fmt.Fprint(w, "payload")
	})
	c := testCache(t)

	code, body := get(t, c, upstream, "k", 0)
	if code != 200 || body != "payload" {
		t.Fatalf("miss: code=%d body=%q", code, body)
	}
	code, body = get(t, c, upstream, "k", 0)
	if code != 200 || body != "payload" {
		t.Fatalf("hit: code=%d body=%q", code, body)
	}
	if hits.Load() != 1 {
		t.Errorf("upstream hits = %d, want 1 (second request should be cached)", hits.Load())
	}
}

func TestTTLRefetch(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "v%d", hits.Add(1))
	})
	c := testCache(t)

	// Seed the cache with an object stored an hour ago.
	if err := c.Store.Put(t.Context(), "k", strings.NewReader("v0"),
		storage.ObjectMeta{StoredAt: time.Now().Add(-time.Hour)}); err != nil {
		t.Fatal(err)
	}
	_, body := get(t, c, upstream, "k", time.Minute)
	if body != "v1" {
		t.Errorf("expired entry should refetch, body = %q", body)
	}
	// The refetch refreshed the entry: next request is served from cache.
	_, body = get(t, c, upstream, "k", time.Minute)
	if body != "v1" || hits.Load() != 1 {
		t.Errorf("refreshed entry should hit cache, body = %q hits = %d", body, hits.Load())
	}
}

// TestNoStoreServesWithoutCaching covers upstream cache: false — the body
// reaches the client on every request but never lands in storage.
func TestNoStoreServesWithoutCaching(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		fmt.Fprint(w, "payload")
	})
	c := testCache(t)

	for i := 0; i < 2; i++ {
		code, body := getPolicy(t, c, upstream, "k", 0, true)
		if code != http.StatusOK || body != "payload" {
			t.Fatalf("request %d: code = %d body = %q, want 200 %q", i, code, body, "payload")
		}
	}
	if got := hits.Load(); got != 2 {
		t.Errorf("upstream hits = %d, want 2 (nothing may be cached)", got)
	}
	if _, err := c.Store.Stat(t.Context(), "k"); err != storage.ErrNotFound {
		t.Error("no-store response must not be written to storage")
	}
}

// TestNoStoreStillServesCachedCopy: the flag suppresses writes, not reads, so
// an object cached earlier (e.g. by another upstream of the same vrepo) is
// still served without touching the upstream.
func TestNoStoreStillServesCachedCopy(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		fmt.Fprint(w, "payload")
	})
	c := testCache(t)

	get(t, c, upstream, "k", 0) // cached by a normal (cacheable) upstream
	if _, body := getPolicy(t, c, upstream, "k", 0, true); body != "payload" {
		t.Errorf("body = %q, want %q", body, "payload")
	}
	if got := hits.Load(); got != 1 {
		t.Errorf("upstream hits = %d, want 1 (second request served from cache)", got)
	}
}

func TestFetchNoStore(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		fmt.Fprint(w, "value")
	})
	c := testCache(t)

	for i := 0; i < 2; i++ {
		body, err := fetchBodyPolicy(c, "k", time.Hour, upstream, true)
		if err != nil || string(body) != "value" {
			t.Fatalf("request %d: body = %q err = %v", i, body, err)
		}
	}
	if got := hits.Load(); got != 2 {
		t.Errorf("upstream hits = %d, want 2", got)
	}
	if _, err := c.Store.Stat(context.Background(), "k"); err != storage.ErrNotFound {
		t.Error("no-store fetch must not be written to storage")
	}
}

func TestUpstreamErrorNotCached(t *testing.T) {
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "gone", http.StatusNotFound)
	})
	c := testCache(t)
	code, _ := get(t, c, upstream, "k", 0)
	if code != http.StatusNotFound {
		t.Errorf("code = %d, want 404 passthrough", code)
	}
	if _, err := c.Store.Stat(t.Context(), "k"); err != storage.ErrNotFound {
		t.Error("404 response should not be cached")
	}
}

func fetchBody(c *Cache, key string, ttl time.Duration, upstream http.Handler, stats ...*metrics.VRepoStats) ([]byte, error) {
	return fetchBodyPolicy(c, key, ttl, upstream, false, stats...)
}

func fetchBodyPolicy(c *Cache, key string, ttl time.Duration, upstream http.Handler, noStore bool, stats ...*metrics.VRepoStats) ([]byte, error) {
	ctx := context.Background()
	if len(stats) > 0 {
		ctx = metrics.WithVRepo(ctx, stats[0])
	}
	return c.Fetch(ctx, key, ttl, func(ctx context.Context) (Result, error) {
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "", nil)
		rec := httptest.NewRecorder()
		upstream.ServeHTTP(rec, req)
		return Result{Resp: rec.Result(), NoStore: noStore}, nil
	})
}

func TestFetchMissThenHit(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		fmt.Fprint(w, "value")
	})
	c := testCache(t)

	body, err := fetchBody(c, "k", 0, upstream)
	if err != nil || string(body) != "value" {
		t.Fatalf("miss: body=%q err=%v", body, err)
	}
	body, err = fetchBody(c, "k", 0, upstream)
	if err != nil || string(body) != "value" {
		t.Fatalf("hit: body=%q err=%v", body, err)
	}
	if hits.Load() != 1 {
		t.Errorf("upstream hits = %d, want 1", hits.Load())
	}
}

func TestFetchTTLRefetch(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "v%d", hits.Add(1))
	})
	c := testCache(t)
	if err := c.Store.Put(t.Context(), "k", strings.NewReader("v0"),
		storage.ObjectMeta{StoredAt: time.Now().Add(-time.Hour)}); err != nil {
		t.Fatal(err)
	}
	body, err := fetchBody(c, "k", time.Minute, upstream)
	if err != nil || string(body) != "v1" {
		t.Errorf("expired entry should refetch, body=%q err=%v", body, err)
	}
	body, _ = fetchBody(c, "k", time.Minute, upstream)
	if string(body) != "v1" || hits.Load() != 1 {
		t.Errorf("refreshed entry should hit cache, body=%q hits=%d", body, hits.Load())
	}
}

func TestFetchNon200NotCached(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		http.Error(w, "not here", http.StatusNotFound)
	})
	c := testCache(t)
	_, err := fetchBody(c, "k", 0, upstream)
	var ue *UpstreamError
	if !errors.As(err, &ue) || ue.StatusCode != http.StatusNotFound {
		t.Fatalf("err = %v, want *UpstreamError with 404", err)
	}
	if _, serr := c.Store.Stat(t.Context(), "k"); serr != storage.ErrNotFound {
		t.Error("non-200 response should not be cached")
	}
	// Next Fetch retries upstream.
	if _, err := fetchBody(c, "k", 0, upstream); err == nil || hits.Load() != 2 {
		t.Errorf("hits = %d, want 2 (no caching of failures)", hits.Load())
	}
}

// failingStore fails every Put immediately, without reading the body.
type failingStore struct{ storage.Store }

func (failingStore) Put(context.Context, string, io.Reader, storage.ObjectMeta) error {
	return errors.New("storage is read-only")
}

func TestMetricsMissThenHit(t *testing.T) {
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "payload")
	})
	c := testCache(t)
	stats := testVRepo()

	get(t, c, upstream, "k", 0, stats)
	get(t, c, upstream, "k", 0, stats)
	hits, misses, errors := stats.Snapshot()
	if hits != 1 || misses != 1 || errors != 0 {
		t.Errorf("snapshot = hits:%d misses:%d errors:%d, want 1/1/0", hits, misses, errors)
	}
}

func TestMetricsFetchError(t *testing.T) {
	c := testCache(t)
	stats := testVRepo()
	rec := httptest.NewRecorder()
	req := markVRepo(httptest.NewRequest(http.MethodGet, "/", nil), []*metrics.VRepoStats{stats})
	c.Serve(rec, req, "k", 0, func(ctx context.Context) (Result, error) {
		return Result{}, errors.New("upstream down")
	})
	_, misses, errors := stats.Snapshot()
	if misses != 1 || errors != 1 {
		t.Errorf("snapshot misses:%d errors:%d, want 1/1", misses, errors)
	}
}

func TestMetricsStoreError(t *testing.T) {
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "payload")
	})
	c := New(failingStore{testCache(t).Store})
	stats := testVRepo()
	get(t, c, upstream, "k", 0, stats)
	_, misses, errors := stats.Snapshot()
	if misses != 1 || errors != 1 {
		t.Errorf("snapshot misses:%d errors:%d, want 1/1", misses, errors)
	}
}

func TestMetricsFetchPath(t *testing.T) {
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "value")
	})
	c := testCache(t)
	stats := testVRepo()
	fetchBody(c, "k", 0, upstream, stats)
	fetchBody(c, "k", 0, upstream, stats)
	hits, misses, errors := stats.Snapshot()
	if hits != 1 || misses != 1 || errors != 0 {
		t.Errorf("snapshot = hits:%d misses:%d errors:%d, want 1/1/0", hits, misses, errors)
	}
}

func TestPutFailureDoesNotHang(t *testing.T) {
	payload := strings.Repeat("x", 1<<20) // larger than any pipe buffer
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, payload)
	})
	c := New(failingStore{testCache(t).Store})

	done := make(chan string, 1)
	go func() {
		_, body := get(t, c, upstream, "k", 0)
		done <- body
	}()
	select {
	case body := <-done:
		if body != payload {
			t.Errorf("body length = %d, want %d (client should still get the full body)", len(body), len(payload))
		}
	case <-time.After(10 * time.Second):
		t.Fatal("request hung: Put failing before reading the pipe must not deadlock the client")
	}
}

// serveResult serves key with method, answering misses with a fixed result
// built by res, and counts how often the fetch ran.
func serveResult(c *Cache, method, key string, calls *atomic.Int32, res func() Result) *httptest.ResponseRecorder {
	rec := httptest.NewRecorder()
	c.Serve(rec, httptest.NewRequest(method, "/", nil), key, 0, func(ctx context.Context) (Result, error) {
		calls.Add(1)
		return res(), nil
	})
	return rec
}

func notFoundResult(ttl time.Duration) func() Result {
	return func() Result {
		rec := httptest.NewRecorder()
		http.NotFound(rec, nil)
		return Result{Resp: rec.Result(), NotFoundTTL: ttl}
	}
}

func TestNotFoundRemembered(t *testing.T) {
	c := testCache(t)
	var calls atomic.Int32
	for i := 0; i < 3; i++ {
		if rec := serveResult(c, http.MethodGet, "k", &calls, notFoundResult(time.Minute)); rec.Code != http.StatusNotFound {
			t.Fatalf("request %d: got %d, want 404", i, rec.Code)
		}
	}
	if got := calls.Load(); got != 1 {
		t.Errorf("fetch calls = %d, want 1", got)
	}
}

func TestNotFoundWithoutTTLNotRemembered(t *testing.T) {
	c := testCache(t)
	var calls atomic.Int32
	serveResult(c, http.MethodGet, "k", &calls, notFoundResult(0))
	serveResult(c, http.MethodGet, "k", &calls, notFoundResult(0))
	if got := calls.Load(); got != 2 {
		t.Errorf("fetch calls = %d, want 2", got)
	}
}

func TestNotFoundExpires(t *testing.T) {
	c := testCache(t)
	var calls atomic.Int32
	serveResult(c, http.MethodGet, "k", &calls, notFoundResult(20*time.Millisecond))
	time.Sleep(40 * time.Millisecond)
	serveResult(c, http.MethodGet, "k", &calls, notFoundResult(20*time.Millisecond))
	if got := calls.Load(); got != 2 {
		t.Errorf("fetch calls = %d, want 2 once the miss expired", got)
	}
}

func TestFetchNotFoundRemembered(t *testing.T) {
	c := testCache(t)
	var calls atomic.Int32
	for i := 0; i < 2; i++ {
		_, err := c.Fetch(t.Context(), "k", 0, func(ctx context.Context) (Result, error) {
			calls.Add(1)
			return notFoundResult(time.Minute)(), nil
		})
		var ue *UpstreamError
		if !errors.As(err, &ue) || ue.StatusCode != http.StatusNotFound {
			t.Fatalf("call %d: err = %v, want 404 UpstreamError", i, err)
		}
	}
	if got := calls.Load(); got != 1 {
		t.Errorf("fetch calls = %d, want 1", got)
	}
}

func TestBodyForHeadStores(t *testing.T) {
	c := testCache(t)
	var calls atomic.Int32
	ok := func() Result {
		rec := httptest.NewRecorder()
		rec.Header().Set("Content-Length", "4")
		rec.WriteString("data")
		return Result{Resp: rec.Result(), BodyForHead: true}
	}
	rec := httptest.NewRecorder()
	storable := false
	c.ServeStoringHead(rec, httptest.NewRequest(http.MethodHead, "/", nil), "k", 0, func(ctx context.Context) (Result, error) {
		calls.Add(1)
		storable = HeadStorable(ctx)
		return ok(), nil
	})
	if !storable {
		t.Fatal("HeadStorable = false for a HEAD under ServeStoringHead")
	}
	if rec.Code != http.StatusOK || rec.Body.Len() != 0 {
		t.Fatalf("HEAD: got %d with %d body bytes, want 200 and none", rec.Code, rec.Body.Len())
	}
	rec = serveResult(c, http.MethodGet, "k", &calls, ok)
	if rec.Code != http.StatusOK || rec.Body.String() != "data" {
		t.Fatalf("GET: got %d %q, want 200 data", rec.Code, rec.Body.String())
	}
	if got := calls.Load(); got != 1 {
		t.Errorf("fetch calls = %d, want 1 (HEAD stored the body)", got)
	}
}

// Without BodyForHead a HEAD miss is relayed and not stored, as before.
func TestHeadWithoutBodyForHeadNotStored(t *testing.T) {
	c := testCache(t)
	var calls atomic.Int32
	ok := func() Result {
		rec := httptest.NewRecorder()
		rec.WriteString("data")
		return Result{Resp: rec.Result()}
	}
	serveResult(c, http.MethodHead, "k", &calls, ok)
	serveResult(c, http.MethodHead, "k", &calls, ok)
	if got := calls.Load(); got != 2 {
		t.Errorf("fetch calls = %d, want 2", got)
	}
}

// With every background slot taken, a HEAD is not offered the GET-and-store
// path: its fetch is told to send the HEAD as is.
func TestHeadStoreSlotsBounded(t *testing.T) {
	c := testCache(t)
	for i := 0; i < maxHeadStores; i++ {
		c.headSlots <- struct{}{}
	}
	storable := true
	c.ServeStoringHead(httptest.NewRecorder(), httptest.NewRequest(http.MethodHead, "/", nil), "k", 0, func(ctx context.Context) (Result, error) {
		storable = HeadStorable(ctx)
		return notFoundResult(0)(), nil
	})
	if storable {
		t.Error("HeadStorable = true with no slot free")
	}
}

// Plain Serve never detaches a HEAD or offers it the store path.
func TestServeHeadNotStorable(t *testing.T) {
	c := testCache(t)
	storable := true
	c.Serve(httptest.NewRecorder(), httptest.NewRequest(http.MethodHead, "/", nil), "k", 0, func(ctx context.Context) (Result, error) {
		storable = HeadStorable(ctx)
		return notFoundResult(0)(), nil
	})
	if storable {
		t.Error("HeadStorable = true under plain Serve")
	}
}

// A remembered miss tells whoever chains in front how long it has left.
func TestRememberedMissCarriesTTL(t *testing.T) {
	c := testCache(t)
	var calls atomic.Int32
	serveResult(c, http.MethodGet, "k", &calls, notFoundResult(90*time.Second))
	rec := serveResult(c, http.MethodGet, "k", &calls, notFoundResult(90*time.Second))
	if got := rec.Header().Get(MissTTLHeader); got != "90" {
		t.Errorf("%s = %q, want 90", MissTTLHeader, got)
	}
}