pkg-proxy/internal/database/hits_test.go

package database

import (
	"bytes"
	"log/slog"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

var discardLogger = slog.New(slog.DiscardHandler)

func seedHitTestArtifact(t *testing.T, db *DB) (string, string) {
	t.Helper()
	versionPURL, filename := "pkg:npm/lodash@4.17.21", "lodash-4.17.21.tgz"
	if err := db.UpsertPackage(&Package{PURL: "pkg:npm/lodash", Ecosystem: "npm", Name: "lodash"}); err != nil {
		t.Fatalf("UpsertPackage failed: %v", err)
	}
	if err := db.UpsertVersion(&Version{PURL: versionPURL, PackagePURL: "pkg:npm/lodash"}); 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)
	}
	return versionPURL, filename
}

func hitCount(t *testing.T, db *DB, versionPURL, filename string) int64 {
	t.Helper()
	a, err := db.GetArtifact(versionPURL, filename)
	if err != nil || a == nil {
		t.Fatalf("GetArtifact failed: %v", err)
	}
	return a.HitCount
}

func TestBatchHitsWritesOnFlush(t *testing.T) {
	runWithBothDatabases(t, func(t *testing.T, db *DB) {
		versionPURL, filename := seedHitTestArtifact(t, db)
		db.BatchHits(time.Hour, discardLogger)

		for range 3 {
			if err := db.RecordArtifactHit(versionPURL, filename); err != nil {
				t.Fatalf("RecordArtifactHit failed: %v", err)
			}
		}
		if got := hitCount(t, db, versionPURL, filename); got != 0 {
			t.Fatalf("hit count before flush = %d, want 0", got)
		}

		if err := db.flushHits(); err != nil {
			t.Fatalf("flushHits failed: %v", err)
		}
		a, err := db.GetArtifact(versionPURL, filename)
		if err != nil {
			t.Fatalf("GetArtifact failed: %v", err)
		}
		if a.HitCount != 3 {
			t.Errorf("hit count after flush = %d, want 3", a.HitCount)
		}
		if !a.LastAccessedAt.Valid {
			t.Error("expected last_accessed_at to be set")
		}
	})
}

func TestBatchHitsWritesOnClose(t *testing.T) {
	path := filepath.Join(t.TempDir(), "test.db")
	db, err := Create(path)
	if err != nil {
		t.Fatalf("Create failed: %v", err)
	}
	versionPURL, filename := seedHitTestArtifact(t, db)
	db.BatchHits(time.Hour, discardLogger)

	for range 2 {
		_ = db.RecordArtifactHit(versionPURL, filename)
	}
	if err := db.Close(); err != nil {
		t.Fatalf("Close failed: %v", err)
	}

	db, err = Open(path)
	if err != nil {
		t.Fatalf("Open failed: %v", err)
	}
	defer func() { _ = db.Close() }()
	if got := hitCount(t, db, versionPURL, filename); got != 2 {
		t.Errorf("hit count after Close = %d, want 2", got)
	}
}

// TestBatchHitsKeepNewerAccessTime flushes a batch older than a hit already
// written, as when proxies share a database. The count still adds up, but
// neither timestamp may move backwards.
func TestBatchHitsKeepNewerAccessTime(t *testing.T) {
	runWithBothDatabases(t, func(t *testing.T, db *DB) {
		versionPURL, filename := seedHitTestArtifact(t, db)
		k := hitKey{versionPURL, filename}
		now := time.Now()
		read := func() *Artifact {
			t.Helper()
			a, err := db.GetArtifact(versionPURL, filename)
			if err != nil || a == nil {
				t.Fatalf("GetArtifact failed: %v", err)
			}
			return a
		}
		write := func(last time.Time) {
			t.Helper()
			if err := db.writeHits(map[hitKey]hitEntry{k: {count: 1, last: last}}); err != nil {
				t.Fatalf("writeHits failed: %v", err)
			}
		}

		write(now)
		newer := read()
		if !newer.LastAccessedAt.Valid {
			t.Fatal("last_accessed_at not set from NULL")
		}

		write(now.Add(-time.Minute))
		got := read()
		if got.HitCount != 2 {
			t.Errorf("hit count = %d, want 2", got.HitCount)
		}
		if !got.LastAccessedAt.Time.Equal(newer.LastAccessedAt.Time) {
			t.Errorf("last_accessed_at moved from %v to %v", newer.LastAccessedAt.Time, got.LastAccessedAt.Time)
		}
		if !got.UpdatedAt.Equal(newer.UpdatedAt) {
			t.Errorf("updated_at moved from %v to %v", newer.UpdatedAt, got.UpdatedAt)
		}

		write(now.Add(time.Minute))
		got = read()
		if !got.LastAccessedAt.Time.After(newer.LastAccessedAt.Time) || !got.UpdatedAt.After(newer.UpdatedAt) {
			t.Error("a newer batch did not advance the timestamps")
		}
	})
}

