Optimize cached artifact serving (#245)

Author: andrew

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{

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(`

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"`

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)
 	}
 }