promptkit/internal/artifact/reader_benchmark_test.go

package artifact

import (
	"bytes"
	"context"
	"crypto/sha256"
	"fmt"
	"mime"
	"os"
	"path/filepath"
	"runtime"
	"testing"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
)

var artifactFileReadBenchmarkSink *domain.Artifact

func BenchmarkArtifactFileRead(b *testing.B) {
	for _, fixture := range makeArtifactReadBenchmarkFiles(b) {
		b.Run(fixture.name, func(b *testing.B) {
			reader := &fileReader{open: openArtifactFile}
			ref := domain.ArtifactRef{Type: domain.ArtifactRefFile, URI: fixture.path}
			checkArtifactFileRead(b, reader, ref, fixture.body)

			b.ReportAllocs()
			b.SetBytes(int64(len(fixture.body)))
			b.ResetTimer()
			for range b.N {
				artifact, err := reader.Read(context.Background(), ref)
				if err != nil {
					b.Fatalf("read artifact: %v", err)
				}
				artifactFileReadBenchmarkSink = artifact
			}
			b.StopTimer()
			runtime.KeepAlive(artifactFileReadBenchmarkSink)
		})
	}
}

// BenchmarkRetainedArtifactFileRead measures live artifact bodies when an
// opened file reports a size larger than the content that reaches EOF.
func BenchmarkRetainedArtifactFileRead(b *testing.B) {
	for _, fixture := range []struct {
		name string
		body []byte
	}{
		{name: "small", body: bytes.Repeat([]byte("x"), 256<<10)},
		{name: "empty", body: []byte{}},
	} {
		b.Run(fixture.name, func(b *testing.B) {
			benchmarkRetainedArtifactFileRead(b, fixture.body)
		})
	}
}

func benchmarkRetainedArtifactFileRead(b *testing.B, body []byte) {
	path := filepath.Join(b.TempDir(), "overreported.txt")
	if err := os.WriteFile(path, body, 0o600); err != nil {
		b.Fatal(err)
	}
	mime.TypeByExtension(filepath.Ext(path))

	reader := &fileReader{open: func(path string) (artifactFile, error) {
		file, err := os.Open(path)
		if err != nil {
			return nil, err
		}
		return &sizeReportingArtifactFile{artifactFile: file, reportedSize: 16 << 20}, nil
	}}
	ref := domain.ArtifactRef{Type: domain.ArtifactRefFile, URI: path}
	checkArtifactFileRead(b, reader, ref, body)

	const retainedBatchSize = 4
	retained := make([]*domain.Artifact, retainedBatchSize)
	artifactFileReadBenchmarkSink = nil
	runtime.GC()
	var setupMemory runtime.MemStats
	runtime.ReadMemStats(&setupMemory)
	runtime.KeepAlive(reader)
	runtime.KeepAlive(ref)

	b.ReportAllocs()
	b.SetBytes(int64(len(body)))
	b.ResetTimer()
	for iteration := 0; iteration < b.N; iteration++ {
		artifact, err := reader.Read(context.Background(), ref)
		if err != nil {
			b.Fatalf("read retained artifact: %v", err)
		}
		retained[iteration%len(retained)] = artifact
	}
	b.StopTimer()
	for i, artifact := range retained {
		if artifact != nil {
			continue
		}
		artifact, err := reader.Read(context.Background(), ref)
		if err != nil {
			b.Fatalf("fill retained artifact batch: %v", err)
		}
		retained[i] = artifact
	}
	runtime.GC()
	var retainedMemory runtime.MemStats
	runtime.ReadMemStats(&retainedMemory)
	b.ReportMetric(float64(int64(retainedMemory.HeapAlloc)-int64(setupMemory.HeapAlloc)), "retained-B")
	runtime.KeepAlive(retained)
	runtime.KeepAlive(reader)
	runtime.KeepAlive(ref)
}

type artifactReadBenchmarkFile struct {
	name string
	path string
	body []byte
}

func makeArtifactReadBenchmarkFiles(b *testing.B) []artifactReadBenchmarkFile {
	b.Helper()
	fixtures := []struct {
		name string
		size int
	}{
		{name: "small_4KiB", size: 4 << 10},
		{name: "1MiB", size: 1 << 20},
		{name: "large_4MiB", size: 4 << 20},
	}
	results := make([]artifactReadBenchmarkFile, 0, len(fixtures))
	for _, fixture := range fixtures {
		body := bytes.Repeat([]byte("x"), fixture.size)
		path := filepath.Join(b.TempDir(), fixture.name+".txt")
		if err := os.WriteFile(path, body, 0o600); err != nil {
			b.Fatal(err)
		}
		mime.TypeByExtension(filepath.Ext(path))
		results = append(results, artifactReadBenchmarkFile{
			name: fixture.name,
			path: path,
			body: body,
		})
	}
	return results
}

func checkArtifactFileRead(b *testing.B, reader Reader, ref domain.ArtifactRef, want []byte) {
	b.Helper()
	artifact, err := reader.Read(context.Background(), ref)
	if err != nil {
		b.Fatalf("check artifact read: %v", err)
	}
	if !bytes.Equal(artifact.Body, want) || artifact.Size != int64(len(want)) {
		b.Fatalf("artifact size=%d body length=%d, want %d", artifact.Size, len(artifact.Body), len(want))
	}
	wantHash := fmt.Sprintf("%x", sha256.Sum256(want))
	if artifact.Hash != wantHash {
		b.Fatalf("artifact hash=%q, want %q", artifact.Hash, wantHash)
	}
}

type sizeReportingArtifactFile struct {
	artifactFile
	reportedSize int64
}

func (file *sizeReportingArtifactFile) Stat() (os.FileInfo, error) {
	info, err := file.artifactFile.Stat()
	if err != nil {
		return nil, err
	}
	return reportedArtifactFileInfo{FileInfo: info, size: file.reportedSize}, nil
}

type reportedArtifactFileInfo struct {
	os.FileInfo
	size int64
}

func (info reportedArtifactFileInfo) Size() int64 { return info.size }