func TestBatchHitsZeroIntervalWritesImmediately(t *testing.T) {
	runWithBothDatabases(t, func(t *testing.T, db *DB) {
		versionPURL, filename := seedHitTestArtifact(t, db)
		db.BatchHits(0, discardLogger)

		_ = db.RecordArtifactHit(versionPURL, filename)
		if got := hitCount(t, db, versionPURL, filename); got != 1 {
			t.Errorf("hit count = %d, want 1", got)
		}
	})
}

func TestBatchHitsKeepsHitsWhenWriteFails(t *testing.T) {
	db := createTestDB(t)
	versionPURL, filename := seedHitTestArtifact(t, db)
	db.BatchHits(time.Hour, discardLogger)

	for range 2 {
		_ = db.RecordArtifactHit(versionPURL, filename)
	}
	_ = db.DB.Close()
	if err := db.flushHits(); err == nil {
		t.Fatal("expected flushHits to fail on a closed database")
	}
	if got := db.hits.take()[hitKey{versionPURL, filename}].count; got != 2 {
		t.Errorf("pending hits after failed flush = %d, want 2", got)
	}
	_ = db.Close()
}

// TestBatchHitsCountsEveryHitWhileFlushing records hits from several goroutines
// while another flushes, so no hit may be lost between taking the pending
// hits and writing them.
func TestBatchHitsCountsEveryHitWhileFlushing(t *testing.T) {
	runWithBothDatabases(t, func(t *testing.T, db *DB) {
		versionPURL, filename := seedHitTestArtifact(t, db)
		db.BatchHits(time.Hour, discardLogger)

		const writers, hitsEach = 8, 100
		stop := make(chan struct{})
		flushed := make(chan struct{})
		go func() {
			defer close(flushed)
			for {
				select {
				case <-stop:
					return
				default:
					if err := db.flushHits(); err != nil {
						t.Errorf("flushHits failed: %v", err)
						return
					}
				}
			}
		}()
		var wg sync.WaitGroup
		for range writers {
			wg.Go(func() {
				for range hitsEach {
					if err := db.RecordArtifactHit(versionPURL, filename); err != nil {
						t.Errorf("RecordArtifactHit failed: %v", err)
					}
				}
			})
		}
		wg.Wait()
		close(stop)
		<-flushed

		if err := db.flushHits(); err != nil {
			t.Fatalf("final flushHits failed: %v", err)
		}
		if got := hitCount(t, db, versionPURL, filename); got != writers*hitsEach {
			t.Errorf("hit count = %d, want %d", got, writers*hitsEach)
		}
	})
}

func TestBatchHitsLogsHitsDroppedOnClose(t *testing.T) {
	db := createTestDB(t)
	versionPURL, filename := seedHitTestArtifact(t, db)
	var logs bytes.Buffer
	db.BatchHits(time.Hour, slog.New(slog.NewTextHandler(&logs, nil)))

	_ = db.RecordArtifactHit(versionPURL, filename)
	_ = db.DB.Close()
	_ = db.Close()

	if !strings.Contains(logs.String(), "dropping them") {
		t.Errorf("no error logged for hits dropped on close, logs: %q", logs.String())
	}
}

func TestCloseTwiceWithBatchHits(t *testing.T) {
	db := createTestDB(t)
	db.BatchHits(time.Hour, discardLogger)
	_ = db.Close()
	_ = db.Close()
}