refactor(storage): consolidate IO reader wrappers

packages/orchestrator/cmd/inspect-build/validate.go

@@ -92,7 +92,7 @@ func validateBuild(ctx context.Context, storagePath, buildID, artifact string, e
 	// fetch still runs and reportValidation shows the checksum as n/a.
 	expected := h.Builds[h.Metadata.BuildId].Checksum
 
-	chunker, obj, size, cleanup, err := openChunker(ctx, storagePath, buildID, artifact, h, ft)
+	chunker, upstream, size, cleanup, err := openChunker(ctx, storagePath, buildID, artifact, h, ft)
 	if err != nil {
 		return err
 	}
@@ -119,7 +119,7 @@ func validateBuild(ctx context.Context, storagePath, buildID, artifact string, e
 				}
 				c := chunks[ci]
 				for off := c.lo; off < c.hi; off += blockSize {
-					if _, err := chunker.Slice(egCtx, off, min(blockSize, c.hi-off), obj, ft); err != nil {
+					if _, err := chunker.Slice(egCtx, off, min(blockSize, c.hi-off), upstream, ft); err != nil {
 						return fmt.Errorf("fetch block at %d: %w", off, err)
 					}
 				}
@@ -133,7 +133,7 @@ func validateBuild(ctx context.Context, storagePath, buildID, artifact string, e
 	// The cache is warm — sweep it one block at a time to hash the image.
 	hasher := sha256.New()
 	for off := int64(0); off < size; off += blockSize {
-		b, err := chunker.Slice(ctx, off, min(blockSize, size-off), obj, ft)
+		b, err := chunker.Slice(ctx, off, min(blockSize, size-off), upstream, ft)
 		if err != nil {
 			return fmt.Errorf("read block at %d: %w", off, err)
 		}
@@ -146,9 +146,8 @@ func validateBuild(ctx context.Context, storagePath, buildID, artifact string, e
 }
 
 // openChunker wires a production block.Chunker over the build's data file,
-// returning the chunker, the upstream storage object, the image's uncompressed
-// size, and a cleanup function.
-func openChunker(ctx context.Context, storagePath, buildID, artifact string, h *header.Header, ft *storage.FrameTable) (*block.Chunker, storage.Seekable, int64, func(), error) {
+// returning the chunker, the image's uncompressed size, and a cleanup function.
+func openChunker(ctx context.Context, storagePath, buildID, artifact string, h *header.Header, ft *storage.FrameTable) (*block.Chunker, storage.RangeOpener, int64, func(), error) {
 	if err := cmdutil.SetupStorage(storagePath); err != nil {
 		return nil, nil, 0, nil, err
 	}

packages/orchestrator/pkg/sandbox/block/streaming_chunk.go

@@ -241,7 +241,7 @@ func (c *Chunker) progressiveRead(ctx context.Context, s *fetchSession, mmapSlic
 		return 0, fmt.Errorf("failed to open range reader at %d: %w", s.chunkOff, err)
 	}
 	defer func() {
-		if closeErr := reader.Close(); closeErr != nil && err == nil {
+		if closeErr := reader.Close(context.WithoutCancel(ctx)); closeErr != nil && err == nil {
 			err = closeErr
 		}
 	}()

packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go

@@ -82,7 +82,7 @@ func (s *fakeSeekable) StoreFile(context.Context, string, ...storage.PutOption)
 	panic("not used")
 }
 
-func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length int64, frameTable *storage.FrameTable) (io.ReadCloser, error) {
+func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length int64, frameTable *storage.FrameTable) (storage.RangeReader, error) {
 	s.fetchCount.Add(1)
 
 	if s.ctrl != nil {
@@ -127,10 +127,10 @@ func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length
 
 	r := io.Reader(bytes.NewReader(s.data[fetchOff:end]))
 	if frameTable.IsCompressed() {
-		return storage.NewDecompressingReader(r, frameTable.CompressionType())
+		return storage.NewDecompressingReader(storage.NewRangeReader(io.NopCloser(r)), frameTable.CompressionType())
 	}
 
-	return io.NopCloser(r), nil
+	return storage.NewRangeReader(io.NopCloser(r)), nil
 }
 
 func makeCompressedTestData(tb testing.TB, data []byte) (*storage.FrameTable, *fakeSeekable) {
@@ -429,7 +429,7 @@ func (s *panicSeekable) StoreFile(context.Context, string, ...storage.PutOption)
 	panic("not used")
 }
 
-func (s *panicSeekable) OpenRangeReader(_ context.Context, off int64, length int64, _ *storage.FrameTable) (io.ReadCloser, error) {
+func (s *panicSeekable) OpenRangeReader(_ context.Context, off int64, length int64, _ *storage.FrameTable) (storage.RangeReader, error) {
 	end := min(off+length, int64(len(s.data)))
 
 	return &panicReader{
@@ -460,7 +460,7 @@ func (r *panicReader) Read(p []byte) (int, error) {
 	return n, nil
 }
 
-func (r *panicReader) Close() error {
+func (r *panicReader) Close(context.Context) error {
 	return nil
 }
 
@@ -587,7 +587,7 @@ func (r *controlledReader) Read(p []byte) (int, error) {
 	return n, nil
 }
 
-func (r *controlledReader) Close() error {
+func (r *controlledReader) Close(context.Context) error {
 	select {
 	case r.closed <- struct{}{}:
 	default:

packages/orchestrator/pkg/sandbox/build/storage_diff_test.go

@@ -188,10 +188,10 @@ func TestStorageDiff_NoRefreshOnFinalizedHeader(t *testing.T) {
 	uncompressedSeekable := storage.NewMockSeekable(t)
 	uncompressedSeekable.EXPECT().
 		OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
-		RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (io.ReadCloser, error) {
+		RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (storage.RangeReader, error) {
 			end := min(off+length, int64(len(payload)))
 
-			return io.NopCloser(bytes.NewReader(payload[off:end])), nil
+			return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), nil
 		})
 	provider.EXPECT().
 		OpenSeekable(mock.Anything, uncompressedPath, mock.Anything).
@@ -235,10 +235,10 @@ func TestStorageDiff_SkipsHeaderRefreshWhenPeerActive(t *testing.T) {
 		Return(int64(payloadSize), nil).Once()
 	uncompressedSeekable.EXPECT().
 		OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
-		RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (io.ReadCloser, error) {
+		RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (storage.RangeReader, error) {
 			end := min(off+length, int64(len(payload)))
 
-			return io.NopCloser(bytes.NewReader(payload[off:end])), nil
+			return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), nil
 		})
 	provider.EXPECT().
 		OpenSeekable(mock.Anything, aPaths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything).
@@ -299,10 +299,10 @@ func TestStorageDiff_V3AncestorFallsBackToUncompressed(t *testing.T) {
 		Return(int64(payloadSize), nil).Once()
 	rawSeekable.EXPECT().
 		OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
-		RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (io.ReadCloser, error) {
+		RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (storage.RangeReader, error) {
 			end := min(off+length, int64(len(payload)))
 
-			return io.NopCloser(bytes.NewReader(payload[off:end])), nil
+			return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), nil
 		})
 	provider.EXPECT().
 		OpenSeekable(mock.Anything, aPaths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything).
@@ -348,10 +348,10 @@ func TestStorageDiff_ReloadSourceLatchesV3AsUncompressed(t *testing.T) {
 	rawSeekable := storage.NewMockSeekable(t)
 	rawSeekable.EXPECT().
 		OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
-		RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (io.ReadCloser, error) {
+		RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (storage.RangeReader, error) {
 			end := min(off+length, int64(len(payload)))
 
-			return io.NopCloser(bytes.NewReader(payload[off:end])), nil
+			return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), nil
 		})
 	provider.EXPECT().
 		OpenSeekable(mock.Anything, aPaths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything).
@@ -434,14 +434,14 @@ func buildHeader(t *testing.T, selfID uuid.UUID, size int64, mapsTo uuid.UUID) *
 // requested U-offset in the caller's frame table, slices the compressed
 // payload, and streams it through a decompressor. Mirrors what a real
 // Seekable does over zstd-compressed data.
-func decompressingRangeReader(compressed []byte) func(context.Context, int64, int64, *storage.FrameTable) (io.ReadCloser, error) {
-	return func(_ context.Context, offsetU, _ int64, ft *storage.FrameTable) (io.ReadCloser, error) {
+func decompressingRangeReader(compressed []byte) func(context.Context, int64, int64, *storage.FrameTable) (storage.RangeReader, error) {
+	return func(_ context.Context, offsetU, _ int64, ft *storage.FrameTable) (storage.RangeReader, error) {
 		r, err := ft.LocateCompressed(offsetU)
 		if err != nil {
 			return nil, err
 		}
 		end := min(r.Offset+int64(r.Length), int64(len(compressed)))
 
-		return storage.NewDecompressingReader(bytes.NewReader(compressed[r.Offset:end]), ft.CompressionType())
+		return storage.NewDecompressingReader(storage.NewRangeReader(io.NopCloser(bytes.NewReader(compressed[r.Offset:end]))), ft.CompressionType())
 	}
 }

packages/orchestrator/pkg/sandbox/template/peerclient/blob.go

@@ -65,7 +65,7 @@ func (b *peerBlob) WriteTo(ctx context.Context, dst io.Writer) (int64, error) {
 			}
 
 			reader := newPeerStreamReader(recv, cancel)
-			defer reader.Close()
+			defer reader.Close(context.WithoutCancel(ctx))
 
 			n, err := io.Copy(dst, reader)
 			if err != nil {

packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"sync"
 	"sync/atomic"
 
@@ -98,9 +97,9 @@ func (s *peerSeekable) Size(ctx context.Context) (int64, error) {
 	return 0, &storage.PeerTransitionedError{}
 }
 
-func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *storage.FrameTable) (io.ReadCloser, error) {
+func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *storage.FrameTable) (storage.RangeReader, error) {
 	res, err := tryPeer(ctx, &s.peerHandle, "peer-seekable-open-range-reader", attrOpRangeReader,
-		func(ctx context.Context) (peerAttempt[io.ReadCloser], error) {
+		func(ctx context.Context) (peerAttempt[storage.RangeReader], error) {
 			streamCtx, cancel := context.WithCancel(ctx)
 
 			recv, err := openPeerSeekableStream(streamCtx, s.client, &orchestrator.ReadAtBuildSeekableRequest{
@@ -113,10 +112,10 @@ func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length in
 				logger.L().Warn(ctx, "failed to open range reader from peer", logger.WithBuildID(s.buildID), zap.Int64("off", off), zap.Int64("length", length), zap.Error(err))
 				cancel()
 
-				return peerAttempt[io.ReadCloser]{}, nil
+				return peerAttempt[storage.RangeReader]{}, nil
 			}
 
-			return peerAttempt[io.ReadCloser]{
+			return peerAttempt[storage.RangeReader]{
 				value: newPeerStreamReader(recv, cancel),
 				hit:   true,
 			}, nil

packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go

@@ -75,7 +75,7 @@ func TestPeerSeekable_OpenRangeReader_PeerSucceeds(t *testing.T) {
 	s := &peerSeekable{peerHandle: peerHandle{client: client, buildID: "build-1", name: storage.MemfileName, uploaded: &atomic.Bool{}}}
 	rc, err := s.OpenRangeReader(t.Context(), 10, int64(len(data)), nil)
 	require.NoError(t, err)
-	defer rc.Close()
+	defer rc.Close(t.Context())
 
 	got, err := io.ReadAll(rc)
 	require.NoError(t, err)
@@ -90,7 +90,7 @@ func TestPeerSeekable_OpenRangeReader_PeerError_FallsBackToBase(t *testing.T) {
 	client.EXPECT().ReadAtBuildSeekable(mock.Anything, mock.Anything).Return(nil, errors.New("peer unavailable"))
 
 	baseSeekable := storage.NewMockSeekable(t)
-	baseSeekable.EXPECT().OpenRangeReader(mock.Anything, int64(0), int64(len(baseData)), (*storage.FrameTable)(nil)).Return(io.NopCloser(bytes.NewReader(baseData)), nil)
+	baseSeekable.EXPECT().OpenRangeReader(mock.Anything, int64(0), int64(len(baseData)), (*storage.FrameTable)(nil)).Return(storage.NewRangeReader(io.NopCloser(bytes.NewReader(baseData))), nil)
 
 	base := storage.NewMockStorageProvider(t)
 	base.EXPECT().OpenSeekable(mock.Anything, "build-1/memfile", storage.MemfileObjectType).Return(baseSeekable, nil)
@@ -107,7 +107,7 @@ func TestPeerSeekable_OpenRangeReader_PeerError_FallsBackToBase(t *testing.T) {
 	}
 	rc, err := s.OpenRangeReader(t.Context(), 0, int64(len(baseData)), nil)
 	require.NoError(t, err)
-	defer rc.Close()
+	defer rc.Close(t.Context())
 
 	got, err := io.ReadAll(rc)
 	require.NoError(t, err)
@@ -185,7 +185,7 @@ func TestPeerStorageProvider_TransitionEmitsError(t *testing.T) {
 	require.NoError(t, err)
 	got, err := io.ReadAll(rc)
 	require.NoError(t, err)
-	require.NoError(t, rc.Close())
+	require.NoError(t, rc.Close(t.Context()))
 	assert.Equal(t, prePeerBytes, got)
 	require.True(t, uploaded.Load(), "uploaded flag should be set after peer EOF with UseStorage")
 

packages/orchestrator/pkg/sandbox/template/peerclient/storage.go

@@ -273,9 +273,9 @@ func tryPeer[T any](
 	return peerAttempt[T]{}, nil
 }
 
-var _ io.ReadCloser = (*peerStreamReader)(nil)
+var _ storage.RangeReader = (*peerStreamReader)(nil)
 
-// peerStreamReader wraps a gRPC streaming recv function as an io.ReadCloser.
+// peerStreamReader wraps a gRPC streaming recv function as a storage.RangeReader.
 // cancel is called on Close to signal the server to terminate the stream.
 type peerStreamReader struct {
 	recv    func() ([]byte, error)
@@ -317,7 +317,7 @@ func (r *peerStreamReader) Read(p []byte) (int, error) {
 	}
 }
 
-func (r *peerStreamReader) Close() error {
+func (r *peerStreamReader) Close(context.Context) error {
 	r.cancel()
 
 	return nil

packages/shared/pkg/storage/compress_decode.go

@@ -1,6 +1,7 @@
 package storage
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"sync"
@@ -9,6 +10,8 @@ import (
 	lz4 "github.com/pierrec/lz4/v4"
 )
 
+var _ RangeReader = (*decompressReader)(nil)
+
 var lz4DecoderPool sync.Pool
 
 func getLZ4Decoder(r io.Reader) *lz4.Reader {
@@ -53,78 +56,47 @@ func putZstdDecoder(dec *zstd.Decoder) {
 	zstdDecoderPool.Put(dec)
 }
 
-// NewDecompressingReader wraps a reader with the appropriate decompressor.
-// Close releases the decompressor back to its pool but does NOT close the
-// underlying reader — the caller is responsible for closing it.
-func NewDecompressingReader(raw io.Reader, ct CompressionType) (io.ReadCloser, error) {
+// decompressReader decompresses inner on Read; Close releases the codec back
+// to its pool and closes inner.
+type decompressReader struct {
+	inner        RangeReader
+	dec          io.Reader
+	releaseCodec func()
+}
+
+func NewDecompressingReader(inner RangeReader, ct CompressionType) (RangeReader, error) {
+	var dec io.Reader
+	var releaseCodec func()
+
 	switch ct {
 	case CompressionLZ4:
-		dec := getLZ4Decoder(raw)
-
-		return &pooledDecoder{
-			Reader: dec,
-			close:  func() { putLZ4Decoder(dec) },
-		}, nil
+		d := getLZ4Decoder(inner)
+		dec, releaseCodec = d, func() { putLZ4Decoder(d) }
 
 	case CompressionZstd:
-		dec, err := getZstdDecoder(raw)
+		d, err := getZstdDecoder(inner)
 		if err != nil {
 			return nil, fmt.Errorf("failed to create zstd decoder: %w", err)
 		}
-
-		return &pooledDecoder{
-			Reader: dec,
-			close:  func() { putZstdDecoder(dec) },
-		}, nil
+		dec, releaseCodec = d, func() { putZstdDecoder(d) }
 
 	default:
 		return nil, fmt.Errorf("unsupported compression type: %s", ct)
 	}
-}
-
-// pooledDecoder wraps a decompressor from a sync.Pool.
-// Close returns the decompressor to the pool.
-type pooledDecoder struct {
-	io.Reader
 
-	close func()
+	return &decompressReader{
+		inner:        inner,
+		dec:          dec,
+		releaseCodec: releaseCodec,
+	}, nil
 }
 
-func (r *pooledDecoder) Close() error {
-	r.close()
-
-	return nil
+func (r *decompressReader) Read(p []byte) (int, error) {
+	return r.dec.Read(p)
 }
 
-// newDecompressingReadCloser wraps raw with the appropriate decompressor and
-// takes ownership: Close releases the decompressor back to the pool AND closes raw.
-func newDecompressingReadCloser(raw io.ReadCloser, ct CompressionType) (io.ReadCloser, error) {
-	dec, err := NewDecompressingReader(raw, ct)
-	if err != nil {
-		return nil, err
-	}
-
-	return &decompressingReadCloser{dec: dec, raw: raw}, nil
-}
-
-// decompressingReadCloser reads from the decompressor and closes both the
-// decompressor (returning it to the pool) and the underlying raw stream.
-type decompressingReadCloser struct {
-	dec io.ReadCloser // decompressor — reads from raw
-	raw io.Closer     // underlying stream
-}
-
-func (c *decompressingReadCloser) Read(p []byte) (int, error) {
-	return c.dec.Read(p)
-}
-
-func (c *decompressingReadCloser) Close() error {
-	decErr := c.dec.Close()
-	rawErr := c.raw.Close()
-
-	if decErr != nil {
-		return decErr
-	}
+func (r *decompressReader) Close(ctx context.Context) error {
+	r.releaseCodec()
 
-	return rawErr
+	return r.inner.Close(ctx)
 }