perf(storage): buffer regular-file PutStream copies

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

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

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 {