diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go index a927f882be..70c9443ce5 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go +++ b/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,7 +148,7 @@ 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) @@ -111,24 +156,16 @@ func (c *Chunker) Slice(ctx context.Context, off, length int64, upstream storage 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), diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_close_test.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_close_test.go new file mode 100644 index 0000000000..70b21f09d0 --- /dev/null +++ b/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() +}