fix(orchestrator): guard chunk reads against cache unmapping

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

@@ -68,31 +68,76 @@ func NewChunker(
 // ReadAt and Slice take {upstream, ft} as a paired snapshot from the caller.
 // The caller is responsible for keeping them consistent.
 func (c *Chunker) ReadAt(ctx context.Context, b []byte, off int64, upstream storage.RangeOpener, ft *storage.FrameTable) (int, error) {
-	slice, err := c.Slice(ctx, off, int64(len(b)), upstream, ft)
+	if err := c.ensureData(ctx, off, int64(len(b)), upstream, ft); err != nil {
+		return 0, fmt.Errorf("failed to ensure cache at %d-%d: %w", off, off+int64(len(b)), err)
+	}
+
+	// Fetching may need its own cache read lock. Acquire the close guard only
+	// after readiness, and retain it until the copy has finished. A completed
+	// prefix is readable even before the entire chunk is marked cached.
+	slice, release, err := c.acquireReadySlice(off, int64(len(b)))
 	if err != nil {
-		return 0, fmt.Errorf("failed to slice cache at %d-%d: %w", off, off+int64(len(b)), err)
+		return 0, fmt.Errorf("failed to read ready cache at %d-%d: %w", off, off+int64(len(b)), err)
 	}
+	defer release()
 
 	return copy(b, slice), nil
 }
 
 func (c *Chunker) Slice(ctx context.Context, off, length int64, upstream storage.RangeOpener, ft *storage.FrameTable) ([]byte, error) {
+	if err := c.ensureData(ctx, off, length, upstream, ft); err != nil {
+		return nil, err
+	}
+
+	// The returned slice is borrowed; callers must keep the cache alive while
+	// using it. The guard protects reading the mmap header against Close.
+	slice, release, err := c.acquireReadySlice(off, length)
+	if err != nil {
+		return nil, fmt.Errorf("failed to slice ready cache at %d-%d: %w", off, off+length, err)
+	}
+	defer release()
+
+	return slice, nil
+}
+
+// acquireReadySlice requires ensureData to have completed. It protects the mmap
+// header and lets ReadAt keep Close blocked until the borrowed bytes are copied.
+func (c *Chunker) acquireReadySlice(off, length int64) ([]byte, func(), error) {
+	c.cache.mu.RLock()
+	slice, err := c.cache.Slice(off, length)
+	if errors.As(err, &BytesNotAvailableError{}) {
+		// A streamed prefix can be ready before the full chunk is marked cached.
+		slice, err = c.cache.sliceDirect(off, length)
+	}
+	if err != nil {
+		c.cache.mu.RUnlock()
+
+		return nil, nil, err
+	}
+
+	return slice, c.cache.mu.RUnlock, nil
+}
+
+// ensureData waits for the requested bytes without exposing a borrowed mmap slice.
+func (c *Chunker) ensureData(ctx context.Context, off, length int64, upstream storage.RangeOpener, ft *storage.FrameTable) error {
 	ct := ft.CompressionType()
 	sliceStart := time.Now()
 
 	// Fast path: already cached.
-	b, err := c.cache.Slice(off, length)
+	c.cache.mu.RLock()
+	_, err := c.cache.Slice(off, length)
+	c.cache.mu.RUnlock()
 	if err == nil {
 		c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), length, storage.OKAttrs(c.objType, storage.SourceMmap, ct))
 		storage.ReportSource(ctx, storage.SourceMmap)
 
-		return b, nil
+		return nil
 	}
 
 	if !errors.As(err, &BytesNotAvailableError{}) {
 		c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), 0, storage.ErrAttrs(c.objType, storage.SourceMmap, ct, err))
 
-		return nil, fmt.Errorf("failed read from cache at offset %d: %w", off, err)
+		return fmt.Errorf("failed read from cache at offset %d: %w", off, err)
 	}
 
 	// Fetch every chunk the range spans (one fetch session per chunk).
@@ -103,32 +148,24 @@ func (c *Chunker) Slice(ctx context.Context, off, length int64, upstream storage
 		if lerr != nil {
 			c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), 0, storage.ErrAttrs(c.objType, src, ct, lerr))
 
-			return nil, fmt.Errorf("failed to locate chunk for offset %d: %w", cur, lerr)
+			return fmt.Errorf("failed to locate chunk for offset %d: %w", cur, lerr)
 		}
 		chunkEnd := chunkOff + chunkLen
 		rangeEnd := min(end, chunkEnd)
 		s, err := c.fetch(ctx, cur, rangeEnd-cur, upstream, ft)
 		if err != nil {
 			c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), 0, storage.ErrAttrs(c.objType, s, ct, err))
 
-			return nil, fmt.Errorf("failed to ensure data at %d-%d: %w", cur, rangeEnd, err)
+			return fmt.Errorf("failed to ensure data at %d-%d: %w", cur, rangeEnd, err)
 		}
 		src = max(src, s)
 		cur = chunkEnd
 	}
 
