pixa/internal/handlers/cache_hit_internal_test.go

package handlers

import (
	"bytes"
	"context"
	"database/sql"
	"database/sql/driver"
	"image"
	"image/jpeg"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"sync/atomic"
	"testing"
	"testing/fstest"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/prometheus/client_golang/prometheus"

	"sneak.berlin/go/pixa/internal/allowlist"
	"sneak.berlin/go/pixa/internal/database"
	"sneak.berlin/go/pixa/internal/encurl"
	"sneak.berlin/go/pixa/internal/imageprocessor"
	"sneak.berlin/go/pixa/internal/imageproxy"
	"sneak.berlin/go/pixa/internal/imagerequest"
	"sneak.berlin/go/pixa/internal/imgcache"
	"sneak.berlin/go/pixa/internal/seal"
	"sneak.berlin/go/pixa/internal/signature"
)

// countingHost is the allowlisted host whose photoPath the server of
// newCountingDatabaseServer serves.
const countingHost = "counting.example.com"

// countingConnector opens connections to the SQLite database file name with
// sqliteDriver, each one counting in statements the statements run on it.
type countingConnector struct {
	sqliteDriver driver.Driver
	name         string
	statements   *atomic.Int64
}

func (c countingConnector) Connect(context.Context) (driver.Conn, error) {
	conn, err := c.sqliteDriver.Open(c.name)
	if err != nil {
		return nil, err
	}

	return countingConn{Conn: conn, statements: c.statements}, nil
}

func (c countingConnector) Driver() driver.Driver {
	return c.sqliteDriver
}

// countingConn counts the statements run on its connection. It has only the
// methods of driver.Conn, so database/sql prepares every query and every
// other statement with Prepare, which counts it.
type countingConn struct {
	driver.Conn

	statements *atomic.Int64
}

func (c countingConn) Prepare(query string) (driver.Stmt, error) {
	c.statements.Add(1)

	return c.Conn.Prepare(query)
}

// newCountingDatabase returns a database with pixa's schema, in a file of
// its own and with one connection, and the count of the statements run on
// it once the schema is in place.
func newCountingDatabase(tb testing.TB) (*sql.DB, *atomic.Int64) {
	tb.Helper()

	path := filepath.Join(tb.TempDir(), "state.sqlite3")

	schemaDB, err := sql.Open("sqlite", path)
	if err != nil {
		tb.Fatalf("opening the database: %v", err)
	}

	defer func() { _ = schemaDB.Close() }()

	err = database.ApplyMigrations(tb.Context(), schemaDB, nil)
	if err != nil {
		tb.Fatalf("applying the migrations: %v", err)
	}

	statements := new(atomic.Int64)

	db := sql.OpenDB(countingConnector{
		sqliteDriver: schemaDB.Driver(),
		name:         path,
		statements:   statements,
	})
	db.SetMaxOpenConns(1)

	tb.Cleanup(func() { _ = db.Close() })

	return db, statements
}

// newCountingDatabaseServer returns a router for both image routes whose
// image cache uses newCountingDatabase, the count of the statements run on
// that database, and a URL on each route for a different image of
// photoPath on countingHost. Nothing writes the cache's access times in the
// background, so every statement counted comes from a request.
func newCountingDatabaseServer(
	tb testing.TB,
) (http.Handler, *atomic.Int64, []string) {
	tb.Helper()

	db, statements := newCountingDatabase(tb)

	var photo bytes.Buffer

	err := jpeg.Encode(&photo, image.NewRGBA(image.Rect(0, 0, 100, 100)), nil)
	if err != nil {
		tb.Fatalf("encoding the photo: %v", err)
	}

	cache, err := imgcache.NewCache(db, imgcache.Config{
		StateDir:    tb.TempDir(),
		NegativeTTL: 5 * time.Minute,
		MaxBytes:    1 << 30,
	}, prometheus.NewRegistry())
	if err != nil {
		tb.Fatalf("imgcache.NewCache() error = %v", err)
	}

	svc, err := imageproxy.New(&imageproxy.Config{
		Cache: cache,
		Fetcher: newMockFetcher(fstest.MapFS{
			countingHost + photoPath: &fstest.MapFile{Data: photo.Bytes()},
		}),
		Processor: imageprocessor.New(imageprocessor.Params{
			MaxInputBytes:           50 << 20,
			MaxConcurrentProcessing: 2,
		}),
		MaxResponseSize: 50 << 20,
		Logger:          slog.New(slog.DiscardHandler),
	}, prometheus.NewRegistry())
	if err != nil {
		tb.Fatalf("imageproxy.New() error = %v", err)
	}

	keys, err := seal.DeriveKeys(testSigningKey)
	if err != nil {
		tb.Fatalf("seal.DeriveKeys() error = %v", err)
	}

	encGen := encurl.NewGenerator(keys.EncryptedURLKey)

	token, err := encGen.Generate(&imagerequest.Request{
		Source: imagerequest.Source{Host: countingHost, Path: photoPath},
		Transform: imagerequest.Transform{
			Size:    imagerequest.Size{Width: 40, Height: 40},
			Format:  imagerequest.FormatJPEG,
			Quality: 85,
			FitMode: imagerequest.FitCover,
		},
	}, time.Time{})
	if err != nil {
		tb.Fatalf("encGen.Generate() error = %v", err)
	}

	h := &Handlers{
		log:       slog.New(slog.DiscardHandler),
		imgSvc:    svc,
		encGen:    encGen,
		signer:    signature.New(testSigningKey),
		allowlist: allowlist.New([]string{countingHost}),
	}

	r := chi.NewRouter()
	r.Get("/v1/image/*", h.HandleImage())
	r.Get("/v1/e/{token}/*", h.HandleImageEnc())

	return r, statements, []string{
		"/v1/image/" + countingHost + photoPath + "/50x50.jpeg",
		"/v1/e/" + token + "/img.jpg",
	}
}

