refactor(storage): consolidate IO reader wrappers

Unify the scattered io.ReadCloser wrapper types into a single
io_wrappers.go file with interface assertions as a TOC:
offsetReader, instrumentedReader, cancelReader, sectionReader.

Author: levb

packages/shared/pkg/storage/io_wrappers.go

@@ -0,0 +1,126 @@
+package storage
+
+import (
+	"context"
+	"io"
+	"os"
+
+	"go.opentelemetry.io/otel/trace"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+)
+
+// io.ReadCloser / io.Reader wrappers used in the storage read path.
+var (
+	_ io.Reader     = (*offsetReader)(nil)       // ReaderAt → sequential Reader adapter
+	_ io.ReadCloser = (*instrumentedReader)(nil) // OTEL timer or span on Close
+	_ io.ReadCloser = (*cancelReader)(nil)       // context cancel on Close
+	_ io.ReadCloser = (*sectionReader)(nil)      // os.File SectionReader for range reads
+)
+
+// offsetReader adapts an io.ReaderAt into a sequential io.Reader
+// starting at the given offset.
+type offsetReader struct {
+	wrapped io.ReaderAt
+	offset  int64
+}
+
+func (r *offsetReader) Read(p []byte) (n int, err error) {
+	n, err = r.wrapped.ReadAt(p, r.offset)
+	r.offset += int64(n)
+
+	return
+}
+
+func newOffsetReader(reader io.ReaderAt, offset int64) *offsetReader {
+	return &offsetReader{reader, offset}
+}
+
+// instrumentedReader wraps an io.ReadCloser with OTEL instrumentation.
+// Counts bytes read, tracks read errors, and on Close records the timer
+// as Success/Failure and/or ends the span. Construct via newTimedReader
+// or newSpanReader; exactly one of timer/span is set per call site.
+type instrumentedReader struct {
+	inner   io.ReadCloser
+	timer   *telemetry.Stopwatch
+	span    trace.Span
+	ctx     context.Context //nolint:containedctx // needed to record timer/span on Close
+	bytes   int64
+	readErr error
+}
+
+func newTimedReader(inner io.ReadCloser, timer *telemetry.Stopwatch, ctx context.Context) *instrumentedReader {
+	return &instrumentedReader{inner: inner, timer: timer, ctx: ctx}
+}
+
+func newSpanReader(inner io.ReadCloser, span trace.Span, ctx context.Context) *instrumentedReader {
+	return &instrumentedReader{inner: inner, span: span, ctx: ctx}
+}
+
+func (r *instrumentedReader) Read(p []byte) (int, error) {
+	n, err := r.inner.Read(p)
+	r.bytes += int64(n)
+
+	if err != nil && err != io.EOF {
+		r.readErr = err
+	}
+
+	return n, err
+}
+
+func (r *instrumentedReader) Close() error {
+	closeErr := r.inner.Close()
+
+	if r.timer != nil {
+		if r.readErr != nil || closeErr != nil {
+			r.timer.Failure(r.ctx, r.bytes)
+		} else {
+			r.timer.Success(r.ctx, r.bytes)
+		}
+	}
+
+	if r.span != nil {
+		if closeErr != nil {
+			recordError(r.span, closeErr)
+		} else if r.readErr != nil {
+			recordError(r.span, r.readErr)
+		}
+
+		r.span.End()
+	}
+
+	return closeErr
+}
+
+// cancelReader calls a CancelFunc on Close, ensuring the context used
+// to create the reader is cleaned up.
+type cancelReader struct {
+	io.ReadCloser
+
+	cancel context.CancelFunc
+}
+
+func (r *cancelReader) Close() error {
+	defer r.cancel()
+
+	return r.ReadCloser.Close()
+}
+
+// sectionReader exposes a bounded section of an os.File as a ReadCloser,
+// closing the underlying file on Close.
+type sectionReader struct {
+	*io.SectionReader
+
+	file *os.File
+}
+
+func newSectionReader(f *os.File, off, length int64) *sectionReader {
+	return &sectionReader{
+		SectionReader: io.NewSectionReader(f, off, length),
+		file:          f,
+	}
+}
+
+func (r *sectionReader) Close() error {
+	return r.file.Close()
+}

packages/shared/pkg/storage/offset_reader.go

@@ -83,27 +83,30 @@ func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length
 		attribute.Bool("compressed", compressed),
 	))
 