-	// sliceDirect skips isCached — the waiter already confirmed the data is in the mmap.
-	b, cacheErr := c.cache.sliceDirect(off, length)
-	if cacheErr != nil {
-		c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), 0, storage.ErrAttrs(c.objType, src, ct, cacheErr))
-
-		return nil, fmt.Errorf("failed to read from cache after ensuring data at %d-%d: %w", off, off+length, cacheErr)
-	}
-
 	c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), length, storage.OKAttrs(c.objType, src, ct))
 	storage.ReportSource(ctx, src)
 
-	return b, nil
+	return nil
 }
 
 // getOrCreateSession returns a fetch session for the chunk at [off, off+length),

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

@@ -0,0 +1,127 @@
+//go:build linux
+
+package block
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/metric"
+	"go.opentelemetry.io/otel/metric/noop"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+)
+
+// Closing from the Slice metric deterministically schedules eviction after the
+// cached range was found but before ReadAt copies it. No timing or sleeps needed.
+func TestChunker_ReadAtClosedAfterReadiness(t *testing.T) {
+	t.Parallel()
+
+	chunker := newTestChunker(t, testBlockSize)
+	payload := makeTestData(int(testBlockSize))
+	_, err := chunker.cache.WriteAt(payload, 0)
+	require.NoError(t, err)
+
+	closed := false
+	t.Cleanup(func() {
+		if !closed {
+			require.NoError(t, chunker.Close())
+		}
+	})
+	meter := closeOnRecordMeter{
+		Meter: noop.NewMeterProvider().Meter("test"),
+		onRecord: func() {
+			require.False(t, closed)
+			require.NoError(t, chunker.Close())
+			closed = true
+		},
+	}
+	chunker.metrics.ChunkSliceTimerFactory, err = telemetry.NewFloatTimerFactory(meter, "test.slice", "duration", "bytes")
+	require.NoError(t, err)
+
+	dst := bytes.Repeat([]byte{0xff}, int(testBlockSize))
+	var n int
+	err = RunFaultSafe(t.Context(), func() error {
+		var readErr error
+		n, readErr = chunker.ReadAt(t.Context(), dst, 0, nil, nil)
+
+		return readErr
+	})
+	require.True(t, closed)
+	var closedErr *CacheClosedError
+	require.ErrorAs(t, err, &closedErr)
+	require.Zero(t, n)
+	require.Equal(t, bytes.Repeat([]byte{0xff}, int(testBlockSize)), dst)
+}
+
+// ReadAt must still return a streamed prefix while the rest of the chunk is
+// blocked upstream: the full-chunk cache bitmap is not a readiness requirement.
+func TestChunker_ReadAtPartialChunk(t *testing.T) {
+	t.Parallel()
+
+	payload := makeTestData(testFileSize)
+	chunker := newControlledChunker(t, payload)
+	t.Cleanup(func() {
+		close(chunker.advance)
+		<-chunker.closed
+		require.NoError(t, chunker.Close())
+	})
+
+	dst := make([]byte, testBlockSize)
+	type result struct {
+		n   int
+		err error
+	}
+	done := make(chan result, 1)
+	go func() {
+		n, err := chunker.ReadAt(t.Context(), dst, 0, chunker.file, nil)
+		done <- result{n: n, err: err}
+	}()
+	chunker.advance <- struct{}{}
+	<-chunker.consumed
+	r := <-done
+	require.NoError(t, r.err)
+	require.Equal(t, len(dst), r.n)
+	require.Equal(t, payload[:testBlockSize], dst)
+	require.False(t, chunker.cache.isCached(0, int64(len(payload))))
+}
+
+func TestChunker_EmptyReadAtEnd(t *testing.T) {
+	t.Parallel()
+
+	chunker := newTestChunker(t, testBlockSize)
+	t.Cleanup(func() { require.NoError(t, chunker.Close()) })
+	n, err := chunker.ReadAt(t.Context(), nil, testBlockSize, nil, nil)
+	require.NoError(t, err)
+	require.Zero(t, n)
+	slice, err := chunker.Slice(t.Context(), testBlockSize, 0, nil, nil)
+	require.NoError(t, err)
+	require.Empty(t, slice)
+}
+
+type closeOnRecordMeter struct {
+	metric.Meter
+
+	onRecord func()
+}
+
+func (m closeOnRecordMeter) Float64Histogram(name string, options ...metric.Float64HistogramOption) (metric.Float64Histogram, error) {
+	h, err := m.Meter.Float64Histogram(name, options...)
+	if err != nil {
+		return nil, err
+	}
+
+	return closeOnRecordHistogram{Float64Histogram: h, onRecord: m.onRecord}, nil
+}
+
+type closeOnRecordHistogram struct {
+	metric.Float64Histogram
+
+	onRecord func()
+}
+
+func (h closeOnRecordHistogram) Record(context.Context, float64, ...metric.RecordOption) {
+	h.onRecord()
+}