// serveImage sends srv a GET for target and returns the response, failing
// the test unless it is a 200 with an image.
func serveImage(
	tb testing.TB, srv http.Handler, target string,
) *httptest.ResponseRecorder {
	tb.Helper()

	rec := httptest.NewRecorder()
	srv.ServeHTTP(rec,
		httptest.NewRequestWithContext(tb.Context(), http.MethodGet, target, nil))

	if rec.Code != http.StatusOK || rec.Body.Len() == 0 {
		tb.Errorf("GET %s: %d with %d bytes, want 200 with an image",
			target, rec.Code, rec.Body.Len())
	}

	return rec
}

// TestImageRoutes_CacheHitsMakeNoDatabaseCall serves an image on each image
// route once, which makes and caches it, then serves it again from the cache.
// Serving those cache hits must run no statement on the database.
func TestImageRoutes_CacheHitsMakeNoDatabaseCall(t *testing.T) {
	t.Parallel()

	srv, statements, targets := newCountingDatabaseServer(t)

	for _, target := range targets {
		before := statements.Load()

		serveImage(t, srv, target)

		afterMiss := statements.Load()
		t.Logf("GET %s, a miss: %d statements", target, afterMiss-before)

		if afterMiss == before {
			t.Fatalf("GET %s, a miss: no statement counted, want some", target)
		}

		for range 3 {
			rec := serveImage(t, srv, target)

			cacheStatus := rec.Header().Get("X-Pixa-Cache")
			if cacheStatus != string(imageproxy.CacheHit) {
				t.Errorf("GET %s: X-Pixa-Cache = %q, want %s",
					target, cacheStatus, imageproxy.CacheHit)
			}
		}

		if hitStatements := statements.Load() - afterMiss; hitStatements != 0 {
			t.Errorf("GET %s: 3 cache hits ran %d statements, want 0",
				target, hitStatements)
		}
	}
}

// BenchmarkHandleImage_CacheHits serves one cached image through the
// /v1/image/ route from b.RunParallel's goroutines, and reports the cache
// hits served per second. It fails if serving them ran any statement on the
// database. script/bench runs it.
func BenchmarkHandleImage_CacheHits(b *testing.B) {
	srv, statements, targets := newCountingDatabaseServer(b)
	target := targets[0]

	serveImage(b, srv, target)

	before := statements.Load()

	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			rec := serveImage(b, srv, target)

			cacheStatus := rec.Header().Get("X-Pixa-Cache")
			if cacheStatus != string(imageproxy.CacheHit) {
				b.Errorf("GET %s: X-Pixa-Cache = %q, want %s",
					target, cacheStatus, imageproxy.CacheHit)

				return
			}
		}
	})

	b.ReportMetric(float64(b.N)/b.Elapsed().Seconds(), "hits/s")

	if hitStatements := statements.Load() - before; hitStatements != 0 {
		b.Fatalf("%d cache hits ran %d statements, want 0", b.N, hitStatements)
	}
}