-	if compressed {
-		rc, err := c.openReaderCompressed(ctx, off, frameTable)
-		if err != nil {
-			recordError(span, err)
-			span.End()
-
-			return nil, err
+	var rc io.ReadCloser
+	var err error
+	switch {
+	case compressed:
+		rc, err = c.openReaderCompressed(ctx, off, frameTable)
+	default:
+		if err = c.validateReadParams(length, off); err == nil {
+			rc, err = c.openReaderUncompressed(ctx, off, length)
 		}
-
-		rc = withSpan(rc, span)
-
-		return rc, nil
 	}
 
-	if err := c.validateReadParams(length, off); err != nil {
+	if err != nil {
 		recordError(span, err)
 		span.End()
 
 		return nil, err
 	}
 
+	return newSpanReader(rc, span, ctx), nil
+}
+
+// openReaderUncompressed is the sequential NFS-then-remote path for
+// uncompressed chunks.
+func (c *cachedSeekable) openReaderUncompressed(ctx context.Context, off, length int64) (io.ReadCloser, error) {
 	timer := cacheSlabReadTimerFactory.Begin(
 		attribute.String(nfsCacheOperationAttr, nfsCacheOperationAttrReadAt),
 		attribute.Bool("compressed", false),
@@ -116,10 +119,7 @@ func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length
 		recordCacheRead(ctx, true, length, cacheTypeSeekable, cacheOpOpenRangeReader)
 		timer.Success(ctx, length)
 
-		rc := io.ReadCloser(&fsRangeReadCloser{Reader: io.NewSectionReader(fp, 0, length), file: fp})
-		rc = withSpan(rc, span)
-
-		return rc, nil
+		return newSectionReader(fp, 0, length), nil
 	}
 
 	if !os.IsNotExist(err) {
@@ -130,50 +130,23 @@ func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length
 
 	rc, err := c.inner.OpenRangeReader(ctx, off, length, nil)
 	if err != nil {
-		recordError(span, err)
-		span.End()
-
 		return nil, fmt.Errorf("failed to open inner range reader: %w", err)
 	}
 
 	recordCacheRead(ctx, false, length, cacheTypeSeekable, cacheOpOpenRangeReader)
 
 	if !skipCacheWriteback(ctx) {
-		rc = newCacheWriteThroughReader(rc, c, ctx, off, length, chunkPath)
+		rc = newCacheWritebackReader(rc, c, ctx, off, length, chunkPath)
 	}
 
-	rc = withSpan(rc, span)
-
 	return rc, nil
 }
 
-// withSpan wraps a reader with an OTEL span that ends on Close.
-func withSpan(rc io.ReadCloser, span trace.Span) io.ReadCloser {
-	return &spanReadCloser{inner: rc, span: span}
-}
-
-type spanReadCloser struct {
-	inner io.ReadCloser
-	span  trace.Span
-}
-
-func (r *spanReadCloser) Read(p []byte) (int, error) {
-	return r.inner.Read(p)
-}
-
-func (r *spanReadCloser) Close() error {
-	err := r.inner.Close()
-	recordError(r.span, err)
-	r.span.End()
-
-	return err
-}
-
-// newCacheWriteThroughReader wraps a reader, buffering all data read through it.
-// On Close, it asynchronously writes the buffered data to the NFS cache only
-// if the total bytes read match the expected length (to avoid caching truncated data).
-func newCacheWriteThroughReader(inner io.ReadCloser, cache *cachedSeekable, ctx context.Context, off, expectedLen int64, chunkPath string) io.ReadCloser {
-	return &cacheWriteThroughReader{
+// writebackReader buffers all data read through it. On Close, it
+// asynchronously writes the buffered data to the NFS cache only if the
+// total bytes read match the expected length (to avoid caching truncated data).
+func newCacheWritebackReader(inner io.ReadCloser, cache *cachedSeekable, ctx context.Context, off, expectedLen int64, chunkPath string) io.ReadCloser {
+	return &cacheWritebackReader{
 		inner:       inner,
 		buf:         bytes.NewBuffer(make([]byte, 0, expectedLen)),
 		cache:       cache,
@@ -184,7 +157,7 @@ func newCacheWriteThroughReader(inner io.ReadCloser, cache *cachedSeekable, ctx
 	}
 }
 
-type cacheWriteThroughReader struct {
+type cacheWritebackReader struct {
 	inner       io.ReadCloser
 	buf         *bytes.Buffer
 	cache       *cachedSeekable
@@ -194,7 +167,7 @@ type cacheWriteThroughReader struct {
 	chunkPath   string
 }
 
-func (r *cacheWriteThroughReader) Read(p []byte) (int, error) {
+func (r *cacheWritebackReader) Read(p []byte) (int, error) {
 	n, err := r.inner.Read(p)
 	if n > 0 {
 		r.buf.Write(p[:n])
@@ -203,12 +176,11 @@ func (r *cacheWriteThroughReader) Read(p []byte) (int, error) {
 	return n, err
 }
 
-func (r *cacheWriteThroughReader) Close() error {
+func (r *cacheWritebackReader) Close() error {
 	closeErr := r.inner.Close()
 
 	// Only cache when the total bytes read match the expected length.
-	// Unlike ReadAt where io.EOF can justify a short read (last chunk),
-	// a streaming reader always ends with EOF regardless of whether the
+	// A streaming reader always ends with EOF regardless of whether the
 	// data was truncated, so the byte count is the only reliable check.
 	if isCompleteRead(r.buf.Len(), int(r.expectedLen), nil) {
 		data := make([]byte, r.buf.Len())

packages/shared/pkg/storage/storage_cache_seekable.go

@@ -10,6 +10,8 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 )
 
+var _ io.ReadCloser = (*decompressingCacheReader)(nil) // decompress on Read, cache compressed bytes on Close
+
 // Precomputed OTEL attributes for compressed cache reads (avoids per-read allocation).
 var compressedCacheReadAttrs = []attribute.KeyValue{
 	attribute.String(nfsCacheOperationAttr, nfsCacheOperationAttrReadAt),
@@ -69,8 +71,8 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64
 	return rc, nil
 }
 
-// newDecompressingCacheReader creates a reader that decompresses on Read and
-// writes the accumulated compressed bytes to the NFS cache on Close.
+// decompressingCacheReader decompresses on Read and writes the accumulated
+// compressed bytes to the NFS cache on Close.
 func newDecompressingCacheReader(
 	raw io.ReadCloser,
 	ct CompressionType,