diff --git a/embedded/putstream_regularfile_benchmark_test.go b/embedded/putstream_regularfile_benchmark_test.go new file mode 100644 index 00000000..7a10f1e9 --- /dev/null +++ b/embedded/putstream_regularfile_benchmark_test.go @@ -0,0 +1,120 @@ +// Copyright 2026 Tigris Data, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package embedded + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "testing" + "time" + + "github.com/rs/zerolog" + stor "github.com/tigrisdata/ocache/storage" +) + +const ( + // Keep values above the default 64 MiB compaction threshold. + embeddedPutStreamRegularFileBenchmarkSize = 256 << 20 + // Match the replay's fresh-key batch while bounding queued cleanup. + embeddedPutStreamRegularFileBenchmarkBatchSize = 8 +) + +// BenchmarkEmbeddedClientPutStreamRegularFile measures owner-local PutStream +// calls for raw-file inputs. Each bounded batch uses fresh keys; delete and +// client teardown happen outside the timer before the next batch starts. +func BenchmarkEmbeddedClientPutStreamRegularFile(b *testing.B) { + previousLogLevel := zerolog.GlobalLevel() + zerolog.SetGlobalLevel(zerolog.Disabled) + b.Cleanup(func() { zerolog.SetGlobalLevel(previousLogLevel) }) + + sourcePath := filepath.Join(b.TempDir(), "input.bin") + out, err := os.Create(sourcePath) + if err != nil { + b.Fatal(err) + } + block := bytes.Repeat([]byte("abcdefgh"), (1<<20)/8) + for written := 0; written < embeddedPutStreamRegularFileBenchmarkSize; written += len(block) { + n, err := out.Write(block) + if err != nil { + _ = out.Close() + b.Fatal(err) + } + if n != len(block) { + _ = out.Close() + b.Fatalf("write source file: wrote %d bytes, want %d", n, len(block)) + } + } + if err := out.Close(); err != nil { + b.Fatal(err) + } + + input, err := os.Open(sourcePath) + if err != nil { + b.Fatal(err) + } + defer input.Close() + + diskPath := b.TempDir() + keys := make([]string, b.N) + for i := range keys { + keys[i] = fmt.Sprintf("benchmark:regular-file:%d", i) + } + + ctx := context.Background() + b.SetBytes(embeddedPutStreamRegularFileBenchmarkSize) + b.ReportAllocs() + b.ResetTimer() + b.StopTimer() + for first := 0; first < b.N; first += embeddedPutStreamRegularFileBenchmarkBatchSize { + end := min(first+embeddedPutStreamRegularFileBenchmarkBatchSize, b.N) + client, err := New(&Config{ + DiskPath: diskPath, + TTL: time.Hour, + Storage: &stor.StorageConfig{ + MetadataCacheSize: 16 << 20, + MetadataBackgroundJobs: 2, + }, + }) + if err != nil { + _ = os.RemoveAll(diskPath) + b.Fatalf("create embedded client: %v", err) + } + + var putErr error + for i := first; i < end; i++ { + if _, err := input.Seek(0, io.SeekStart); err != nil { + putErr = fmt.Errorf("rewind source file: %w", err) + break + } + b.StartTimer() + err := client.PutStream(ctx, keys[i], input, int64(time.Hour.Seconds())) + b.StopTimer() + if err != nil { + putErr = fmt.Errorf("PutStream(%q): %w", keys[i], err) + break + } + } + + var cleanupErr error + for _, key := range keys[first:end] { + if err := client.Delete(ctx, key); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("Delete(%q): %w", key, err)) + } + } + if err := client.Close(); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("close embedded client: %w", err)) + } + if err := os.RemoveAll(diskPath); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("remove benchmark storage: %w", err)) + } + if err := errors.Join(putErr, cleanupErr); err != nil { + b.Fatal(err) + } + } +} diff --git a/storage/prefix_reader_test.go b/storage/prefix_reader_test.go index 14783f64..792edaac 100644 --- a/storage/prefix_reader_test.go +++ b/storage/prefix_reader_test.go @@ -5,8 +5,13 @@ package storage import ( "bytes" + "hash/crc32" "io" + "os" + "path/filepath" "testing" + + "github.com/tigrisdata/ocache/common/bufferpool" ) type prefixReaderReadRecorder struct { @@ -98,6 +103,79 @@ func TestJoinPrefixUsesCopyBufferForReaderOnlyRemainder(t *testing.T) { } } +func TestJoinPrefixRegularFileUsesPooledBuffer(t *testing.T) { + const total = 1 << 20 + prefixSize := DefaultInlineThreshold + 1 + want := bytes.Repeat([]byte("abcdefgh"), total/8) + + dir := t.TempDir() + source, err := os.Create(filepath.Join(dir, "source")) + if err != nil { + t.Fatal(err) + } + if _, err := source.Write(want); err != nil { + source.Close() + t.Fatal(err) + } + info, err := source.Stat() + if err != nil { + source.Close() + t.Fatal(err) + } + if !info.Mode().IsRegular() { + source.Close() + t.Fatalf("source mode = %v, want regular file", info.Mode()) + } + if _, err := source.Seek(int64(prefixSize), io.SeekStart); err != nil { + source.Close() + t.Fatal(err) + } + defer source.Close() + + output, err := os.Create(filepath.Join(dir, "output")) + if err != nil { + t.Fatal(err) + } + var writeSizes []int + counter := writerFunc(func(p []byte) (int, error) { + writeSizes = append(writeSizes, len(p)) + return output.Write(p) + }) + reader := joinPrefix(want[:prefixSize], source) + _, hasWriterTo := reader.(io.WriterTo) + buffer, release := bufferpool.AcquireBuffer(1 << 20) + defer release() + hash := crc32.NewIEEE() + + bytesWritten, err := io.CopyBuffer(io.MultiWriter(counter, hash), reader, buffer) + if err != nil { + output.Close() + t.Fatalf("copy regular-file remainder: %v", err) + } + if err := output.Close(); err != nil { + t.Fatal(err) + } + got, err := os.ReadFile(filepath.Join(dir, "output")) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, want) { + t.Fatal("regular-file copy changed payload order or identity") + } + if bytesWritten != int64(len(want)) { + t.Errorf("bytes written = %d, want %d", bytesWritten, len(want)) + } + if gotChecksum, wantChecksum := hash.Sum32(), crc32.ChecksumIEEE(want); gotChecksum != wantChecksum { + t.Errorf("checksum = %08x, want %08x", gotChecksum, wantChecksum) + } + if hasWriterTo { + t.Error("regular-file remainder unexpectedly preserves io.WriterTo") + } + if len(writeSizes) != 2 || writeSizes[0] != prefixSize || writeSizes[1] != total-prefixSize { + t.Errorf("file-boundary writes = %d calls with sizes %v, want 2 calls [%d %d]", len(writeSizes), writeSizes, prefixSize, total-prefixSize) + } +} + func TestJoinPrefixPreservesRemainderWriterTo(t *testing.T) { prefix := []byte("first chunk") remainder := []byte("remainder") diff --git a/storage/storage.go b/storage/storage.go index a4fca820..e32fbfed 100644 --- a/storage/storage.go +++ b/storage/storage.go @@ -1148,12 +1148,15 @@ func (r *prefixWriterToReader) WriteTo(writer io.Writer) (int64, error) { } // joinPrefix returns the already-read prefix followed by reader. Non-WriterTo -// streams stay reader-only so FileManager's pooled buffer drives the remainder. -// A WriterTo remainder keeps its direct-write path, which is used by unary -// byte-backed Puts. +// streams and *os.File remainders stay reader-only so FileManager's pooled +// buffer controls the copy granularity. Other WriterTo remainders keep their +// direct-write path, which is used by unary byte-backed Puts. func joinPrefix(prefix []byte, reader io.Reader) io.Reader { prefixed := prefixReader{prefix: prefix, reader: reader} if writerTo, ok := reader.(io.WriterTo); ok { + if _, ok := reader.(*os.File); ok { + return &prefixed + } return &prefixWriterToReader{prefixReader: prefixed, writerTo: writerTo} } return &prefixed @@ -1206,9 +1209,9 @@ func (s *Storage) Put(key string, body io.Reader, ttl int) error { // the value length exceeds the small-value threshold. if n > s.inlineThreshold { storageType = "raw_file" - // Join the bytes already read with the remaining stream. A non-WriterTo - // stream uses FileManager's pooled copy buffer, while byte-backed unary - // readers retain their direct-write path. + // Join the bytes already read with the remaining stream. Reader-only + // streams and *os.File inputs use FileManager's pooled buffer; other + // WriterTo readers retain their direct-write path. reader := joinPrefix(firstChunk[:n], body) filePath, checksum, bytesWritten, err := s.fileManager.Write(key, reader) if err != nil {