fix: preserve unread input in local streaming writes

docs/embedded.md

@@ -269,7 +269,11 @@ keys, nextToken, hasMore, err = cache.ListPage(ctx, prefix, limit, continuationT
 
 ### Streaming Operations
 
-For large values, use streaming to avoid loading everything into memory:
+For large values, use streaming to avoid loading everything into memory.
+`PutStream` reads from the reader's current position through EOF. A local
+retryable failure restarts from that call-entry position when it can be obtained
+through `io.Seeker`; otherwise the local write is attempted once. A failed
+retry seek returns an error.
 
 ```go
 // Stream data to cache

embedded/client.go

@@ -417,7 +417,11 @@ func (c *Client) ListPageWithValues(ctx context.Context, prefix string, limit in
 	return entries, token, more, nil
 }
 
-// PutStream stores data from a reader for the given key.
+// PutStream stores bytes from r's current position through EOF. On local writes,
+// the first attempt does not rewind r; retryable failures restart from the
+// call-entry position when it can be obtained through io.Seeker. Readers with
+// no obtainable position are attempted once. A failed retry seek returns an
+// error.
 func (c *Client) PutStream(ctx context.Context, key string, r io.Reader, ttlSeconds int64) error {
 	return c.ops.Put(ctx, key, r, int(ttlSeconds))
 }

embedded/put_stream_test.go

@@ -0,0 +1,203 @@
+// Copyright 2026 Tigris Data, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package embedded
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	stor "github.com/tigrisdata/ocache/storage"
+	storageErrors "github.com/tigrisdata/ocache/storage/errors"
+)
+
+const putStreamReaderPositionInlineThreshold = 8
+
+// readerOnly hides the seek methods of its underlying reader.
+type readerOnly struct {
+	io.Reader
+}
+
+type blockingPositionSeeker struct {
+	*bytes.Reader
+	seekEntered chan struct{}
+	releaseSeek <-chan struct{}
+}
+
+func (r *blockingPositionSeeker) Seek(offset int64, whence int) (int64, error) {
+	if offset == 0 && whence == io.SeekCurrent {
+		select {
+		case r.seekEntered <- struct{}{}:
+		default:
+		}
+		<-r.releaseSeek
+	}
+	return r.Reader.Seek(offset, whence)
+}
+
+func TestPutStreamReaderPosition(t *testing.T) {
+	client, err := New(&Config{
+		DiskPath:        t.TempDir(),
+		TTL:             time.Hour,
+		InlineThreshold: putStreamReaderPositionInlineThreshold,
+		Storage: &stor.StorageConfig{
+			CleanupInterval:     time.Hour,
+			DisableRecompaction: true,
+		},
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = client.Close() })
+
+	ctx := context.Background()
+	storeAndCheck := func(t *testing.T, key string, body io.Reader, want []byte) {
+		t.Helper()
+		require.NoError(t, client.PutStream(ctx, key, body, 0))
+		got, err := client.Get(ctx, key)
+		require.NoError(t, err)
+		require.NotNil(t, got, "PutStream should create the key even when the unread suffix is empty")
+		require.Equal(t, want, got)
+	}
+
+	t.Run("bytes reader starts at current position", func(t *testing.T) {
+		data := []byte("consumed-prefix|unread-bytes")
+		position := int64(len("consumed-prefix|"))
+		body := bytes.NewReader(data)
+		_, err := body.Seek(position, io.SeekStart)
+		require.NoError(t, err)
+		storeAndCheck(t, "bytes-position", body, data[position:])
+	})
+
+	t.Run("file starts at current position", func(t *testing.T) {
+		data := []byte("file-prefix|remaining-file-data")
+		position := int64(len("file-prefix|"))
+		path := filepath.Join(t.TempDir(), "body")
+		require.NoError(t, os.WriteFile(path, data, 0o600))
+		body, err := os.Open(path)
+		require.NoError(t, err)
+		t.Cleanup(func() { _ = body.Close() })
+		_, err = body.Seek(position, io.SeekStart)
+		require.NoError(t, err)
+		storeAndCheck(t, "file-position", body, data[position:])
+	})
+
+	t.Run("zero offset", func(t *testing.T) {
+		data := []byte("whole value from zero")
+		storeAndCheck(t, "zero-position", bytes.NewReader(data), data)
+	})
+
+	t.Run("end of input", func(t *testing.T) {
+		body := bytes.NewReader([]byte("already consumed"))
+		_, err := body.Seek(0, io.SeekEnd)
+		require.NoError(t, err)
+		storeAndCheck(t, "end-position", body, []byte{})
+	})
+
+	t.Run("non-seekable reader starts at current position", func(t *testing.T) {
+		data := []byte("wrapped-prefix|non-seekable-remainder")
+		position := int64(len("wrapped-prefix|"))
+		body := bytes.NewReader(data)
+		_, err := body.Seek(position, io.SeekStart)
+		require.NoError(t, err)
+		storeAndCheck(t, "non-seekable-position", readerOnly{Reader: body}, data[position:])
+	})
+
+	t.Run("pipe-backed file is readable when seek is unsupported", func(t *testing.T) {
+		data := []byte("pipe data from the reader's current position")
+		readEnd, writeEnd, err := os.Pipe()
+		require.NoError(t, err)
+		writerDone := make(chan error, 1)
+		go func() {
+			n, writeErr := writeEnd.Write(data)
+			closeErr := writeEnd.Close()
+			if writeErr == nil && n != len(data) {
+				writeErr = io.ErrShortWrite
+			}
+			if writeErr != nil {
+				writerDone <- writeErr
+				return
+			}
+			writerDone <- closeErr
+		}()
+
+		putErr := client.PutStream(ctx, "pipe-file", readEnd, 0)
+		readCloseErr := readEnd.Close()
+		writeErr := <-writerDone
+		require.NoError(t, putErr)
+		require.NoError(t, readCloseErr)
+		require.NoError(t, writeErr)
+		got, err := client.Get(ctx, "pipe-file")
+		require.NoError(t, err)
+		require.Equal(t, data, got)
+	})
+
+	t.Run("canceled request does not consume or store an unpositionable seeker", func(t *testing.T) {
+		data := []byte("already-canceled pipe payload")
+		readEnd, writeEnd, err := os.Pipe()
+		require.NoError(t, err)
+		t.Cleanup(func() { _ = readEnd.Close() })
+		t.Cleanup(func() { _ = writeEnd.Close() })
+
+		n, err := writeEnd.Write(data)
+		require.NoError(t, err)
+		require.Equal(t, len(data), n)
+		require.NoError(t, writeEnd.Close())
+
+		canceled, cancel := context.WithCancel(ctx)
+		cancel()
+		putErr := client.PutStream(canceled, "canceled-pipe", readEnd, 0)
+		gotType, ok := storageErrors.GetType(putErr)
+		require.True(t, ok, "a canceled seekable input must retain its timeout error")
+		require.Equal(t, storageErrors.TypeTimeout, gotType)
+
+		remaining, err := io.ReadAll(readEnd)
+		require.NoError(t, err)
+		require.Equal(t, data, remaining, "cancellation before the storage attempt must not consume the reader")
+		got, err := client.Get(ctx, "canceled-pipe")
+		require.NoError(t, err)
+		require.Nil(t, got, "cancellation before the storage attempt must not create the key")
+	})
+
+	t.Run("pre-canceled request returns before a blocking position query", func(t *testing.T) {
+		data := []byte("already-canceled blocking-seeker payload")
+		releaseSeek := make(chan struct{})
+		body := &blockingPositionSeeker{
+			Reader:      bytes.NewReader(data),
+			seekEntered: make(chan struct{}, 1),
+			releaseSeek: releaseSeek,
+		}
+		canceled, cancel := context.WithCancel(ctx)
+		cancel()
+		done := make(chan error, 1)
+		go func() {
+			done <- client.PutStream(canceled, "canceled-blocking-seeker", body, 0)
+		}()
+
+		select {
+		case putErr := <-done:
+			gotType, ok := storageErrors.GetType(putErr)
+			require.True(t, ok, "a canceled seekable input must retain its timeout error")
+			require.Equal(t, storageErrors.TypeTimeout, gotType)
+		case <-body.seekEntered:
+			close(releaseSeek)
+			putErr := <-done
+			t.Fatalf("PutStream queried the reader before the canceled-context check: %v", putErr)
+		case <-time.After(5 * time.Second):
+			close(releaseSeek)
+			putErr := <-done
+			t.Fatalf("PutStream did not return before a blocked position query: %v", putErr)
+		}
+
+		remaining, err := io.ReadAll(body.Reader)
+		require.NoError(t, err)
+		require.Equal(t, data, remaining, "the canceled request must not consume the reader")
+		got, err := client.Get(ctx, "canceled-blocking-seeker")
+		require.NoError(t, err)
+		require.Nil(t, got, "the canceled request must not create the key")
+	})
+}

server/operations/put.go

@@ -73,10 +73,10 @@ func (o *Operations) OverBudget() bool {
 	return o.storage.OverShedRatio()
 }
 
-// PutLocal stores data in local storage directly.
-// If the body implements io.Seeker (e.g., bytes.Reader), retries are enabled with automatic reset.
-// For non-seekable readers (e.g., io.PipeReader from streaming), no retries are attempted
-// since the reader cannot be rewound.
+// PutLocal stores data in local storage directly. If the reader's current
+// position can be obtained through io.Seeker, retryable failures restart from
+// that call-entry position. Otherwise, storage is attempted once from the
+// reader's current position.
 func (o *Operations) PutLocal(ctx context.Context, key string, body io.Reader, ttl int) error {
 	return o.PutLocalWithOptions(ctx, key, body, ttl, PutOptions{})
 }
@@ -88,20 +88,46 @@ func (o *Operations) PutLocalWithOptions(ctx context.Context, key string, body i
 		return storageErrors.NewOverBudgetError("Put")
 	}
 
-	// Check if we can retry (reader is seekable)
-	seeker, canRetry := body.(io.Seeker)
-
-	if !canRetry {
-		// Non-seekable reader (e.g., io.PipeReader) - no retry possible
+	return putLocalWithRetry(ctx, key, body, func() error {
 		return o.storage.Put(key, body, ttl)
+	})
+}
+
+// putLocalWithRetry attempts an unpositionable reader once. When its initial
+// offset is available, the first attempt reads from the current position and
+// each retry restores that same offset.
+func putLocalWithRetry(ctx context.Context, key string, body io.Reader, put func() error) error {
+	seeker, ok := body.(io.Seeker)
+	if !ok {
+		return put()
+	}
+	if ctx.Err() != nil {
+		// Match DoWithKey's pre-attempt check before calling a caller-supplied Seek.
+		return storageErrors.NewTimeoutError("Put", key)
 	}
 
-	// Seekable reader (e.g., bytes.Reader) - retry with reset
+	start, err := seeker.Seek(0, io.SeekCurrent)
+	if err != nil || start < 0 {
+		// Keep DoWithKey's pre-attempt cancellation check, but do not retry a
+		// reader whose original position cannot be restored.
+		cfg := retry.DefaultConfig()
+		cfg.MaxRetries = 0
+		return retry.DoWithKey(ctx, cfg, "Put", key, put)
+	}
+
+	attempt := 0
 	return retry.DoWithKey(ctx, retry.DefaultConfig(), "Put", key, func() error {
-		if _, err := seeker.Seek(0, io.SeekStart); err != nil {
-			return err
+		if attempt > 0 {
+			if _, err := seeker.Seek(start, io.SeekStart); err != nil {
+				if storageErrors.IsRetryable(err) {
+					// A failed reposition is terminal even if its error resembles a storage error.
+					return errors.New(err.Error())
+				}
+				return err
+			}
 		}
-		return o.storage.Put(key, body, ttl)
+		attempt++
+		return put()
 	})
 }