diff --git a/internal/artifact/reader_benchmark_test.go b/internal/artifact/reader_benchmark_test.go new file mode 100644 index 0000000..39fa7cd --- /dev/null +++ b/internal/artifact/reader_benchmark_test.go @@ -0,0 +1,167 @@ +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) { + body := bytes.Repeat([]byte("x"), 256<<10) + 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 }