From aa205221ea8127fa3e0acf977810c33469b89b61 Mon Sep 17 00:00:00 2001 From: Andrew Nesbitt Date: Wed, 12 Aug 2026 19:47:38 +0100 Subject: [PATCH] Optimize cached artifact serving --- internal/database/database_test.go | 80 +++++++ internal/database/queries.go | 22 ++ internal/database/types.go | 10 + internal/handler/handler.go | 65 +++--- internal/handler/handler_bench_test.go | 300 +++++++++++++++++++++++++ internal/handler/handler_test.go | 4 +- 6 files changed, 444 insertions(+), 37 deletions(-) create mode 100644 internal/handler/handler_bench_test.go diff --git a/internal/database/database_test.go b/internal/database/database_test.go index 6fca4ea..3e92b91 100644 --- a/internal/database/database_test.go +++ b/internal/database/database_test.go @@ -239,6 +239,86 @@ func TestArtifactCRUD(t *testing.T) { }) } +func TestGetCachedArtifact(t *testing.T) { + runWithBothDatabases(t, func(t *testing.T, db *DB) { + const ( + packagePURL = "pkg:npm/lodash" + versionPURL = "pkg:npm/lodash@4.17.21" + filename = "lodash-4.17.21.tgz" + ) + seedCachedArtifactTestData(t, db, packagePURL, versionPURL, filename) + + cached, err := db.GetCachedArtifact(packagePURL, versionPURL, filename) + if err != nil { + t.Fatalf("GetCachedArtifact before cache failed: %v", err) + } + if cached != nil { + t.Fatalf("expected no cached artifact, got %+v", cached) + } + + if err := db.MarkArtifactCached(versionPURL, filename, "/cache/npm/"+filename, + "sha256-abc", 12345, "application/gzip"); err != nil { + t.Fatalf("MarkArtifactCached failed: %v", err) + } + + cached, err = db.GetCachedArtifact(packagePURL, versionPURL, filename) + if err != nil { + t.Fatalf("GetCachedArtifact failed: %v", err) + } + if cached == nil { + t.Fatal("expected cached artifact, got nil") + } + if cached.Ecosystem != "npm" { + t.Errorf("expected npm ecosystem, got %q", cached.Ecosystem) + } + if cached.StoragePath != "/cache/npm/"+filename { + t.Errorf("expected cached storage path, got %q", cached.StoragePath) + } + if cached.ContentHash.String != "sha256-abc" { + t.Errorf("expected cached content hash, got %q", cached.ContentHash.String) + } + if cached.Size.Int64 != 12345 { + t.Errorf("expected cached size 12345, got %d", cached.Size.Int64) + } + if cached.ContentType.String != "application/gzip" { + t.Errorf("expected cached content type, got %q", cached.ContentType.String) + } + if cached.Integrity.String != "sha512-abc123" { + t.Errorf("expected cached integrity, got %q", cached.Integrity.String) + } + + cached, err = db.GetCachedArtifact("pkg:npm/other", versionPURL, filename) + if err != nil { + t.Fatalf("GetCachedArtifact with wrong package failed: %v", err) + } + if cached != nil { + t.Fatalf("expected package mismatch to miss cache, got %+v", cached) + } + }) +} + +func seedCachedArtifactTestData(t *testing.T, db *DB, packagePURL, versionPURL, filename string) { + t.Helper() + + if err := db.UpsertPackage(&Package{PURL: packagePURL, Ecosystem: "npm", Name: "lodash"}); err != nil { + t.Fatalf("UpsertPackage failed: %v", err) + } + if err := db.UpsertVersion(&Version{ + PURL: versionPURL, + PackagePURL: packagePURL, + Integrity: sql.NullString{String: "sha512-abc123", Valid: true}, + }); err != nil { + t.Fatalf("UpsertVersion failed: %v", err) + } + if err := db.UpsertArtifact(&Artifact{ + VersionPURL: versionPURL, + Filename: filename, + UpstreamURL: "https://registry.npmjs.org/lodash/-/" + filename, + }); err != nil { + t.Fatalf("UpsertArtifact failed: %v", err) + } +} + func TestCacheManagement(t *testing.T) { runWithBothDatabases(t, func(t *testing.T, db *DB) { pkg := &Package{ diff --git a/internal/database/queries.go b/internal/database/queries.go index f8b76fe..b01abe4 100644 --- a/internal/database/queries.go +++ b/internal/database/queries.go @@ -191,6 +191,28 @@ func (db *DB) GetArtifact(versionPURL, filename string) (*Artifact, error) { return &a, nil } +// GetCachedArtifact returns the fields needed to serve a cached artifact. +func (db *DB) GetCachedArtifact(packagePURL, versionPURL, filename string) (*CachedArtifact, error) { + var artifact CachedArtifact + query := db.Rebind(` + SELECT packages.ecosystem, artifacts.storage_path, artifacts.content_hash, artifacts.size, + artifacts.content_type, versions.integrity + FROM artifacts + JOIN versions ON versions.purl = artifacts.version_purl + JOIN packages ON packages.purl = versions.package_purl + WHERE packages.purl = ? AND artifacts.version_purl = ? AND artifacts.filename = ? + AND artifacts.storage_path IS NOT NULL AND artifacts.fetched_at IS NOT NULL + `) + err := db.Get(&artifact, query, packagePURL, versionPURL, filename) + if err == sql.ErrNoRows { + return nil, nil + } + if err != nil { + return nil, err + } + return &artifact, nil +} + func (db *DB) GetArtifactByPath(storagePath string) (*Artifact, error) { var a Artifact query := db.Rebind(` diff --git a/internal/database/types.go b/internal/database/types.go index 7128f12..3826c7a 100644 --- a/internal/database/types.go +++ b/internal/database/types.go @@ -76,6 +76,16 @@ func (a *Artifact) IsCached() bool { return a.StoragePath.Valid && a.FetchedAt.Valid } +// CachedArtifact contains the fields needed to serve a cached artifact. +type CachedArtifact struct { + Ecosystem string `db:"ecosystem"` + StoragePath string `db:"storage_path"` + ContentHash sql.NullString `db:"content_hash"` + Size sql.NullInt64 `db:"size"` + ContentType sql.NullString `db:"content_type"` + Integrity sql.NullString `db:"integrity"` +} + // MetadataCacheEntry represents a cached metadata blob for offline serving. type MetadataCacheEntry struct { ID int64 `db:"id" json:"id"` diff --git a/internal/handler/handler.go b/internal/handler/handler.go index 72b9c28..6a7aab4 100644 --- a/internal/handler/handler.go +++ b/internal/handler/handler.go @@ -13,6 +13,7 @@ import ( "net/url" "strconv" "strings" + "sync" "time" "github.com/git-pkgs/cooldown" @@ -48,6 +49,15 @@ func hasDotDotSegment(path string) bool { const defaultHTTPTimeout = 30 * time.Second +const artifactCopyBufferSize = 32 << 10 + +var artifactCopyBufferPool = sync.Pool{ //nolint:gochecknoglobals // shared across artifact responses + New: func() any { + buffer := make([]byte, artifactCopyBufferSize) + return &buffer + }, +} + // canonicalPackagePURL returns a versionless PURL in canonical form so cooldown // lookups match keys produced by config.CooldownConfig.NormalizedPackages. func canonicalPackagePURL(ecosystem, name string) string { @@ -157,27 +167,11 @@ func (p *Proxy) GetCachedArtifact(ctx context.Context, ecosystem, name, version, // checkCache looks up an artifact in the cache. Returns nil if not cached. func (p *Proxy) checkCache(ctx context.Context, pkgPURL, versionPURL, filename string) (*CacheResult, error) { - pkg, err := p.DB.GetPackageByPURL(pkgPURL) - if err != nil { - return nil, fmt.Errorf("checking package cache: %w", err) - } - if pkg == nil { - return nil, nil - } - - ver, err := p.DB.GetVersionByPURL(versionPURL) - if err != nil { - return nil, fmt.Errorf("checking version cache: %w", err) - } - if ver == nil { - return nil, nil - } - - artifact, err := p.DB.GetArtifact(versionPURL, filename) + artifact, err := p.DB.GetCachedArtifact(pkgPURL, versionPURL, filename) if err != nil { return nil, fmt.Errorf("checking artifact cache: %w", err) } - if artifact == nil || !artifact.IsCached() { + if artifact == nil { return nil, nil } @@ -189,39 +183,39 @@ func (p *Proxy) checkCache(ctx context.Context, pkgPURL, versionPURL, filename s } if p.DirectServe { - signed, err := p.Storage.SignedURL(ctx, artifact.StoragePath.String, p.DirectServeTTL) + signed, err := p.Storage.SignedURL(ctx, artifact.StoragePath, p.DirectServeTTL) if err == nil { result.RedirectURL = rewriteSignedURLHost(signed, p.DirectServeBaseURL) - p.recordCacheHit(pkgPURL, versionPURL, filename) + p.recordCacheHit(artifact.Ecosystem, versionPURL, filename) return result, nil } if !errors.Is(err, storage.ErrSignedURLUnsupported) { p.Logger.Warn("failed to sign storage URL, falling back to streaming", - "path", artifact.StoragePath.String, "error", err) + "path", artifact.StoragePath, "error", err) } } start := time.Now() - reader, err := p.Storage.Open(ctx, artifact.StoragePath.String) + reader, err := p.Storage.Open(ctx, artifact.StoragePath) metrics.RecordStorageOperation("read", time.Since(start)) if err != nil { metrics.RecordStorageError("read") p.Logger.Warn("cached artifact missing from storage, will refetch", - "path", artifact.StoragePath.String, "error", err) + "path", artifact.StoragePath, "error", err) return nil, nil } - result.Reader = newVerifyingReader(reader, artifact.ContentHash.String, ver.Integrity.String, + result.Reader = newVerifyingReader(reader, artifact.ContentHash.String, artifact.Integrity.String, func(reason string) { p.Logger.Error("cached artifact failed integrity check", "purl", versionPURL, "filename", filename, - "path", artifact.StoragePath.String, "reason", reason) - metrics.RecordIntegrityFailure(pkg.Ecosystem) + "path", artifact.StoragePath, "reason", reason) + metrics.RecordIntegrityFailure(artifact.Ecosystem) if err := p.DB.ClearArtifactCache(versionPURL, filename); err != nil { p.Logger.Warn("failed to clear corrupt artifact from cache", "error", err) } }) - p.recordCacheHit(pkgPURL, versionPURL, filename) + p.recordCacheHit(artifact.Ecosystem, versionPURL, filename) return result, nil } @@ -245,11 +239,9 @@ func rewriteSignedURLHost(signed, baseURL string) string { return s.String() } -func (p *Proxy) recordCacheHit(pkgPURL, versionPURL, filename string) { +func (p *Proxy) recordCacheHit(ecosystem, versionPURL, filename string) { _ = p.DB.RecordArtifactHit(versionPURL, filename) - if parsed, err := purl.Parse(pkgPURL); err == nil { - metrics.RecordCacheHit(purl.PURLTypeToEcosystem(parsed.Type)) - } + metrics.RecordCacheHit(purl.NormalizeEcosystem(ecosystem)) } func (p *Proxy) fetchAndCache(ctx context.Context, ecosystem, name, version, filename, pkgPURL, versionPURL string) (*CacheResult, error) { @@ -376,7 +368,7 @@ func ServeArtifact(w http.ResponseWriter, result *CacheResult) { func serveArtifact(w http.ResponseWriter, method string, result *CacheResult) { if result.RedirectURL != "" { if result.Hash != "" { - w.Header().Set("ETag", fmt.Sprintf(`"%s"`, result.Hash)) + w.Header().Set("ETag", `"`+result.Hash+`"`) } w.Header().Set("Location", result.RedirectURL) w.WriteHeader(http.StatusFound) @@ -391,15 +383,18 @@ func serveArtifact(w http.ResponseWriter, method string, result *CacheResult) { w.Header().Set("Content-Type", result.ContentType) } if result.Size > 0 || (method == http.MethodHead && result.Size == 0) { - w.Header().Set("Content-Length", fmt.Sprintf("%d", result.Size)) + w.Header().Set("Content-Length", strconv.FormatInt(result.Size, 10)) } if result.Hash != "" { - w.Header().Set("ETag", fmt.Sprintf(`"%s"`, result.Hash)) + w.Header().Set("ETag", `"`+result.Hash+`"`) } w.WriteHeader(http.StatusOK) if method != http.MethodHead && result.Reader != nil { - _, _ = io.Copy(w, result.Reader) + buffer := artifactCopyBufferPool.Get().(*[]byte) + defer artifactCopyBufferPool.Put(buffer) + // Hide optional ReaderFrom methods so io.CopyBuffer uses the pooled buffer. + _, _ = io.CopyBuffer(struct{ io.Writer }{w}, result.Reader, *buffer) } } diff --git a/internal/handler/handler_bench_test.go b/internal/handler/handler_bench_test.go new file mode 100644 index 0000000..cdbb524 --- /dev/null +++ b/internal/handler/handler_bench_test.go @@ -0,0 +1,300 @@ +package handler + +import ( + "bytes" + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/git-pkgs/proxy/internal/database" + "github.com/git-pkgs/proxy/internal/storage" + "github.com/git-pkgs/purl" + "github.com/git-pkgs/registries/fetch" +) + +const benchmarkArtifactSize = 64 << 10 + +const benchmarkMetadataSize = 1 << 20 + +type benchmarkResponseWriter struct { + header http.Header +} + +func (w *benchmarkResponseWriter) Header() http.Header { + return w.header +} + +func (w *benchmarkResponseWriter) Write(p []byte) (int, error) { + return len(p), nil +} + +func (w *benchmarkResponseWriter) WriteHeader(_ int) {} + +func benchmarkCachedProxy(b *testing.B) (*Proxy, *mockStorage) { + b.Helper() + + proxy, db, store, _ := setupTestProxy(b) + content := strings.Repeat("x", benchmarkArtifactSize) + seedPackage(b, db, store, "npm", "lodash", "4.17.21", "lodash-4.17.21.tgz", content) + + artifact, err := db.GetArtifact("pkg:npm/lodash@4.17.21", "lodash-4.17.21.tgz") + if err != nil { + b.Fatalf("get seeded artifact: %v", err) + } + sum := sha256.Sum256([]byte(content)) + artifact.ContentHash.String = hex.EncodeToString(sum[:]) + if err := db.UpsertArtifact(artifact); err != nil { + b.Fatalf("update seeded artifact hash: %v", err) + } + + return proxy, store +} + +func BenchmarkArtifactCacheHit(b *testing.B) { + ctx := context.Background() + + b.Run("stream-64KiB", func(b *testing.B) { + proxy, _ := benchmarkCachedProxy(b) + w := &benchmarkResponseWriter{header: make(http.Header)} + b.SetBytes(benchmarkArtifactSize) + b.ReportAllocs() + b.ResetTimer() + + for b.Loop() { + result, err := proxy.GetOrFetchArtifact(ctx, "npm", "lodash", "4.17.21", "lodash-4.17.21.tgz") + if err != nil { + b.Fatal(err) + } + ServeArtifact(w, result) + } + }) + + b.Run("direct-serve", func(b *testing.B) { + proxy, store := benchmarkCachedProxy(b) + proxy.DirectServe = true + store.signedURL = "https://storage.example/npm/lodash-4.17.21.tgz?signature=abc" + w := &benchmarkResponseWriter{header: make(http.Header)} + b.ReportAllocs() + b.ResetTimer() + + for b.Loop() { + result, err := proxy.GetOrFetchArtifact(ctx, "npm", "lodash", "4.17.21", "lodash-4.17.21.tgz") + if err != nil { + b.Fatal(err) + } + ServeArtifact(w, result) + } + }) +} + +func BenchmarkArtifactCacheHitParallel(b *testing.B) { + proxy, _ := benchmarkCachedProxy(b) + ctx := context.Background() + b.SetBytes(benchmarkArtifactSize) + b.ReportAllocs() + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + w := &benchmarkResponseWriter{header: make(http.Header)} + for pb.Next() { + result, err := proxy.GetOrFetchArtifact(ctx, "npm", "lodash", "4.17.21", "lodash-4.17.21.tgz") + if err != nil { + b.Error(err) + return + } + ServeArtifact(w, result) + } + }) +} + +func BenchmarkReadMetadata(b *testing.B) { + payload := bytes.Repeat([]byte("x"), benchmarkMetadataSize) + proxy := &Proxy{MetadataMaxSize: benchmarkMetadataSize} + b.SetBytes(benchmarkMetadataSize) + b.ReportAllocs() + + var data []byte + for b.Loop() { + var err error + data, err = proxy.ReadMetadata(bytes.NewReader(payload)) + if err != nil { + b.Fatal(err) + } + } + if len(data) != len(payload) { + b.Fatalf("metadata size = %d, want %d", len(data), len(payload)) + } +} + +func BenchmarkArtifactPURLConstruction(b *testing.B) { + for _, tc := range []struct { + name string + ecosystem string + packageID string + }{ + {"npm", "npm", "lodash"}, + {"scoped-npm", "npm", "@scope/package"}, + {"go", "golang", "github.com/git-pkgs/proxy"}, + } { + b.Run(tc.name, func(b *testing.B) { + b.ReportAllocs() + var packagePURL, versionPURL string + for b.Loop() { + packagePURL = purl.MakePURLString(tc.ecosystem, tc.packageID, "") + versionPURL = purl.MakePURLString(tc.ecosystem, tc.packageID, "1.2.3") + } + if packagePURL == "" || versionPURL == "" { + b.Fatal("empty PURL") + } + }) + } +} + +type benchmarkNPMServer struct { + client *http.Client + requestURL string + db *database.DB + versionPURL string + filename string +} + +func newBenchmarkNPMServer(b *testing.B) *benchmarkNPMServer { + b.Helper() + + ctx := context.Background() + dir := b.TempDir() + db, err := database.Create(filepath.Join(dir, "benchmark.db")) + if err != nil { + b.Fatalf("create database: %v", err) + } + b.Cleanup(func() { _ = db.Close() }) + + store, err := storage.OpenBucket(ctx, "file://"+filepath.Join(dir, "cache")) + if err != nil { + b.Fatalf("open storage: %v", err) + } + b.Cleanup(func() { _ = store.Close() }) + + content := bytes.Repeat([]byte("x"), benchmarkArtifactSize) + storagePath := storage.ArtifactPath("npm", "", "lodash", "4.17.21", "lodash-4.17.21.tgz") + size, hash, err := store.Store(ctx, storagePath, bytes.NewReader(content)) + if err != nil { + b.Fatalf("store artifact: %v", err) + } + + pkg := &database.Package{PURL: "pkg:npm/lodash", Ecosystem: "npm", Name: "lodash"} + if err := db.UpsertPackage(pkg); err != nil { + b.Fatalf("seed package: %v", err) + } + version := &database.Version{PURL: "pkg:npm/lodash@4.17.21", PackagePURL: pkg.PURL} + if err := db.UpsertVersion(version); err != nil { + b.Fatalf("seed version: %v", err) + } + artifact := &database.Artifact{ + VersionPURL: version.PURL, + Filename: "lodash-4.17.21.tgz", + UpstreamURL: "https://registry.npmjs.org/lodash/-/lodash-4.17.21.tgz", + StoragePath: sql.NullString{String: storagePath, Valid: true}, + ContentHash: sql.NullString{String: hash, Valid: true}, + Size: sql.NullInt64{Int64: size, Valid: true}, + ContentType: sql.NullString{String: "application/gzip", Valid: true}, + FetchedAt: sql.NullTime{Time: time.Now(), Valid: true}, + } + if err := db.UpsertArtifact(artifact); err != nil { + b.Fatalf("seed artifact: %v", err) + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + proxy := NewProxy(db, store, &mockFetcher{}, fetch.NewResolver(), logger) + handler := NewNPMHandler(proxy, "http://proxy.example", "https://registry.npmjs.org") + server := httptest.NewServer(handler.Routes()) + b.Cleanup(server.Close) + client := server.Client() + return &benchmarkNPMServer{ + client: client, + requestURL: server.URL + "/lodash/-/lodash-4.17.21.tgz", + db: db, + versionPURL: version.PURL, + filename: artifact.Filename, + } +} + +func (s *benchmarkNPMServer) request() error { + resp, err := s.client.Get(s.requestURL) + if err != nil { + return fmt.Errorf("GET cached artifact: %w", err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("GET cached artifact status = %d, want %d", resp.StatusCode, http.StatusOK) + } + n, err := io.Copy(io.Discard, resp.Body) + if err != nil { + return fmt.Errorf("read cached artifact: %w", err) + } + if n != benchmarkArtifactSize { + return fmt.Errorf("cached artifact size = %d, want %d", n, benchmarkArtifactSize) + } + return nil +} + +func (s *benchmarkNPMServer) hitCount(b *testing.B) int64 { + b.Helper() + artifact, err := s.db.GetArtifact(s.versionPURL, s.filename) + if err != nil { + b.Fatalf("get artifact hit count: %v", err) + } + return artifact.HitCount +} + +func benchmarkNPMArtifactCacheHitHTTP(b *testing.B, parallel bool) { + server := newBenchmarkNPMServer(b) + if err := server.request(); err != nil { + b.Fatal(err) + } + startHits := server.hitCount(b) + + b.SetBytes(benchmarkArtifactSize) + b.ReportAllocs() + b.ResetTimer() + if parallel { + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + if err := server.request(); err != nil { + b.Error(err) + return + } + } + }) + } else { + for b.Loop() { + if err := server.request(); err != nil { + b.Fatal(err) + } + } + } + b.StopTimer() + + if hitCount := server.hitCount(b) - startHits; hitCount != int64(b.N) { + b.Fatalf("new artifact hits = %d, want %d", hitCount, b.N) + } + b.ReportMetric(float64(b.N)/b.Elapsed().Seconds(), "requests/s") +} + +func BenchmarkNPMArtifactCacheHitHTTP(b *testing.B) { + benchmarkNPMArtifactCacheHitHTTP(b, false) +} + +func BenchmarkNPMArtifactCacheHitHTTPParallel(b *testing.B) { + benchmarkNPMArtifactCacheHitHTTP(b, true) +} diff --git a/internal/handler/handler_test.go b/internal/handler/handler_test.go index a3bd9b3..d52e7b6 100644 --- a/internal/handler/handler_test.go +++ b/internal/handler/handler_test.go @@ -128,7 +128,7 @@ func (f *mockFetcher) Head(_ context.Context, _ string) (int64, string, error) { } // setupTestProxy creates a Proxy with a real DB (SQLite in temp dir) and mock storage/fetcher. -func setupTestProxy(t *testing.T) (*Proxy, *database.DB, *mockStorage, *mockFetcher) { +func setupTestProxy(t testing.TB) (*Proxy, *database.DB, *mockStorage, *mockFetcher) { t.Helper() dir := t.TempDir() @@ -148,7 +148,7 @@ func setupTestProxy(t *testing.T) (*Proxy, *database.DB, *mockStorage, *mockFetc } // seedPackage creates a package, version, and cached artifact in the test DB and storage. -func seedPackage(t *testing.T, db *database.DB, store *mockStorage, ecosystem, name, version, filename, content string) { +func seedPackage(t testing.TB, db *database.DB, store *mockStorage, ecosystem, name, version, filename, content string) { t.Helper() pkg := &database.Package{