From 596fdb4cd14dbedba4be3bde472ec0d4960ea98b Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Wed, 3 Jun 2026 12:12:49 -0700 Subject: [PATCH 1/6] WIP --- .../pkg/sandbox/block/streaming_chunk.go | 41 ++- .../pkg/sandbox/block/streaming_chunk_test.go | 68 ++-- .../orchestrator/pkg/sandbox/build/build.go | 98 +++--- .../orchestrator/pkg/sandbox/build/diff.go | 12 +- .../pkg/sandbox/build/mocks/mockdiff.go | 20 +- .../pkg/sandbox/build/storage_diff.go | 119 ++++++- .../pkg/sandbox/build/storage_diff_test.go | 322 ++++++++++++++++++ packages/orchestrator/pkg/server/chunks.go | 8 +- packages/shared/pkg/storage/header/header.go | 12 +- .../pkg/storage/header/serialization_test.go | 13 +- packages/shared/pkg/storage/storage.go | 12 +- .../pkg/storage/storage_cache_seekable.go | 4 +- packages/shared/pkg/storage/storage_fs.go | 6 +- packages/shared/pkg/storage/storage_google.go | 6 +- 14 files changed, 593 insertions(+), 148 deletions(-) create mode 100644 packages/orchestrator/pkg/sandbox/build/storage_diff_test.go diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go index f920cdccd0..c21eedbd19 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go @@ -26,7 +26,6 @@ const ( ) type Chunker struct { - upstream storage.StreamingReader cache *Cache metrics metrics.Metrics fetchTimeout time.Duration @@ -38,15 +37,9 @@ type Chunker struct { fetchSessions []*fetchSession } -var ( - _ FramedReader = (*Chunker)(nil) - _ FramedSlicer = (*Chunker)(nil) -) - func NewChunker( ff *featureflags.Client, size, blockSize int64, - upstream storage.StreamingReader, cachePath string, metrics metrics.Metrics, ) (*Chunker, error) { @@ -57,7 +50,6 @@ func NewChunker( return &Chunker{ size: size, - upstream: upstream, cache: cache, metrics: metrics, featureFlags: ff, @@ -65,8 +57,10 @@ func NewChunker( }, nil } -func (c *Chunker) ReadAt(ctx context.Context, b []byte, off int64, ft *storage.FrameTable) (int, error) { - slice, err := c.Slice(ctx, off, int64(len(b)), ft) +// 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 != nil { return 0, fmt.Errorf("failed to slice cache at %d-%d: %w", off, off+int64(len(b)), err) } @@ -74,7 +68,7 @@ func (c *Chunker) ReadAt(ctx context.Context, b []byte, off int64, ft *storage.F return copy(b, slice), nil } -func (c *Chunker) Slice(ctx context.Context, off, length int64, ft *storage.FrameTable) ([]byte, error) { +func (c *Chunker) Slice(ctx context.Context, off, length int64, upstream storage.RangeOpener, ft *storage.FrameTable) ([]byte, error) { attrs := chunkerAttrs if ft.IsCompressed() { attrs = chunkerAttrsCompressed @@ -106,7 +100,7 @@ func (c *Chunker) Slice(ctx context.Context, off, length int64, ft *storage.Fram } chunkEnd := chunkOff + chunkLen rangeEnd := min(end, chunkEnd) - if err := c.fetch(ctx, cur, rangeEnd-cur, ft); err != nil { + if err := c.fetch(ctx, cur, rangeEnd-cur, upstream, ft); err != nil { timer.RecordRaw(ctx, length, attrs.failRemoteFetch) return nil, fmt.Errorf("failed to ensure data at %d-%d: %w", cur, rangeEnd, err) @@ -129,7 +123,7 @@ func (c *Chunker) Slice(ctx context.Context, off, length int64, ft *storage.Fram // getOrCreateSession returns a fetch session for the chunk at [off, off+length), // or (nil, true) if the data is already fully cached. -func (c *Chunker) getOrCreateSession(ctx context.Context, off, length int64, ft *storage.FrameTable) (_ *fetchSession, cached bool) { +func (c *Chunker) getOrCreateSession(ctx context.Context, off, length int64, upstream storage.RangeOpener, ft *storage.FrameTable) (_ *fetchSession, cached bool) { c.fetchMu.Lock() defer c.fetchMu.Unlock() @@ -152,8 +146,9 @@ func (c *Chunker) getOrCreateSession(ctx context.Context, off, length int64, ft // Detach from the caller's cancel signal so the shared fetch goroutine // continues even if the first caller's context is cancelled. Trace/value - // context is preserved for metrics. - go c.runFetch(context.WithoutCancel(ctx), s, ft) + // context is preserved for metrics. The (upstream, ft) pair is captured + // by value here — in-flight sessions are unaffected by later swaps. + go c.runFetch(context.WithoutCancel(ctx), s, upstream, ft) return s, false } @@ -161,13 +156,13 @@ func (c *Chunker) getOrCreateSession(ctx context.Context, off, length int64, ft // fetch ensures the chunk for [off, off+length) is fetched and waits // for every block the range spans (a span can cross block boundaries // after dedup; waiting only on the start block leaves the tail unfetched). -func (c *Chunker) fetch(ctx context.Context, off, length int64, ft *storage.FrameTable) error { +func (c *Chunker) fetch(ctx context.Context, off, length int64, upstream storage.RangeOpener, ft *storage.FrameTable) error { chunkOff, chunkLen, err := c.locateChunk(off, ft) if err != nil { return fmt.Errorf("failed to locate chunk for offset %d: %w", off, err) } - session, justGotCached := c.getOrCreateSession(ctx, chunkOff, chunkLen, ft) + session, justGotCached := c.getOrCreateSession(ctx, chunkOff, chunkLen, upstream, ft) if justGotCached { return nil } @@ -189,7 +184,7 @@ func (c *Chunker) fetch(ctx context.Context, off, length int64, ft *storage.Fram } // runFetch fetches data from storage into the mmap cache. Runs in a background goroutine. -func (c *Chunker) runFetch(ctx context.Context, s *fetchSession, ft *storage.FrameTable) { +func (c *Chunker) runFetch(ctx context.Context, s *fetchSession, upstream storage.RangeOpener, ft *storage.FrameTable) { ctx, cancel := context.WithTimeout(ctx, c.fetchTimeout) defer cancel() @@ -222,7 +217,7 @@ func (c *Chunker) runFetch(ctx context.Context, s *fetchSession, ft *storage.Fra } fetchTimer := c.metrics.RemoteReadsTimerFactory.Begin() - readBytes, err := c.progressiveRead(ctx, s, mmapSlice, ft) + readBytes, err := c.progressiveRead(ctx, s, mmapSlice, upstream, ft) if err != nil { fetchTimer.RecordRaw(ctx, readBytes, attrs.remoteFailure) @@ -240,8 +235,8 @@ func (c *Chunker) runFetch(ctx context.Context, s *fetchSession, ft *storage.Fra s.setDone() } -func (c *Chunker) progressiveRead(ctx context.Context, s *fetchSession, mmapSlice []byte, ft *storage.FrameTable) (totalRead int64, err error) { - reader, err := c.upstream.OpenRangeReader(ctx, s.chunkOff, s.chunkLen, ft) +func (c *Chunker) progressiveRead(ctx context.Context, s *fetchSession, mmapSlice []byte, upstream storage.RangeOpener, ft *storage.FrameTable) (totalRead int64, err error) { + reader, err := upstream.OpenRangeReader(ctx, s.chunkOff, s.chunkLen, ft) if err != nil { return 0, fmt.Errorf("failed to open range reader at %d: %w", s.chunkOff, err) } @@ -330,6 +325,10 @@ func (c *Chunker) IsCached(_ context.Context, off, length int64) bool { return c.cache.isCached(off, length) } +func (c *Chunker) Size() int64 { + return c.size +} + func (c *Chunker) FileSize(ctx context.Context) (int64, error) { return c.cache.FileSize(ctx) } diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go index 2263dc7c4e..7edbbe508c 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go @@ -66,9 +66,9 @@ type testControl struct { onOpen func() // optional callback on OpenRangeReader } -func newTestChunker(t *testing.T, file storage.Seekable, size int64) *Chunker { +func newTestChunker(t *testing.T, size int64) *Chunker { t.Helper() - c, err := NewChunker(&featureflags.Client{}, size, testBlockSize, file, t.TempDir()+"/cache", newTestMetrics(t)) + c, err := NewChunker(&featureflags.Client{}, size, testBlockSize, t.TempDir()+"/cache", newTestMetrics(t)) require.NoError(t, err) return c @@ -151,25 +151,25 @@ func makeCompressedTestData(tb testing.TB, data []byte) (*storage.FrameTable, *f type chunkerTestCase struct { name string - newChunker func(t *testing.T, data []byte) (*Chunker, *storage.FrameTable) + newChunker func(t *testing.T, data []byte) (*Chunker, storage.RangeOpener, *storage.FrameTable) } var allChunkerTestCases = []chunkerTestCase{ { name: "Compressed", - newChunker: func(t *testing.T, data []byte) (*Chunker, *storage.FrameTable) { + newChunker: func(t *testing.T, data []byte) (*Chunker, storage.RangeOpener, *storage.FrameTable) { t.Helper() ft, getter := makeCompressedTestData(t, data) - return newTestChunker(t, getter, int64(len(data))), ft + return newTestChunker(t, int64(len(data))), getter, ft }, }, { name: "Uncompressed", - newChunker: func(t *testing.T, data []byte) (*Chunker, *storage.FrameTable) { + newChunker: func(t *testing.T, data []byte) (*Chunker, storage.RangeOpener, *storage.FrameTable) { t.Helper() - return newTestChunker(t, &fakeSeekable{data: data}, int64(len(data))), nil + return newTestChunker(t, int64(len(data))), &fakeSeekable{data: data}, nil }, }, } @@ -182,10 +182,10 @@ func TestChunker_BasicSlice(t *testing.T) { t.Parallel() data := makeTestData(testFileSize) - chunker, ft := tc.newChunker(t, data) + chunker, file, ft := tc.newChunker(t, data) defer chunker.Close() - slice, err := chunker.Slice(t.Context(), 0, testBlockSize, ft) + slice, err := chunker.Slice(t.Context(), 0, testBlockSize, file, ft) require.NoError(t, err) require.Equal(t, data[:testBlockSize], slice) }) @@ -201,11 +201,11 @@ func TestChunker_CacheHit(t *testing.T) { // Uncompressed only — we need direct access to the fakeSeekable to count fetches. file := &fakeSeekable{data: data} - chunker := newTestChunker(t, file, int64(len(data))) + chunker := newTestChunker(t, int64(len(data))) defer chunker.Close() // First read triggers a fetch. - slice1, err := chunker.Slice(t.Context(), 0, testBlockSize, nil) + slice1, err := chunker.Slice(t.Context(), 0, testBlockSize, file, nil) require.NoError(t, err) require.Equal(t, data[:testBlockSize], slice1) @@ -213,7 +213,7 @@ func TestChunker_CacheHit(t *testing.T) { require.Positive(t, firstFetches) // Second read of the same block — should hit cache. - slice2, err := chunker.Slice(t.Context(), 0, testBlockSize, nil) + slice2, err := chunker.Slice(t.Context(), 0, testBlockSize, file, nil) require.NoError(t, err) require.Equal(t, data[:testBlockSize], slice2) require.Equal(t, firstFetches, file.fetchCount.Load(), "expected no additional upstream fetch") @@ -230,17 +230,17 @@ func TestChunker_FullChunkCachedAfterPartialRequest(t *testing.T) { t.Parallel() data := makeTestData(testFileSize) - chunker, ft := tc.newChunker(t, data) + chunker, file, ft := tc.newChunker(t, data) defer chunker.Close() - _, err := chunker.Slice(t.Context(), 0, testBlockSize, ft) + _, err := chunker.Slice(t.Context(), 0, testBlockSize, file, ft) require.NoError(t, err) // The second Slice joins the in-flight session (or hits // cache if the fetch already completed). Either way it blocks // until the data is available — no polling needed. lastOff := int64(testFileSize) - testBlockSize - slice, err := chunker.Slice(t.Context(), lastOff, testBlockSize, ft) + slice, err := chunker.Slice(t.Context(), lastOff, testBlockSize, file, ft) require.NoError(t, err) require.Equal(t, data[lastOff:lastOff+testBlockSize], slice) }) @@ -342,14 +342,15 @@ func TestChunker_ErrorKeepsPartialData(t *testing.T) { data := makeTestData(testFileSize) - chunker := newTestChunker(t, &fakeSeekable{data: data, failAfter: int64(testFileSize / 2)}, int64(len(data))) + file := &fakeSeekable{data: data, failAfter: int64(testFileSize / 2)} + chunker := newTestChunker(t, int64(len(data))) defer chunker.Close() lastOff := int64(testFileSize) - testBlockSize - _, err := chunker.Slice(t.Context(), lastOff, testBlockSize, nil) + _, err := chunker.Slice(t.Context(), lastOff, testBlockSize, file, nil) require.Error(t, err) - slice, err := chunker.Slice(t.Context(), 0, testBlockSize, nil) + slice, err := chunker.Slice(t.Context(), 0, testBlockSize, file, nil) require.NoError(t, err) require.Equal(t, data[:testBlockSize], slice) } @@ -399,13 +400,13 @@ func TestChunker_LastBlockPartial(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() - chunker, ft := tc.newChunker(t, data) + chunker, file, ft := tc.newChunker(t, data) defer chunker.Close() lastBlockOff := (int64(size) / testBlockSize) * testBlockSize remaining := int64(size) - lastBlockOff - slice, err := chunker.Slice(t.Context(), lastBlockOff, remaining, ft) + slice, err := chunker.Slice(t.Context(), lastBlockOff, remaining, file, ft) require.NoError(t, err) require.Equal(t, data[lastBlockOff:], slice) }) @@ -469,16 +470,17 @@ func TestChunker_PanicRecovery(t *testing.T) { data := makeTestData(testFileSize) panicAt := int64(testFileSize / 2) - chunker := newTestChunker(t, &panicSeekable{data: data, panicAfter: panicAt}, int64(len(data))) + file := &panicSeekable{data: data, panicAfter: panicAt} + chunker := newTestChunker(t, int64(len(data))) defer chunker.Close() // Request data past the panic point — should get an error, not hang or crash lastOff := int64(testFileSize) - testBlockSize - _, err := chunker.Slice(t.Context(), lastOff, testBlockSize, nil) + _, err := chunker.Slice(t.Context(), lastOff, testBlockSize, file, nil) require.Error(t, err) // Data before the panic point should still be cached - slice, err := chunker.Slice(t.Context(), 0, testBlockSize, nil) + slice, err := chunker.Slice(t.Context(), 0, testBlockSize, file, nil) require.NoError(t, err) require.Equal(t, data[:testBlockSize], slice) } @@ -491,7 +493,7 @@ func TestChunker_ConcurrentStress(t *testing.T) { t.Parallel() data := makeTestData(testFileSize) - chunker, ft := tc.newChunker(t, data) + chunker, file, ft := tc.newChunker(t, data) defer chunker.Close() const numGoroutines = 50 @@ -504,7 +506,7 @@ func TestChunker_ConcurrentStress(t *testing.T) { eg.Go(func() error { for j := range opsPerGoroutine { off := int64(((i*opsPerGoroutine)+j)%(len(data)/int(readLen))) * readLen - slice, err := chunker.Slice(t.Context(), off, readLen, ft) + slice, err := chunker.Slice(t.Context(), off, readLen, file, ft) if err != nil { return fmt.Errorf("goroutine %d op %d: %w", i, j, err) } @@ -522,11 +524,19 @@ func TestChunker_ConcurrentStress(t *testing.T) { } } -// controlledChunker wraps a Chunker with channel-based flow control for tests. -// advance gates reads; opened/consumed/closed signal fetch lifecycle events. +// controlledChunker bundles a Chunker, its upstream, and the channels +// gating reads through that upstream. type controlledChunker struct { *Chunker *testControl + + file *fakeSeekable +} + +// Slice forwards to the embedded Chunker with the bundled upstream — saves +// each test from passing cc.file by hand. +func (cc *controlledChunker) Slice(ctx context.Context, off, length int64, ft *storage.FrameTable) ([]byte, error) { + return cc.Chunker.Slice(ctx, off, length, cc.file, ft) } func newControlledChunker(t *testing.T, data []byte) *controlledChunker { @@ -538,12 +548,12 @@ func newControlledChunker(t *testing.T, data []byte) *controlledChunker { opened: make(chan struct{}, 10), closed: make(chan struct{}, 10), } - file := &fakeSeekable{data: data, ctrl: ctrl} return &controlledChunker{ - Chunker: newTestChunker(t, file, int64(len(data))), + Chunker: newTestChunker(t, int64(len(data))), testControl: ctrl, + file: file, } } diff --git a/packages/orchestrator/pkg/sandbox/build/build.go b/packages/orchestrator/pkg/sandbox/build/build.go index c4814e9558..09b2e43092 100644 --- a/packages/orchestrator/pkg/sandbox/build/build.go +++ b/packages/orchestrator/pkg/sandbox/build/build.go @@ -83,12 +83,6 @@ func (b *File) ReadAt(ctx context.Context, p []byte, off int64) (int, error) { if errors.As(err, &closed) { continue } - // A peer transition swaps the header to the finalized one; retry against it. - if retry, swapErr := b.retryOnTransition(ctx, err); retry { - continue - } else if swapErr != nil { - return 0, swapErr - } return 0, err } @@ -99,7 +93,9 @@ type readSegment struct { srcOff int64 length int64 diff Diff - ft *storage.FrameTable + // ft uses the nil-vs-empty convention: nil = no entry (hint unknown), + // empty &FrameTable{} = authoritatively uncompressed, non-empty = compressed. + ft *storage.FrameTable } func (b *File) readSegments(ctx context.Context, p []byte, segments []readSegment, maxParallel int) error { @@ -123,10 +119,24 @@ func (b *File) readSegments(ctx context.Context, p []byte, segments []readSegmen return nil } +// readSegment reads one segment, recovering from PeerTransitionedError or a +// wrong-CT 404 by refreshing the FT from storage and retrying once. func (b *File) readSegment(ctx context.Context, p []byte, s readSegment) error { - n, err := s.diff.ReadAt(ctx, p[s.dstOff:s.dstOff+int(s.length)], s.srcOff, s.ft) + refresher, _ := s.diff.(FrameTableRefresher) + dst := p[s.dstOff : s.dstOff+int(s.length)] + + n, err := s.diff.ReadAt(ctx, dst, s.srcOff, s.ft) if err != nil { - return err + if refresher == nil { + return err + } + if err = b.specialErrorHandler(ctx, refresher, s, err); err != nil { + return err + } + n, err = s.diff.ReadAt(ctx, dst, s.srcOff, s.ft) + if err != nil { + return err + } } if int64(n) != s.length { return io.ErrUnexpectedEOF @@ -135,6 +145,33 @@ func (b *File) readSegment(ctx context.Context, p []byte, s readSegment) error { return nil } +func (b *File) specialErrorHandler(ctx context.Context, refresher FrameTableRefresher, s readSegment, err error) error { + var transitionErr *storage.PeerTransitionedError + switch { + case errors.As(err, &transitionErr): + if waitErr := waitTransitionBackoff(ctx, b.fileType, transitionErr); waitErr != nil { + return waitErr + } + case errors.Is(err, storage.ErrObjectNotExist) && s.ft == nil: + // fall through to refresh + default: + return err + } + + h, rerr := refresher.RefreshFrameTable(ctx) + if rerr != nil { + return fmt.Errorf("recover after %w: %w", err, rerr) + } + // If the just-loaded header is *this* file's own finalized header, promote + // it: every subsequent read picks up authoritative Builds[ancestor] hints + // instead of paying a separate refresh per ancestor. + if h != nil && h.Metadata.BuildId == b.Header().Metadata.BuildId { + b.SwapHeader(h) + } + + return nil +} + // planRead resolves the segments covering p, zero-filling uuid.Nil regions. // A returned byte count below len(p) means the mappings ran out (EOF). func (b *File) planRead(ctx context.Context, p []byte, off int64) ([]readSegment, int, error) { @@ -193,9 +230,7 @@ func (b *File) cachedBuild(ctx context.Context, h *header.Header, buildID uuid.U } } - // CompressionType is nil-safe (nil frame table -> CompressionNone). - ct := h.GetBuildFrameData(buildID).CompressionType() - diff, err := b.getBuild(ctx, buildID, b.buildFileSize(h, buildID), ct) + diff, err := b.getBuild(ctx, buildID, b.buildFileSize(h, buildID), h.GetBuildFrameData(buildID).CompressionType()) if err != nil { return nil, fmt.Errorf("failed to get build: %w", err) } @@ -230,7 +265,7 @@ func (b *File) Slice(ctx context.Context, off, length int64) ([]byte, error) { } logger.L().Warn(ctx, "failed to slice build fast path", zap.Error(sErr)) } - // Errors fall through to ReadAt's retry-on-transition path. + // Errors fall through to ReadAt. } } } @@ -278,42 +313,21 @@ func (b *File) IsCached(ctx context.Context, off, length int64) bool { return true } -// retryOnTransition catches a PeerTransitionedError and swaps the header from -// storage. Returns (true, nil) to signal the caller should continue the loop, -// or (false, swapErr) if the swap itself failed. RetryAfter backs off repeated -// post-transition storage 404s. -// -// The transition is signaled only after the source upload has finalized, so -// the header object already exists in storage. A single LoadHeader is enough; -// polling here would multiply GCS reads under high peer-transition rates. -func (b *File) retryOnTransition(ctx context.Context, err error) (bool, error) { - var transErr *storage.PeerTransitionedError - if !errors.As(err, &transErr) { - return false, nil - } +// waitTransitionBackoff honors the peer's RetryAfter hint before the caller +// retries against base storage. Returns ctx.Err() if cancelled during sleep. +func waitTransitionBackoff(ctx context.Context, fileType DiffType, transErr *storage.PeerTransitionedError) error { if transErr.RetryAfter > 0 { timer := time.NewTimer(transErr.RetryAfter) defer timer.Stop() - select { case <-timer.C: case <-ctx.Done(): - return false, ctx.Err() + return ctx.Err() } } + logger.L().Info(ctx, "peer transition detected, recovering", zap.String("file_type", string(fileType))) - logger.L().Info(ctx, "peer transition detected, swapping header", - zap.String("file_type", string(b.fileType)), - ) - - hdrPath := storage.Paths{BuildID: b.Header().Metadata.BuildId.String()}.HeaderFile(string(b.fileType)) - h, loadErr := header.LoadHeader(ctx, b.persistence, hdrPath) - if loadErr != nil { - return false, fmt.Errorf("failed to swap header: %w", loadErr) - } - b.SwapHeader(h) - - return true, nil + return nil } // buildFileSize returns the uncompressed file size for a build. Returns 0 for @@ -326,7 +340,7 @@ func (b *File) buildFileSize(h *header.Header, buildID uuid.UUID) int64 { return 0 } -func (b *File) getBuild(ctx context.Context, buildID uuid.UUID, uncompressedSize int64, ct storage.CompressionType) (Diff, error) { +func (b *File) getBuild(ctx context.Context, buildID uuid.UUID, uncompressedSize int64, hintCT storage.CompressionType) (Diff, error) { storageDiff, err := newStorageDiff( b.store.cachePath, buildID.String(), @@ -334,7 +348,7 @@ func (b *File) getBuild(ctx context.Context, buildID uuid.UUID, uncompressedSize int64(b.Header().Metadata.BlockSize), b.metrics, b.persistence, - uncompressedSize, ct, + uncompressedSize, hintCT, b.store.flags, ) if err != nil { diff --git a/packages/orchestrator/pkg/sandbox/build/diff.go b/packages/orchestrator/pkg/sandbox/build/diff.go index 0e3761a4f7..62109928bf 100644 --- a/packages/orchestrator/pkg/sandbox/build/diff.go +++ b/packages/orchestrator/pkg/sandbox/build/diff.go @@ -11,6 +11,7 @@ import ( "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block" "github.com/e2b-dev/infra/packages/shared/pkg/id" "github.com/e2b-dev/infra/packages/shared/pkg/storage" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" ) type DiffType string @@ -28,15 +29,24 @@ const ( type Diff interface { io.Closer - storage.SeekableReader + block.FramedReader block.FramedSlicer CacheKey() DiffStoreKey CachePath(ctx context.Context) (string, error) + // Size returns the logical (uncompressed, U-space) file size. + Size(ctx context.Context) (int64, error) + // FileSize returns the number of bytes resident in the local cache file + // on disk. Used by the DiffStore evictor. FileSize(ctx context.Context) (int64, error) BlockSize() int64 Init(ctx context.Context) error } +type FrameTableRefresher interface { + FrameTable() *storage.FrameTable // nil if the authoritative FT isn't installed yet + RefreshFrameTable(ctx context.Context) (*header.Header, error) +} + type NoDiff struct{} var _ Diff = (*NoDiff)(nil) diff --git a/packages/orchestrator/pkg/sandbox/build/mocks/mockdiff.go b/packages/orchestrator/pkg/sandbox/build/mocks/mockdiff.go index 88a18debca..4eedded6ca 100644 --- a/packages/orchestrator/pkg/sandbox/build/mocks/mockdiff.go +++ b/packages/orchestrator/pkg/sandbox/build/mocks/mockdiff.go @@ -345,8 +345,8 @@ func (_c *MockDiff_Init_Call) RunAndReturn(run func(ctx context.Context) error) } // ReadAt provides a mock function for the type MockDiff -func (_mock *MockDiff) ReadAt(ctx context.Context, buffer []byte, off int64, ft *storage.FrameTable) (int, error) { - ret := _mock.Called(ctx, buffer, off, ft) +func (_mock *MockDiff) ReadAt(ctx context.Context, p []byte, off int64, ft *storage.FrameTable) (int, error) { + ret := _mock.Called(ctx, p, off, ft) if len(ret) == 0 { panic("no return value specified for ReadAt") @@ -355,15 +355,15 @@ func (_mock *MockDiff) ReadAt(ctx context.Context, buffer []byte, off int64, ft var r0 int var r1 error if returnFunc, ok := ret.Get(0).(func(context.Context, []byte, int64, *storage.FrameTable) (int, error)); ok { - return returnFunc(ctx, buffer, off, ft) + return returnFunc(ctx, p, off, ft) } if returnFunc, ok := ret.Get(0).(func(context.Context, []byte, int64, *storage.FrameTable) int); ok { - r0 = returnFunc(ctx, buffer, off, ft) + r0 = returnFunc(ctx, p, off, ft) } else { r0 = ret.Get(0).(int) } if returnFunc, ok := ret.Get(1).(func(context.Context, []byte, int64, *storage.FrameTable) error); ok { - r1 = returnFunc(ctx, buffer, off, ft) + r1 = returnFunc(ctx, p, off, ft) } else { r1 = ret.Error(1) } @@ -377,14 +377,14 @@ type MockDiff_ReadAt_Call struct { // ReadAt is a helper method to define mock.On call // - ctx context.Context -// - buffer []byte +// - p []byte // - off int64 // - ft *storage.FrameTable -func (_e *MockDiff_Expecter) ReadAt(ctx interface{}, buffer interface{}, off interface{}, ft interface{}) *MockDiff_ReadAt_Call { - return &MockDiff_ReadAt_Call{Call: _e.mock.On("ReadAt", ctx, buffer, off, ft)} +func (_e *MockDiff_Expecter) ReadAt(ctx interface{}, p interface{}, off interface{}, ft interface{}) *MockDiff_ReadAt_Call { + return &MockDiff_ReadAt_Call{Call: _e.mock.On("ReadAt", ctx, p, off, ft)} } -func (_c *MockDiff_ReadAt_Call) Run(run func(ctx context.Context, buffer []byte, off int64, ft *storage.FrameTable)) *MockDiff_ReadAt_Call { +func (_c *MockDiff_ReadAt_Call) Run(run func(ctx context.Context, p []byte, off int64, ft *storage.FrameTable)) *MockDiff_ReadAt_Call { _c.Call.Run(func(args mock.Arguments) { var arg0 context.Context if args[0] != nil { @@ -417,7 +417,7 @@ func (_c *MockDiff_ReadAt_Call) Return(n int, err error) *MockDiff_ReadAt_Call { return _c } -func (_c *MockDiff_ReadAt_Call) RunAndReturn(run func(ctx context.Context, buffer []byte, off int64, ft *storage.FrameTable) (int, error)) *MockDiff_ReadAt_Call { +func (_c *MockDiff_ReadAt_Call) RunAndReturn(run func(ctx context.Context, p []byte, off int64, ft *storage.FrameTable) (int, error)) *MockDiff_ReadAt_Call { _c.Call.Return(run) return _c } diff --git a/packages/orchestrator/pkg/sandbox/build/storage_diff.go b/packages/orchestrator/pkg/sandbox/build/storage_diff.go index eea5c839f6..226bccb081 100644 --- a/packages/orchestrator/pkg/sandbox/build/storage_diff.go +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff.go @@ -5,18 +5,32 @@ package build import ( "context" "fmt" + "sync" + "sync/atomic" + + "github.com/google/uuid" "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block" blockmetrics "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block/metrics" "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" "github.com/e2b-dev/infra/packages/shared/pkg/storage" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" ) +// authoritativeFullFrameTable: nil = not fetched; empty = uncompressed. +type source struct { + upstream storage.RangeOpener + authoritativeFullFrameTable *storage.FrameTable +} + +// StorageDiff opens at the parent's hinted CT, refreshing from this build's +// own header sidecar if the hint turns out stale. type StorageDiff struct { chunker *block.Chunker cachePath string cacheKey DiffStoreKey - storagePath string + buildID string + diffType DiffType storageObjectType storage.SeekableObjectType blockSize int64 @@ -24,6 +38,11 @@ type StorageDiff struct { persistence storage.StorageProvider featureFlags *featureflags.Client uncompressedSize int64 + // initCT picks Init's storage path; never reaches read decoding. + initCT storage.CompressionType + + source atomic.Pointer[source] + refreshMu sync.Mutex } var _ Diff = (*StorageDiff)(nil) @@ -44,7 +63,7 @@ func newStorageDiff( metrics blockmetrics.Metrics, persistence storage.StorageProvider, uncompressedSize int64, - ct storage.CompressionType, + hintCT storage.CompressionType, ff *featureflags.Client, ) (*StorageDiff, error) { storageObjectType, ok := storageObjectType(diffType) @@ -54,8 +73,9 @@ func newStorageDiff( cachePath := GenerateDiffCachePath(basePath, buildId, diffType) - return &StorageDiff{ - storagePath: storage.Paths{BuildID: buildId}.DataFile(string(diffType), ct), + d := &StorageDiff{ + buildID: buildId, + diffType: diffType, storageObjectType: storageObjectType, cachePath: cachePath, blockSize: blockSize, @@ -64,7 +84,12 @@ func newStorageDiff( featureFlags: ff, uncompressedSize: uncompressedSize, cacheKey: GetDiffStoreKey(buildId, diffType), - }, nil + initCT: hintCT, + } + // Pre-populate so source.Load() is never nil. + d.source.Store(&source{}) + + return d, nil } func storageObjectType(diffType DiffType) (storage.SeekableObjectType, bool) { @@ -83,7 +108,8 @@ func (b *StorageDiff) CacheKey() DiffStoreKey { } func (b *StorageDiff) Init(ctx context.Context) error { - obj, err := b.persistence.OpenSeekable(ctx, b.storagePath, b.storageObjectType) + storagePath := storage.Paths{BuildID: b.buildID}.DataFile(string(b.diffType), b.initCT) + obj, err := b.persistence.OpenSeekable(ctx, storagePath, b.storageObjectType) if err != nil { return err } @@ -96,16 +122,32 @@ func (b *StorageDiff) Init(ctx context.Context) error { } } - c, err := block.NewChunker(b.featureFlags, size, b.blockSize, obj, b.cachePath, b.metrics) + c, err := block.NewChunker(b.featureFlags, size, b.blockSize, b.cachePath, b.metrics) if err != nil { return fmt.Errorf("failed to create chunker: %w", err) } b.chunker = c + b.source.Store(&source{upstream: obj}) return nil } +// resolveSource: authoritative FT (if bound) wins over caller's hint. +func (b *StorageDiff) resolveSource(callerFT *storage.FrameTable) (storage.RangeOpener, *storage.FrameTable) { + cur := b.source.Load() + ft := callerFT + if cur.authoritativeFullFrameTable != nil { + ft = cur.authoritativeFullFrameTable + } + + return cur.upstream, ft +} + +func (b *StorageDiff) FrameTable() *storage.FrameTable { + return b.source.Load().authoritativeFullFrameTable +} + func (b *StorageDiff) Close() error { if b.chunker == nil { return nil @@ -114,12 +156,59 @@ func (b *StorageDiff) Close() error { return b.chunker.Close() } -func (b *StorageDiff) ReadAt(ctx context.Context, p []byte, off int64, ft *storage.FrameTable) (int, error) { - return b.chunker.ReadAt(ctx, p, off, ft) +func (b *StorageDiff) ReadAt(ctx context.Context, p []byte, off int64, callerFT *storage.FrameTable) (int, error) { + up, ft := b.resolveSource(callerFT) + + return b.chunker.ReadAt(ctx, p, off, up, ft) } -func (b *StorageDiff) Slice(ctx context.Context, off, length int64, ft *storage.FrameTable) ([]byte, error) { - return b.chunker.Slice(ctx, off, length, ft) +func (b *StorageDiff) Slice(ctx context.Context, off, length int64, callerFT *storage.FrameTable) ([]byte, error) { + up, ft := b.resolveSource(callerFT) + + return b.chunker.Slice(ctx, off, length, up, ft) +} + +// RefreshFrameTable binds the authoritative full self-FT from this build's +// own header sidecar and reopens the upstream at the resulting CT. Idempotent +// one-shot latch; post-Init only. Returns the loaded header on success so +// the caller can promote it (e.g. swap into File.header on a self-uuid +// match). Returns (nil, nil) when the latch short-circuits. +func (b *StorageDiff) RefreshFrameTable(ctx context.Context) (*header.Header, error) { + b.refreshMu.Lock() + defer b.refreshMu.Unlock() + + if b.source.Load().authoritativeFullFrameTable != nil { + return nil, nil + } + + headerPath := storage.Paths{BuildID: b.buildID}.HeaderFile(string(b.diffType)) + h, err := header.LoadHeader(ctx, b.persistence, headerPath) + if err != nil { + return nil, fmt.Errorf("load own header for build %s: %w", b.buildID, err) + } + + bid, err := uuid.Parse(b.buildID) + if err != nil { + return nil, fmt.Errorf("parse build id %s: %w", b.buildID, err) + } + ft := h.GetBuildFrameData(bid) + if ft == nil { + // Own header missing its own entry — shouldn't happen, but fall + // through to an empty FT so the latch fires. + ft = &storage.FrameTable{} + } + + newPath := storage.Paths{BuildID: b.buildID}.DataFile(string(b.diffType), ft.CompressionType()) + newObj, err := b.persistence.OpenSeekable(ctx, newPath, b.storageObjectType) + if err != nil { + return nil, fmt.Errorf("open upstream at %s: %w", newPath, err) + } + + // ft was just extracted from the header (so is authoritative) for self-UUID + // (so is full, covers the entire diff). + b.source.Store(&source{upstream: newObj, authoritativeFullFrameTable: ft}) + + return h, nil } // The local file might not be synced. @@ -135,8 +224,12 @@ func (b *StorageDiff) FileSize(ctx context.Context) (int64, error) { return b.chunker.FileSize(ctx) } -func (b *StorageDiff) Size(ctx context.Context) (int64, error) { - return b.FileSize(ctx) +func (b *StorageDiff) Size(_ context.Context) (int64, error) { + if b.chunker == nil { + return 0, nil + } + + return b.chunker.Size(), nil } func (b *StorageDiff) BlockSize() int64 { diff --git a/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go b/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go new file mode 100644 index 0000000000..fc43828373 --- /dev/null +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go @@ -0,0 +1,322 @@ +//go:build linux + +package build + +import ( + "bytes" + "context" + "errors" + "io" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/metric/noop" + + "github.com/e2b-dev/infra/packages/orchestrator/pkg/cfg" + blockmetrics "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block/metrics" + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" + "github.com/e2b-dev/infra/packages/shared/pkg/storage" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" +) + +// B mid-upload + no Builds[A] entry: Init opens at the hint's (uncompressed) +// path, the chunker's first read hits ErrObjectNotExist, the readSegment +// recovery fetches A's own header sidecar and reopens at the compressed path, +// and the retry succeeds with the decompressed payload. +func TestStorageDiff_RefreshesStaleHintFromAncestorHeader(t *testing.T) { + t.Parallel() + + const ( + blockSize = int64(4 << 10) + frameSizeKB = 256 + payloadSize = 256 * 1024 // single zstd frame + readLen = blockSize + ) + + ctx := t.Context() + aID := uuid.New() + payload := bytes.Repeat([]byte("ancestor-payload-"), payloadSize/len("ancestor-payload-")+1)[:payloadSize] + + aFrameTable, compressed, _, err := storage.CompressBytes(ctx, payload, storage.CompressConfig{ + Enabled: true, + Type: storage.CompressionZstd.String(), + Level: 2, + EncoderConcurrency: 1, + FrameEncodeWorkers: 1, + FrameSizeKB: frameSizeKB, + MinPartSizeMB: 50, + }) + require.NoError(t, err) + + aHeader := buildHeader(t, aID, payloadSize, aID) + aHeader.SetBuild(aID, header.BuildData{Size: int64(payloadSize), FrameData: aFrameTable}) + aHeaderBytes, err := header.SerializeHeader(aHeader) + require.NoError(t, err) + + bHeader := buildHeader(t, uuid.New(), payloadSize, aID) + bHeader.IncompletePendingUpload = true + + aPaths := storage.Paths{BuildID: aID.String()} + provider := storage.NewMockStorageProvider(t) + + // Init opens the hint (uncompressed) path. The chunker's first range + // fetch surfaces ErrObjectNotExist, triggering the recovery hook. + uncompressedSeekable := storage.NewMockSeekable(t) + uncompressedSeekable.EXPECT().Size(mock.Anything).Return(int64(payloadSize), nil) + uncompressedSeekable.EXPECT(). + OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything). + Return(nil, storage.ErrObjectNotExist).Once() + provider.EXPECT(). + OpenSeekable(mock.Anything, aPaths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything). + Return(uncompressedSeekable, nil).Once() + + // Recovery fetches A's own header from the sidecar. + headerBlob := storage.NewMockBlob(t) + headerBlob.EXPECT(). + WriteTo(mock.Anything, mock.Anything). + RunAndReturn(func(_ context.Context, w io.Writer) (int64, error) { + return io.Copy(w, bytes.NewReader(aHeaderBytes)) + }).Once() + provider.EXPECT(). + OpenBlob(mock.Anything, aPaths.HeaderFile(storage.MemfileName), mock.Anything). + Return(headerBlob, nil).Once() + + // Reopen at the compressed path; the retried read decompresses cleanly. + compressedSeekable := storage.NewMockSeekable(t) + compressedSeekable.EXPECT(). + OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything). + RunAndReturn(decompressingRangeReader(compressed)) + provider.EXPECT(). + OpenSeekable(mock.Anything, aPaths.DataFile(storage.MemfileName, storage.CompressionZstd), mock.Anything). + Return(compressedSeekable, nil).Once() + + require.Equal(t, payload[:readLen], runRead(t, bHeader, provider, readLen)) +} + +// Finalized B + nil FrameData on A means "A is uncompressed and that's the +// authoritative answer." The read must serve uncompressed bytes and must not +// fetch A's header sidecar. +func TestStorageDiff_NoRefreshOnFinalizedHeader(t *testing.T) { + t.Parallel() + + const ( + blockSize = int64(4 << 10) + payloadSize = 8 << 10 + readLen = blockSize + ) + + aID := uuid.New() + payload := bytes.Repeat([]byte("uncompressed-A-"), payloadSize/len("uncompressed-A-")+1)[:payloadSize] + + bHeader := buildHeader(t, uuid.New(), payloadSize, aID) + bHeader.SetBuild(aID, header.BuildData{Size: int64(payloadSize), FrameData: nil}) + + uncompressedPath := storage.Paths{BuildID: aID.String()}.DataFile(storage.MemfileName, storage.CompressionNone) + + provider := storage.NewMockStorageProvider(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) { + end := min(off+length, int64(len(payload))) + + return io.NopCloser(bytes.NewReader(payload[off:end])), nil + }) + provider.EXPECT(). + OpenSeekable(mock.Anything, uncompressedPath, mock.Anything). + Return(uncompressedSeekable, nil) + // No expectation on OpenBlob — any header fetch panics the test. + + require.Equal(t, payload[:readLen], runRead(t, bHeader, provider, readLen)) +} + +// runRead wires up a File over a fresh DiffStore + noop metrics and returns +// the first `n` bytes read from offset 0. +func runRead(t *testing.T, h *header.Header, provider storage.StorageProvider, n int64) []byte { + t.Helper() + store, err := NewDiffStore(cfg.Config{}, &featureflags.Client{}, t.TempDir(), time.Hour, time.Minute) + require.NoError(t, err) + m, err := blockmetrics.NewMetrics(noop.NewMeterProvider()) + require.NoError(t, err) + f := NewFile(h, store, Memfile, provider, m) + + buf := make([]byte, n) + got, err := f.ReadAt(t.Context(), buf, 0) + require.NoError(t, err) + require.Equal(t, int(n), got) + + return buf +} + +// testBlockSize is the block size used by buildHeader. Every header-shape +// test in this file uses the same value; if a future test legitimately +// needs a different one, hoist this back into a parameter. +const testBlockSize = int64(4 << 10) + +// buildHeader constructs a V4 header with a single-build self-cover mapping. +func buildHeader(t *testing.T, selfID uuid.UUID, size int64, mapsTo uuid.UUID) *header.Header { + t.Helper() + meta := header.NewTemplateMetadata(selfID, uint64(testBlockSize), uint64(size)) + meta.Version = header.MetadataVersionV4 + h, err := header.NewHeader(meta, []header.BuildMap{{ + Offset: 0, Length: uint64(size), BuildId: mapsTo, BuildStorageOffset: 0, + }}) + require.NoError(t, err) + + return h +} + +// decompressingRangeReader returns a RunAndReturn callback that locates the +// 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) { + 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()) + } +} + +// fakeRefresher is a stand-in for *StorageDiff in tests that exercise +// specialErrorHandler's decision logic without spinning up the chunker. +// loadedHeader is what RefreshFrameTable returns (nil simulates the +// idempotency-latch short-circuit; non-nil simulates a successful load). +type fakeRefresher struct { + ft *storage.FrameTable + loadedHeader *header.Header + refreshCalls int +} + +func (f *fakeRefresher) FrameTable() *storage.FrameTable { return f.ft } + +func (f *fakeRefresher) RefreshFrameTable(context.Context) (*header.Header, error) { + f.refreshCalls++ + + return f.loadedHeader, nil +} + +func TestFile_SpecialErrorDecisionMatrix(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + err error + callerFT *storage.FrameTable + latchedFT *storage.FrameTable + wantRecovery bool + }{ + { + name: "ObjectNotExist, no caller hint, no latched FT triggers refresh", + err: storage.ErrObjectNotExist, + wantRecovery: true, + }, + { + name: "ObjectNotExist, no caller hint, latched FT still triggers refresh (C1)", + err: storage.ErrObjectNotExist, + latchedFT: &storage.FrameTable{}, + wantRecovery: true, + }, + { + name: "ObjectNotExist with caller-authoritative hint does not refresh", + err: storage.ErrObjectNotExist, + callerFT: &storage.FrameTable{}, + wantRecovery: false, + }, + { + name: "PeerTransitionedError triggers refresh after backoff", + err: &storage.PeerTransitionedError{}, + wantRecovery: true, + }, + { + name: "Unrelated errors do not trigger refresh", + err: errors.New("boom"), + wantRecovery: false, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + selfID := uuid.New() + f := NewFile(buildHeader(t, selfID, 4096, selfID), nil, Memfile, nil, blockmetrics.Metrics{}) + refresher := &fakeRefresher{ft: tc.latchedFT} + seg := readSegment{ft: tc.callerFT} + + herr := f.specialErrorHandler(t.Context(), refresher, seg, tc.err) + if tc.wantRecovery { + require.NoError(t, herr, "expected recovery to succeed and signal retry") + require.Equal(t, 1, refresher.refreshCalls) + } else { + require.ErrorIs(t, herr, tc.err, "expected original error to propagate without refresh") + require.Zero(t, refresher.refreshCalls) + } + }) + } +} + +// TestFile_SpecialErrorHandlerPromotesSelfHeader covers the post-refresh +// SwapHeader hook. When the refreshed header's buildID matches the File's +// own, recovery promotes it; an ancestor header (different buildID) is +// dropped, leaving File.header untouched. +func TestFile_SpecialErrorHandlerPromotesSelfHeader(t *testing.T) { + t.Parallel() + + selfID := uuid.New() + ancestorID := uuid.New() + + cases := []struct { + name string + loadedBuildID uuid.UUID + wantSwapped bool + }{ + { + name: "self uuid match promotes loaded header", + loadedBuildID: selfID, + wantSwapped: true, + }, + { + name: "ancestor uuid leaves File.header untouched", + loadedBuildID: ancestorID, + wantSwapped: false, + }, + { + name: "idempotency-latch nil header is a no-op", + loadedBuildID: uuid.Nil, // sentinel meaning "fakeRefresher returns nil" + wantSwapped: false, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + initial := buildHeader(t, selfID, 4096, selfID) + f := NewFile(initial, nil, Memfile, nil, blockmetrics.Metrics{}) + + var loaded *header.Header + if tc.loadedBuildID != uuid.Nil { + loaded = buildHeader(t, tc.loadedBuildID, 4096, tc.loadedBuildID) + } + refresher := &fakeRefresher{loadedHeader: loaded} + seg := readSegment{} + + require.NoError(t, f.specialErrorHandler(t.Context(), refresher, seg, storage.ErrObjectNotExist)) + require.Equal(t, 1, refresher.refreshCalls) + + if tc.wantSwapped { + require.Same(t, loaded, f.Header(), "expected File.header to be swapped to the loaded header") + } else { + require.Same(t, initial, f.Header(), "expected File.header to remain the initial one") + } + }) + } +} diff --git a/packages/orchestrator/pkg/server/chunks.go b/packages/orchestrator/pkg/server/chunks.go index 0e7c3a2085..77c963e14a 100644 --- a/packages/orchestrator/pkg/server/chunks.go +++ b/packages/orchestrator/pkg/server/chunks.go @@ -52,7 +52,7 @@ func toGRPCError(err error) error { func (s *Server) GetBuildFileSize(ctx context.Context, req *orchestrator.GetBuildFileSizeRequest) (*orchestrator.GetBuildFileSizeResponse, error) { telemetry.SetAttributes(ctx, telemetry.WithBuildID(req.GetBuildId()), attribute.String("file_name", req.GetName())) - if s.uploadedBuilds.Get(req.GetBuildId()) != nil { + if s.uploadedBuilds.Has(req.GetBuildId()) { telemetry.SetAttributes(ctx, attribute.Bool("uploaded", true)) return &orchestrator.GetBuildFileSizeResponse{Availability: peerUseStorage}, nil @@ -80,7 +80,7 @@ func (s *Server) GetBuildFileSize(ctx context.Context, req *orchestrator.GetBuil func (s *Server) GetBuildFileExists(ctx context.Context, req *orchestrator.GetBuildFileExistsRequest) (*orchestrator.GetBuildFileExistsResponse, error) { telemetry.SetAttributes(ctx, telemetry.WithBuildID(req.GetBuildId()), attribute.String("file_name", req.GetName())) - if s.uploadedBuilds.Get(req.GetBuildId()) != nil { + if s.uploadedBuilds.Has(req.GetBuildId()) { telemetry.SetAttributes(ctx, attribute.Bool("uploaded", true)) return &orchestrator.GetBuildFileExistsResponse{Availability: peerUseStorage}, nil @@ -124,7 +124,7 @@ func (s *Server) ReadAtBuildSeekable(req *orchestrator.ReadAtBuildSeekableReques attribute.Int64("length", length), ) - if s.uploadedBuilds.Get(req.GetBuildId()) != nil { + if s.uploadedBuilds.Has(req.GetBuildId()) { telemetry.SetAttributes(ctx, attribute.Bool("uploaded", true)) return stream.Send(&orchestrator.ReadAtBuildSeekableResponse{Availability: peerUseStorage}) @@ -159,7 +159,7 @@ func (s *Server) GetBuildBlob(req *orchestrator.GetBuildBlobRequest, stream orch attribute.String("file_name", req.GetName()), ) - if s.uploadedBuilds.Get(req.GetBuildId()) != nil { + if s.uploadedBuilds.Has(req.GetBuildId()) { telemetry.SetAttributes(ctx, attribute.Bool("uploaded", true)) return stream.Send(&orchestrator.GetBuildBlobResponse{Availability: peerUseStorage}) diff --git a/packages/shared/pkg/storage/header/header.go b/packages/shared/pkg/storage/header/header.go index b1852e2ff9..c50969c317 100644 --- a/packages/shared/pkg/storage/header/header.go +++ b/packages/shared/pkg/storage/header/header.go @@ -178,14 +178,18 @@ func (t *Header) GetShiftedMapping(ctx context.Context, offset int64) (BuildMap, return b, nil } -// GetBuildFrameData returns the FrameTable for a build, or nil. -// nil means the build is uncompressed — the caller reads raw bytes instead. +// GetBuildFrameData returns the FrameTable for a build: nil = no entry, +// empty &FrameTable{} = authoritatively uncompressed, else compressed. func (t *Header) GetBuildFrameData(buildID uuid.UUID) *storage.FrameTable { - if t.Builds == nil { + bd, ok := t.Builds[buildID] + if !ok { return nil } + if bd.FrameData == nil { + return &storage.FrameTable{} + } - return t.Builds[buildID].FrameData + return bd.FrameData } func (t *Header) getMapping(ctx context.Context, offset int64) (BuildMap, int64, error) { diff --git a/packages/shared/pkg/storage/header/serialization_test.go b/packages/shared/pkg/storage/header/serialization_test.go index 83adabcaf8..db664a3050 100644 --- a/packages/shared/pkg/storage/header/serialization_test.go +++ b/packages/shared/pkg/storage/header/serialization_test.go @@ -839,7 +839,7 @@ func TestSerializeDeserialize_V4_Uncompressed_SelfEntry(t *testing.T) { require.Equal(t, uint64(MetadataVersionV4), got.Metadata.Version) require.Len(t, got.Builds, 1) require.Contains(t, got.Builds, buildID) - require.Nil(t, got.GetBuildFrameData(buildID)) + require.Equal(t, &storage.FrameTable{}, got.GetBuildFrameData(buildID)) } // Layered chain V4-uncompressed (self) → V4-compressed (mid) → V4-uncompressed @@ -893,8 +893,8 @@ func TestSerializeDeserialize_V4_MixedChain(t *testing.T) { require.Equal(t, 5, got.Mapping.Len()) require.Len(t, got.Builds, 3) - require.Nil(t, got.GetBuildFrameData(selfID)) - require.Nil(t, got.GetBuildFrameData(olderID)) + require.Equal(t, &storage.FrameTable{}, got.GetBuildFrameData(selfID)) + require.Equal(t, &storage.FrameTable{}, got.GetBuildFrameData(olderID)) gotMidFT := got.GetBuildFrameData(midID) require.NotNil(t, gotMidFT) @@ -990,12 +990,11 @@ func TestSerializeDeserialize_V4_CompressedSelfChain(t *testing.T) { require.Equal(t, tc.wantBuild, m.BuildId, "offset %d build", tc.offset) require.Equal(t, tc.wantOffset, m.Offset, "offset %d storage offset", tc.offset) - ft := got.GetBuildFrameData(m.BuildId) + ct := got.GetBuildFrameData(m.BuildId).CompressionType() if tc.compressed { - require.NotNil(t, ft, "offset %d expected FrameData", tc.offset) - require.Equal(t, storage.CompressionZstd, ft.CompressionType(), "offset %d", tc.offset) + require.Equal(t, storage.CompressionZstd, ct, "offset %d", tc.offset) } else { - require.Nil(t, ft, "offset %d expected no FrameData", tc.offset) + require.Equal(t, storage.CompressionNone, ct, "offset %d", tc.offset) } } diff --git a/packages/shared/pkg/storage/storage.go b/packages/shared/pkg/storage/storage.go index 668ca49fb4..dfa341b851 100644 --- a/packages/shared/pkg/storage/storage.go +++ b/packages/shared/pkg/storage/storage.go @@ -140,14 +140,8 @@ type Blob interface { Exists(ctx context.Context) (bool, error) } -type SeekableReader interface { - // Random slice access, off and buffer length must be aligned to block size - ReadAt(ctx context.Context, buffer []byte, off int64, ft *FrameTable) (int, error) - Size(ctx context.Context) (int64, error) -} - -// StreamingReader supports progressive reads via a streaming range reader. -type StreamingReader interface { +// RangeOpener supports progressive reads via a streaming range reader. +type RangeOpener interface { OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (io.ReadCloser, error) } @@ -157,7 +151,7 @@ type SeekableWriter interface { } type Seekable interface { - StreamingReader + RangeOpener SeekableWriter Size(ctx context.Context) (int64, error) } diff --git a/packages/shared/pkg/storage/storage_cache_seekable.go b/packages/shared/pkg/storage/storage_cache_seekable.go index a3b20c1118..08c95441d9 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable.go +++ b/packages/shared/pkg/storage/storage_cache_seekable.go @@ -76,8 +76,8 @@ type cachedSeekable struct { } var ( - _ Seekable = (*cachedSeekable)(nil) - _ StreamingReader = (*cachedSeekable)(nil) + _ Seekable = (*cachedSeekable)(nil) + _ RangeOpener = (*cachedSeekable)(nil) ) func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *FrameTable) (io.ReadCloser, error) { diff --git a/packages/shared/pkg/storage/storage_fs.go b/packages/shared/pkg/storage/storage_fs.go index 25b08f7b10..d834e2d807 100644 --- a/packages/shared/pkg/storage/storage_fs.go +++ b/packages/shared/pkg/storage/storage_fs.go @@ -34,9 +34,9 @@ type fsObject struct { } var ( - _ Seekable = (*fsObject)(nil) - _ Blob = (*fsObject)(nil) - _ StreamingReader = (*fsObject)(nil) + _ Seekable = (*fsObject)(nil) + _ Blob = (*fsObject)(nil) + _ RangeOpener = (*fsObject)(nil) ) type fsRangeReadCloser struct { diff --git a/packages/shared/pkg/storage/storage_google.go b/packages/shared/pkg/storage/storage_google.go index 0556cecb73..cbee0cd867 100644 --- a/packages/shared/pkg/storage/storage_google.go +++ b/packages/shared/pkg/storage/storage_google.go @@ -96,9 +96,9 @@ type gcpObject struct { } var ( - _ Seekable = (*gcpObject)(nil) - _ Blob = (*gcpObject)(nil) - _ StreamingReader = (*gcpObject)(nil) + _ Seekable = (*gcpObject)(nil) + _ Blob = (*gcpObject)(nil) + _ RangeOpener = (*gcpObject)(nil) ) func NewGCP(ctx context.Context, bucketName string, limiter *limit.Limiter) (StorageProvider, error) { From 01ae0ba930aa8546cdc887f20f7eb68352706439 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Mon, 8 Jun 2026 06:58:50 -0700 Subject: [PATCH 2/6] improvement: do not rely on 404s, ever --- .../orchestrator/pkg/sandbox/build/build.go | 247 +++++++++--- .../orchestrator/pkg/sandbox/build/cache.go | 54 ++- .../pkg/sandbox/build/cache_test.go | 28 +- .../orchestrator/pkg/sandbox/build/diff.go | 18 +- .../pkg/sandbox/build/header_load.go | 32 +- .../pkg/sandbox/build/local_diff.go | 4 +- .../pkg/sandbox/build/mocks/mockdiff.go | 102 ++--- .../pkg/sandbox/build/storage_diff.go | 290 ++++++++------ .../pkg/sandbox/build/storage_diff_test.go | 375 ++++++++++-------- .../sandbox/nbd/testutils/template_rootfs.go | 1 + .../pkg/sandbox/template/cache.go | 1 + .../sandbox/template/peerclient/seekable.go | 46 +-- .../template/peerclient/seekable_test.go | 55 +-- .../sandbox/template/peerclient/storage.go | 13 + .../pkg/sandbox/template/storage.go | 2 +- .../pkg/storage/compress_frame_table.go | 6 + packages/shared/pkg/storage/header/header.go | 5 +- .../pkg/storage/header/serialization.go | 17 +- .../pkg/storage/header/serialization_test.go | 6 +- 19 files changed, 800 insertions(+), 502 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/build/build.go b/packages/orchestrator/pkg/sandbox/build/build.go index 09b2e43092..c9fc8c6387 100644 --- a/packages/orchestrator/pkg/sandbox/build/build.go +++ b/packages/orchestrator/pkg/sandbox/build/build.go @@ -11,15 +11,39 @@ import ( "time" "github.com/google/uuid" + "go.opentelemetry.io/otel" "go.uber.org/zap" "golang.org/x/sync/errgroup" "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block" blockmetrics "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block/metrics" + "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/template/peerclient" "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" "github.com/e2b-dev/infra/packages/shared/pkg/logger" "github.com/e2b-dev/infra/packages/shared/pkg/storage" "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" + "github.com/e2b-dev/infra/packages/shared/pkg/telemetry" + "github.com/e2b-dev/infra/packages/shared/pkg/utils" +) + +var ( + buildMeter = otel.Meter("github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/build") + + // frameTableRefreshTimer measures ancestor-header refresh events as the + // standard (duration, bytes, count) triple. A refresh is triggered when a + // parent had no Builds[self] entry (proactive load at source resolution) + // or a P2P peer transitions to storage mid-read. + // + // Attributes: + // cause — proactive | peer_transitioned + // file_type — memfile | rootfs + // result — success | failure (added by TimerFactory) + frameTableRefreshTimer = utils.Must(telemetry.NewTimerFactory(buildMeter, + "orchestrator.storage.diff.frame_table_refresh", + "Duration of frame-table refresh header loads", + "Bytes loaded during frame-table refreshes", + "Frame-table refresh events", + )) ) type File struct { @@ -94,7 +118,8 @@ type readSegment struct { length int64 diff Diff // ft uses the nil-vs-empty convention: nil = no entry (hint unknown), - // empty &FrameTable{} = authoritatively uncompressed, non-empty = compressed. + // storage.UncompressedFrameTable = authoritatively uncompressed, + // non-empty = compressed. ft *storage.FrameTable } @@ -119,19 +144,23 @@ func (b *File) readSegments(ctx context.Context, p []byte, segments []readSegmen return nil } -// readSegment reads one segment, recovering from PeerTransitionedError or a -// wrong-CT 404 by refreshing the FT from storage and retrying once. +// readSegment reads one segment. On PeerTransitionedError, it waits the +// peer's RetryAfter, refreshes the diff's source against the post-finalize +// header/CT, and retries once. All other errors propagate. func (b *File) readSegment(ctx context.Context, p []byte, s readSegment) error { - refresher, _ := s.diff.(FrameTableRefresher) dst := p[s.dstOff : s.dstOff+int(s.length)] n, err := s.diff.ReadAt(ctx, dst, s.srcOff, s.ft) if err != nil { - if refresher == nil { + var transitionErr *storage.PeerTransitionedError + if !errors.As(err, &transitionErr) { return err } - if err = b.specialErrorHandler(ctx, refresher, s, err); err != nil { - return err + if waitErr := waitTransitionBackoff(ctx, b.fileType, transitionErr); waitErr != nil { + return waitErr + } + if refreshErr := s.diff.RefreshSource(ctx); refreshErr != nil { + return fmt.Errorf("refresh after peer transition: %w", refreshErr) } n, err = s.diff.ReadAt(ctx, dst, s.srcOff, s.ft) if err != nil { @@ -145,33 +174,6 @@ func (b *File) readSegment(ctx context.Context, p []byte, s readSegment) error { return nil } -func (b *File) specialErrorHandler(ctx context.Context, refresher FrameTableRefresher, s readSegment, err error) error { - var transitionErr *storage.PeerTransitionedError - switch { - case errors.As(err, &transitionErr): - if waitErr := waitTransitionBackoff(ctx, b.fileType, transitionErr); waitErr != nil { - return waitErr - } - case errors.Is(err, storage.ErrObjectNotExist) && s.ft == nil: - // fall through to refresh - default: - return err - } - - h, rerr := refresher.RefreshFrameTable(ctx) - if rerr != nil { - return fmt.Errorf("recover after %w: %w", err, rerr) - } - // If the just-loaded header is *this* file's own finalized header, promote - // it: every subsequent read picks up authoritative Builds[ancestor] hints - // instead of paying a separate refresh per ancestor. - if h != nil && h.Metadata.BuildId == b.Header().Metadata.BuildId { - b.SwapHeader(h) - } - - return nil -} - // planRead resolves the segments covering p, zero-filling uuid.Nil regions. // A returned byte count below len(p) means the mappings ran out (EOF). func (b *File) planRead(ctx context.Context, p []byte, off int64) ([]readSegment, int, error) { @@ -206,7 +208,7 @@ func (b *File) planRead(ctx context.Context, p []byte, off int64) ([]readSegment continue } - diff, err := b.cachedBuild(ctx, h, mappedToBuild.BuildId, &cacheIDs, &cacheDiffs) + diff, err := b.cachedBuild(ctx, mappedToBuild.BuildId, &cacheIDs, &cacheDiffs) if err != nil { return nil, 0, err } @@ -223,14 +225,14 @@ func (b *File) planRead(ctx context.Context, p []byte, off int64) ([]readSegment return segments, n, nil } -func (b *File) cachedBuild(ctx context.Context, h *header.Header, buildID uuid.UUID, ids *[]uuid.UUID, diffs *[]Diff) (Diff, error) { +func (b *File) cachedBuild(ctx context.Context, buildID uuid.UUID, ids *[]uuid.UUID, diffs *[]Diff) (Diff, error) { for i, id := range *ids { if id == buildID { return (*diffs)[i], nil } } - diff, err := b.getBuild(ctx, buildID, b.buildFileSize(h, buildID), h.GetBuildFrameData(buildID).CompressionType()) + diff, err := b.getBuild(ctx, buildID) if err != nil { return nil, fmt.Errorf("failed to get build: %w", err) } @@ -253,9 +255,8 @@ func (b *File) Slice(ctx context.Context, off, length int64) ([]byte, error) { return header.EmptyHugePage[:length], nil } if m.BuildId != uuid.Nil { - size := b.buildFileSize(h, m.BuildId) ft := h.GetBuildFrameData(m.BuildId) - diff, derr := b.getBuild(ctx, m.BuildId, size, ft.CompressionType()) + diff, derr := b.getBuild(ctx, m.BuildId) if derr != nil { logger.L().Warn(ctx, "failed to get build for slice fast path", zap.Error(derr)) } else { @@ -330,35 +331,173 @@ func waitTransitionBackoff(ctx context.Context, fileType DiffType, transErr *sto return nil } -// buildFileSize returns the uncompressed file size for a build. Returns 0 for -// V3 headers, which signals the read path to fall back to a Size() RPC. -func (b *File) buildFileSize(h *header.Header, buildID uuid.UUID) int64 { - if bd, ok := h.Builds[buildID]; ok { - return bd.Size - } +// getBuild returns the cached StorageDiff for buildID, constructing one via +// createDiff on miss. The singleflight inside GetOrCreate ensures createDiff +// runs at most once per key across concurrent callers. +func (b *File) getBuild(ctx context.Context, buildID uuid.UUID) (Diff, error) { + key := GetDiffStoreKey(buildID.String(), b.fileType) - return 0 + return b.store.GetOrCreate(ctx, key, func(ctx context.Context) (Diff, error) { + return b.createDiff(ctx, buildID) + }) } -func (b *File) getBuild(ctx context.Context, buildID uuid.UUID, uncompressedSize int64, hintCT storage.CompressionType) (Diff, error) { - storageDiff, err := newStorageDiff( +func (b *File) createDiff(ctx context.Context, buildID uuid.UUID) (Diff, error) { + h := b.Header() + blockSize := int64(h.Metadata.BlockSize) + + objType, ok := storageObjectType(b.fileType) + if !ok { + return nil, UnknownDiffTypeError{b.fileType} + } + + var ( + upstream storage.Seekable + size int64 + initialCT storage.CompressionType + initialFT *storage.FrameTable + ) + + bd, hasEntry := h.Builds[buildID] + switch { + case hasEntry: + // Our header has a Builds entry for this ancestor. Do NOT latch + // bd.FrameData as the StorageDiff's authoritative full-file FT — it is + // filtered to only the frames our header references, not the ancestor's + // full table. + size = bd.Size + initialCT = bd.FrameData.CompressionType() + + case h.Metadata.Version >= header.MetadataVersionV4: + peerActive := b.store.isActivePeer != nil && b.store.isActivePeer(buildID.String()) + if peerActive { + // Peer mode is active for the build. Open at the uncompressed path + // (peers serve uncompressed by basic name regardless of stored CT) + // and ask the peer for size. initFT stays nil (as opposed to {}) + // since we do not know what it is. + var err error + upstream, err = b.openUpstream(ctx, buildID, objType, initialCT) + if err != nil { + return nil, err + } + if peerReportedSize, ok, err := initialSize(ctx, upstream); err != nil { + return nil, fmt.Errorf("createDiff: peer Size for build %s: %w", buildID, err) + } else if ok { + size = peerReportedSize + + break + } + + // fall through to refresh. + } + + // Refresh ancestor and open upstream. + var err error + upstream, size, initialFT, err = b.refreshAncestorAndOpenUpstream(ctx, buildID, objType) + if err != nil { + return nil, err + } + + default: + initialFT = storage.UncompressedFrameTable + } + + if upstream == nil { + var err error + upstream, err = b.openUpstream(ctx, buildID, objType, initialCT) + if err != nil { + return nil, err + } + } + + if size == 0 { + // (d) and degenerate (a) where bd.Size was zero. Ask storage directly. + var err error + size, err = upstream.Size(ctx) + if err != nil { + return nil, fmt.Errorf("createDiff: size lookup for build %s: %w", buildID, err) + } + } + + return newStorageDiff( b.store.cachePath, buildID.String(), b.fileType, - int64(b.Header().Metadata.BlockSize), + objType, + blockSize, b.metrics, b.persistence, - uncompressedSize, hintCT, + b.store.isActivePeer, + upstream, + size, + initialFT, b.store.flags, ) +} + +// openUpstream resolves the data-file path for buildID at ct and opens it. +func (b *File) openUpstream(ctx context.Context, buildID uuid.UUID, objType storage.SeekableObjectType, ct storage.CompressionType) (storage.Seekable, error) { + path := storage.Paths{BuildID: buildID.String()}.DataFile(string(b.fileType), ct) + upstream, err := b.persistence.OpenSeekable(ctx, path, objType) if err != nil { - return nil, fmt.Errorf("failed to create storage diff: %w", err) + return nil, fmt.Errorf("createDiff: open upstream for build %s at %s: %w", buildID, path, err) } - source, err := b.store.Get(ctx, storageDiff) + return upstream, nil +} + +func (b *File) refreshAncestorAndOpenUpstream(ctx context.Context, buildID uuid.UUID, objType storage.SeekableObjectType) (storage.Seekable, int64, *storage.FrameTable, error) { + loaded, err := refreshBuildHeader(ctx, b.persistence, buildID, b.fileType, refreshCauseProactive) + if err != nil { + return nil, 0, nil, fmt.Errorf("createDiff: proactive header load for build %s: %w", buildID, err) + } + + // Promote a self-matching loaded header if authoritative. + if h := b.Header(); loaded.Metadata.BuildId == h.Metadata.BuildId { + if _, hasSelf := loaded.Builds[loaded.Metadata.BuildId]; hasSelf { + b.SwapHeader(loaded) + } + } + + bd := loaded.Builds[buildID] + ft := bd.FrameData + if ft == nil { + ft = storage.UncompressedFrameTable + } + + upstream, err := b.openUpstream(ctx, buildID, objType, ft.CompressionType()) if err != nil { - return nil, fmt.Errorf("failed to get build from store: %w", err) + return nil, 0, nil, err + } + + return upstream, bd.Size, ft, nil +} + +// initialSize is THE only production code path that calls Size on +// a freshly opened upstream. Invoked from createDiff when the V4+ ancestor is +// peer-active: ask the peer wrapper for the size. Four outcomes: +// +// - peer-routed wrapper, peer answered → (size, true, nil) +// - peer-routed wrapper, PeerTransitionedError → (0, false, nil) caller refreshes +// - peer-routed wrapper, peer RPC failure → (0, false, err) +// - NOT peer-routed (resolveProvider cleared between IsActive probe and +// OpenSeekable) → (0, false, nil) caller refreshes +// +// Symmetric with readSegment's PeerTransitionedError handling on the read path: +// peer says "go to storage" → refresh authoritative header → continue. +// No 404-driven recovery. +func initialSize(ctx context.Context, upstream storage.Seekable) (size int64, ok bool, err error) { + if _, peerRouted := upstream.(peerclient.PeerRouted); !peerRouted { + return 0, false, nil + } + size, err = upstream.Size(ctx) + if err == nil { + return size, true, nil + } + var transErr *storage.PeerTransitionedError + if errors.As(err, &transErr) { + return 0, false, nil } - return source, nil + return 0, false, err } diff --git a/packages/orchestrator/pkg/sandbox/build/cache.go b/packages/orchestrator/pkg/sandbox/build/cache.go index 7804d0b4c3..0490e120ca 100644 --- a/packages/orchestrator/pkg/sandbox/build/cache.go +++ b/packages/orchestrator/pkg/sandbox/build/cache.go @@ -39,12 +39,13 @@ type deleteDiff struct { } type DiffStore struct { - cachePath string - cache *ttlcache.Cache[DiffStoreKey, Diff] - initGroup singleflight.Group - cancel func() - config cfg.Config - flags *featureflags.Client + cachePath string + cache *ttlcache.Cache[DiffStoreKey, Diff] + initGroup singleflight.Group + cancel func() + config cfg.Config + flags *featureflags.Client + isActivePeer IsActivePeer // pdSizes is used to keep track of the diff sizes // that are scheduled for deletion, as this won't show up in the disk usage. @@ -60,6 +61,7 @@ func NewDiffStore( flags *featureflags.Client, cachePath string, ttl, delay time.Duration, + isActivePeer IsActivePeer, ) (*DiffStore, error) { err := os.MkdirAll(cachePath, 0o755) if err != nil { @@ -71,13 +73,14 @@ func NewDiffStore( ) ds := &DiffStore{ - cachePath: cachePath, - cache: cache, - cancel: func() {}, - config: config, - flags: flags, - pdSizes: make(map[DiffStoreKey]*deleteDiff), - pdDelay: delay, + cachePath: cachePath, + cache: cache, + cancel: func() {}, + config: config, + flags: flags, + isActivePeer: isActivePeer, + pdSizes: make(map[DiffStoreKey]*deleteDiff), + pdDelay: delay, } cache.OnEviction(func(ctx context.Context, _ ttlcache.EvictionReason, item *ttlcache.Item[DiffStoreKey, Diff]) { @@ -118,8 +121,24 @@ func (s *DiffStore) Close() { s.cache.Stop() } -func (s *DiffStore) Get(ctx context.Context, diff Diff) (Diff, error) { - key := diff.CacheKey() +// Get returns the cached Diff for key, refreshing TTL and cancelling any +// pending eviction. Returns (nil, false) if the key isn't present. Use +// Lookup for a side-effect-free probe (e.g. from IsCached). +func (s *DiffStore) Get(key DiffStoreKey) (Diff, bool) { + s.resetDelete(key) + item := s.cache.Get(key) + if item == nil { + return nil, false + } + + return item.Value(), true +} + +// GetOrCreate returns the cached Diff for key, or calls create inside a +// singleflight to construct + cache a new one. The create closure is invoked +// at most once per key across concurrent callers; on success the returned Diff +// is cached and its insertion time recorded. +func (s *DiffStore) GetOrCreate(ctx context.Context, key DiffStoreKey, create func(context.Context) (Diff, error)) (Diff, error) { s.resetDelete(key) if item := s.cache.Get(key); item != nil { @@ -134,7 +153,8 @@ func (s *DiffStore) Get(ctx context.Context, diff Diff) (Diff, error) { insertTime := time.Now() - if err := diff.Init(ctx); err != nil { + diff, err := create(ctx) + if err != nil { return nil, err } @@ -144,7 +164,7 @@ func (s *DiffStore) Get(ctx context.Context, diff Diff) (Diff, error) { return diff, nil }) if err != nil { - return nil, fmt.Errorf("failed to init source: %w", err) + return nil, fmt.Errorf("failed to create diff: %w", err) } return v.(Diff), nil diff --git a/packages/orchestrator/pkg/sandbox/build/cache_test.go b/packages/orchestrator/pkg/sandbox/build/cache_test.go index 405e407378..af1880e193 100644 --- a/packages/orchestrator/pkg/sandbox/build/cache_test.go +++ b/packages/orchestrator/pkg/sandbox/build/cache_test.go @@ -92,6 +92,7 @@ func TestNewDiffStore(t *testing.T) { cachePath, 25*time.Hour, 60*time.Second, + nil, ) require.NoError(t, err) assert.NotNil(t, store) @@ -114,6 +115,7 @@ func TestDiffStoreTTLEviction(t *testing.T) { cachePath, ttl, delay, + nil, ) require.NoError(t, err) @@ -150,6 +152,7 @@ func TestDiffStoreRefreshTTLEviction(t *testing.T) { cachePath, ttl, delay, + nil, ) require.NoError(t, err) @@ -161,8 +164,8 @@ func TestDiffStoreRefreshTTLEviction(t *testing.T) { // Refresh diff expiration time.Sleep(ttl / 2) - _, err = store.Get(t.Context(), diff) - require.NoError(t, err) + _, ok := store.Get(diff.CacheKey()) + require.True(t, ok) // Try to expire diff time.Sleep(ttl/2 + time.Microsecond) @@ -188,6 +191,7 @@ func TestDiffStoreDelayEviction(t *testing.T) { //nolint:paralleltest // very ti cachePath, ttl, delay, + nil, ) require.NoError(t, err) @@ -234,6 +238,7 @@ func TestDiffStoreDelayEvictionAbort(t *testing.T) { //nolint:paralleltest // ve cachePath, ttl, delay, + nil, ) require.NoError(t, err) @@ -256,8 +261,8 @@ func TestDiffStoreDelayEvictionAbort(t *testing.T) { //nolint:paralleltest // ve assert.True(t, dFound) // Abort removal of diff - _, err = store.Get(t.Context(), diff) - require.NoError(t, err) + _, ok := store.Get(diff.CacheKey()) + require.True(t, ok) found = store.Has(diff) assert.True(t, found) @@ -288,6 +293,7 @@ func TestDiffStoreOldestFromCache(t *testing.T) { cachePath, ttl, delay, + nil, ) require.NoError(t, err) @@ -359,6 +365,7 @@ func TestDiffStoreConcurrentEvictionRace(t *testing.T) { cachePath, ttl, delay, + nil, ) require.NoError(t, err) @@ -400,8 +407,7 @@ func TestDiffStoreConcurrentEvictionRace(t *testing.T) { // Occasionally try to access the item, which calls resetDelete if j%5 == 0 { - _, err := store.Get(t.Context(), diff) - assert.NoError(t, err) + store.Get(diff.CacheKey()) } } }(i) @@ -448,6 +454,7 @@ func TestDiffStoreResetDeleteRace(t *testing.T) { cachePath, ttl, delay, + nil, ) require.NoError(t, err) @@ -485,10 +492,9 @@ func TestDiffStoreResetDeleteRace(t *testing.T) { // Small random delay to desynchronize goroutines slightly time.Sleep(time.Duration(iteration%10) * time.Microsecond) - // This call to Get() will trigger resetDelete, which is where the race occurs - // Multiple goroutines calling resetDelete on the same key can race - _, err = store.Get(t.Context(), iterDiff) - assert.NoError(t, err) + // This call will trigger resetDelete, which is where the race occurs. + // Multiple goroutines calling resetDelete on the same key can race. + store.Get(iterDiff.CacheKey()) // Also try direct resetDelete calls to increase race probability store.resetDelete(iterDiff.CacheKey()) @@ -533,6 +539,7 @@ func TestFileIsCached_UUIDNilMappingReportsCached(t *testing.T) { t.TempDir(), time.Hour, time.Minute, + nil, ) require.NoError(t, err) @@ -560,6 +567,7 @@ func TestFileIsCached_UninitializedChunkerReportsUncached(t *testing.T) { t.TempDir(), time.Hour, time.Minute, + nil, ) require.NoError(t, err) diff --git a/packages/orchestrator/pkg/sandbox/build/diff.go b/packages/orchestrator/pkg/sandbox/build/diff.go index 62109928bf..aadb0b3ce0 100644 --- a/packages/orchestrator/pkg/sandbox/build/diff.go +++ b/packages/orchestrator/pkg/sandbox/build/diff.go @@ -11,11 +11,12 @@ import ( "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block" "github.com/e2b-dev/infra/packages/shared/pkg/id" "github.com/e2b-dev/infra/packages/shared/pkg/storage" - "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" ) type DiffType string +type IsActivePeer func(buildID string) bool + type NoDiffError struct{} func (NoDiffError) Error() string { @@ -39,12 +40,11 @@ type Diff interface { // on disk. Used by the DiffStore evictor. FileSize(ctx context.Context) (int64, error) BlockSize() int64 - Init(ctx context.Context) error -} - -type FrameTableRefresher interface { - FrameTable() *storage.FrameTable // nil if the authoritative FT isn't installed yet - RefreshFrameTable(ctx context.Context) (*header.Header, error) + // RefreshSource synchronously re-resolves the diff's upstream data object + // (path) and the frame table by reloading the build's header and reopening + // upstream at the resulting CT. Called when the caller knows the currently + // latched source is stale. To support P2P header swaps. + RefreshSource(ctx context.Context) error } type NoDiff struct{} @@ -79,9 +79,7 @@ func (n *NoDiff) CacheKey() DiffStoreKey { return "" } -func (n *NoDiff) Init(context.Context) error { - return NoDiffError{} -} +func (n *NoDiff) RefreshSource(_ context.Context) error { return nil } func (n *NoDiff) BlockSize() int64 { return 0 diff --git a/packages/orchestrator/pkg/sandbox/build/header_load.go b/packages/orchestrator/pkg/sandbox/build/header_load.go index badc40c7f8..4e55c77998 100644 --- a/packages/orchestrator/pkg/sandbox/build/header_load.go +++ b/packages/orchestrator/pkg/sandbox/build/header_load.go @@ -9,9 +9,12 @@ import ( "time" "github.com/google/uuid" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" "github.com/e2b-dev/infra/packages/shared/pkg/storage" "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" + "github.com/e2b-dev/infra/packages/shared/pkg/utils" ) const ( @@ -20,6 +23,16 @@ const ( loadV4MaxTransientErrors = 3 ) +// uploadHeaderPollWait measures how long the upload-side poll waited for the +// finalized header to become visible in storage. Result attribute is +// "ok" / "deadline_exceeded" / "transient_errors" / "ctx_cancelled" / +// "upload_failed", file_type is memfile|rootfs. +var uploadHeaderPollWait = utils.Must(buildMeter.Int64Histogram( + "orchestrator.storage.upload.header_poll_wait", + metric.WithDescription("Duration of the upload-side wait for the finalized V4 header to appear"), + metric.WithUnit("ms"), +)) + // PollRemoteStorageForHeader polls storage for the post-upload V4 header for buildID/fileType. // ErrObjectNotExist is retried until the budget expires; other LoadHeader // errors are tolerated up to loadV4MaxTransientErrors consecutive occurrences @@ -39,33 +52,50 @@ func PollRemoteStorageForHeader( hint <-chan error, budget time.Duration, ) (*header.Header, error) { + start := time.Now() + result := "ok" + defer func() { + uploadHeaderPollWait.Record(ctx, time.Since(start).Milliseconds(), metric.WithAttributes( + attribute.String("file_type", string(t)), + attribute.String("result", result), + )) + }() + hdrPath := storage.Paths{BuildID: buildID.String()}.HeaderFile(string(t)) deadline := time.Now().Add(budget) backoff := loadV4InitialBackoff transientErrs := 0 for { - h, err := header.LoadHeader(ctx, store, hdrPath) + h, _, err := header.LoadHeader(ctx, store, hdrPath) if err == nil { return h, nil } if !errors.Is(err, storage.ErrObjectNotExist) { transientErrs++ if transientErrs >= loadV4MaxTransientErrors { + result = "transient_errors" + return nil, fmt.Errorf("load V4 header for %s/%s after %d attempts: %w", buildID, t, transientErrs, err) } } else { transientErrs = 0 } if !time.Now().Before(deadline) { + result = "deadline_exceeded" + return nil, fmt.Errorf("V4 header for %s/%s not visible after %s: %w", buildID, t, budget, err) } select { case <-ctx.Done(): + result = "ctx_cancelled" + return nil, ctx.Err() case hintErr := <-hint: if hintErr != nil { + result = "upload_failed" + return nil, fmt.Errorf("upload signaled failure for %s/%s: %w", buildID, t, hintErr) } backoff = loadV4InitialBackoff diff --git a/packages/orchestrator/pkg/sandbox/build/local_diff.go b/packages/orchestrator/pkg/sandbox/build/local_diff.go index 75a4da9ecf..d88242d84c 100644 --- a/packages/orchestrator/pkg/sandbox/build/local_diff.go +++ b/packages/orchestrator/pkg/sandbox/build/local_diff.go @@ -137,9 +137,7 @@ func (b *localDiff) CacheKey() DiffStoreKey { return b.cacheKey } -func (b *localDiff) Init(context.Context) error { - return nil -} +func (b *localDiff) RefreshSource(_ context.Context) error { return nil } func (b *localDiff) BlockSize() int64 { return b.cache.BlockSize() diff --git a/packages/orchestrator/pkg/sandbox/build/mocks/mockdiff.go b/packages/orchestrator/pkg/sandbox/build/mocks/mockdiff.go index 4eedded6ca..8bb106f1f8 100644 --- a/packages/orchestrator/pkg/sandbox/build/mocks/mockdiff.go +++ b/packages/orchestrator/pkg/sandbox/build/mocks/mockdiff.go @@ -293,57 +293,6 @@ func (_c *MockDiff_FileSize_Call) RunAndReturn(run func(ctx context.Context) (in return _c } -// Init provides a mock function for the type MockDiff -func (_mock *MockDiff) Init(ctx context.Context) error { - ret := _mock.Called(ctx) - - if len(ret) == 0 { - panic("no return value specified for Init") - } - - var r0 error - if returnFunc, ok := ret.Get(0).(func(context.Context) error); ok { - r0 = returnFunc(ctx) - } else { - r0 = ret.Error(0) - } - return r0 -} - -// MockDiff_Init_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Init' -type MockDiff_Init_Call struct { - *mock.Call -} - -// Init is a helper method to define mock.On call -// - ctx context.Context -func (_e *MockDiff_Expecter) Init(ctx interface{}) *MockDiff_Init_Call { - return &MockDiff_Init_Call{Call: _e.mock.On("Init", ctx)} -} - -func (_c *MockDiff_Init_Call) Run(run func(ctx context.Context)) *MockDiff_Init_Call { - _c.Call.Run(func(args mock.Arguments) { - var arg0 context.Context - if args[0] != nil { - arg0 = args[0].(context.Context) - } - run( - arg0, - ) - }) - return _c -} - -func (_c *MockDiff_Init_Call) Return(err error) *MockDiff_Init_Call { - _c.Call.Return(err) - return _c -} - -func (_c *MockDiff_Init_Call) RunAndReturn(run func(ctx context.Context) error) *MockDiff_Init_Call { - _c.Call.Return(run) - return _c -} - // ReadAt provides a mock function for the type MockDiff func (_mock *MockDiff) ReadAt(ctx context.Context, p []byte, off int64, ft *storage.FrameTable) (int, error) { ret := _mock.Called(ctx, p, off, ft) @@ -422,6 +371,57 @@ func (_c *MockDiff_ReadAt_Call) RunAndReturn(run func(ctx context.Context, p []b return _c } +// RefreshSource provides a mock function for the type MockDiff +func (_mock *MockDiff) RefreshSource(ctx context.Context) error { + ret := _mock.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for RefreshSource") + } + + var r0 error + if returnFunc, ok := ret.Get(0).(func(context.Context) error); ok { + r0 = returnFunc(ctx) + } else { + r0 = ret.Error(0) + } + return r0 +} + +// MockDiff_RefreshSource_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RefreshSource' +type MockDiff_RefreshSource_Call struct { + *mock.Call +} + +// RefreshSource is a helper method to define mock.On call +// - ctx context.Context +func (_e *MockDiff_Expecter) RefreshSource(ctx interface{}) *MockDiff_RefreshSource_Call { + return &MockDiff_RefreshSource_Call{Call: _e.mock.On("RefreshSource", ctx)} +} + +func (_c *MockDiff_RefreshSource_Call) Run(run func(ctx context.Context)) *MockDiff_RefreshSource_Call { + _c.Call.Run(func(args mock.Arguments) { + var arg0 context.Context + if args[0] != nil { + arg0 = args[0].(context.Context) + } + run( + arg0, + ) + }) + return _c +} + +func (_c *MockDiff_RefreshSource_Call) Return(err error) *MockDiff_RefreshSource_Call { + _c.Call.Return(err) + return _c +} + +func (_c *MockDiff_RefreshSource_Call) RunAndReturn(run func(ctx context.Context) error) *MockDiff_RefreshSource_Call { + _c.Call.Return(run) + return _c +} + // Size provides a mock function for the type MockDiff func (_mock *MockDiff) Size(ctx context.Context) (int64, error) { ret := _mock.Called(ctx) diff --git a/packages/orchestrator/pkg/sandbox/build/storage_diff.go b/packages/orchestrator/pkg/sandbox/build/storage_diff.go index 226bccb081..8212484a33 100644 --- a/packages/orchestrator/pkg/sandbox/build/storage_diff.go +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff.go @@ -9,6 +9,7 @@ import ( "sync/atomic" "github.com/google/uuid" + "go.opentelemetry.io/otel/attribute" "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block" blockmetrics "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block/metrics" @@ -17,14 +18,23 @@ import ( "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" ) -// authoritativeFullFrameTable: nil = not fetched; empty = uncompressed. +const ( + refreshCauseProactive = "proactive" + refreshCausePeerTransitioned = "peer_transitioned" +) + +// source carries the StorageDiff's current routing state. upstream is always +// non-nil after construction but may be switched once over the lifetime; the ft +// pointer's nil/empty/non-empty state encodes the lifecycle: +// +// ft == nil not authoritative. may trigger refresh logic. +// ft == storage.UncompressedFrameTable authoritatively uncompressed (only set by refresh) +// ft non-empty authoritatively compressed with the bound full-file FT type source struct { - upstream storage.RangeOpener - authoritativeFullFrameTable *storage.FrameTable + upstream storage.RangeOpener + fullDiffFrameTable *storage.FrameTable } -// StorageDiff opens at the parent's hinted CT, refreshing from this build's -// own header sidecar if the hint turns out stale. type StorageDiff struct { chunker *block.Chunker cachePath string @@ -33,13 +43,10 @@ type StorageDiff struct { diffType DiffType storageObjectType storage.SeekableObjectType - blockSize int64 - metrics blockmetrics.Metrics - persistence storage.StorageProvider - featureFlags *featureflags.Client - uncompressedSize int64 - // initCT picks Init's storage path; never reaches read decoding. - initCT storage.CompressionType + blockSize int64 + metrics blockmetrics.Metrics + persistence storage.StorageProvider + isActivePeer IsActivePeer source atomic.Pointer[source] refreshMu sync.Mutex @@ -55,39 +62,43 @@ func (e UnknownDiffTypeError) Error() string { return fmt.Sprintf("unknown diff type: %s", e.DiffType) } +// newStorageDiff assembles a StorageDiff from a fully-resolved upstream, size, +// and full-file FrameTable. All regime decisioning (peer-active bootstrap, V3 +// fallback, authoritative-refresh recovery) lives in the caller (createDiff); +// this constructor is intentionally pure. func newStorageDiff( basePath string, - buildId string, + buildID string, diffType DiffType, + storageObjectType storage.SeekableObjectType, blockSize int64, metrics blockmetrics.Metrics, persistence storage.StorageProvider, + isActivePeer IsActivePeer, + upstream storage.Seekable, uncompressedSize int64, - hintCT storage.CompressionType, + initFT *storage.FrameTable, ff *featureflags.Client, ) (*StorageDiff, error) { - storageObjectType, ok := storageObjectType(diffType) - if !ok { - return nil, UnknownDiffTypeError{diffType} + cachePath := GenerateDiffCachePath(basePath, buildID, diffType) + c, err := block.NewChunker(ff, uncompressedSize, blockSize, cachePath, metrics) + if err != nil { + return nil, fmt.Errorf("create chunker for build %s: %w", buildID, err) } - cachePath := GenerateDiffCachePath(basePath, buildId, diffType) - d := &StorageDiff{ - buildID: buildId, + buildID: buildID, diffType: diffType, storageObjectType: storageObjectType, cachePath: cachePath, blockSize: blockSize, metrics: metrics, persistence: persistence, - featureFlags: ff, - uncompressedSize: uncompressedSize, - cacheKey: GetDiffStoreKey(buildId, diffType), - initCT: hintCT, + isActivePeer: isActivePeer, + chunker: c, + cacheKey: GetDiffStoreKey(buildID, diffType), } - // Pre-populate so source.Load() is never nil. - d.source.Store(&source{}) + d.source.Store(&source{upstream: upstream, fullDiffFrameTable: initFT}) return d, nil } @@ -107,47 +118,6 @@ func (b *StorageDiff) CacheKey() DiffStoreKey { return b.cacheKey } -func (b *StorageDiff) Init(ctx context.Context) error { - storagePath := storage.Paths{BuildID: b.buildID}.DataFile(string(b.diffType), b.initCT) - obj, err := b.persistence.OpenSeekable(ctx, storagePath, b.storageObjectType) - if err != nil { - return err - } - - size := b.uncompressedSize - if size == 0 { - size, err = obj.Size(ctx) - if err != nil { - return fmt.Errorf("failed to get object size: %w", err) - } - } - - c, err := block.NewChunker(b.featureFlags, size, b.blockSize, b.cachePath, b.metrics) - if err != nil { - return fmt.Errorf("failed to create chunker: %w", err) - } - - b.chunker = c - b.source.Store(&source{upstream: obj}) - - return nil -} - -// resolveSource: authoritative FT (if bound) wins over caller's hint. -func (b *StorageDiff) resolveSource(callerFT *storage.FrameTable) (storage.RangeOpener, *storage.FrameTable) { - cur := b.source.Load() - ft := callerFT - if cur.authoritativeFullFrameTable != nil { - ft = cur.authoritativeFullFrameTable - } - - return cur.upstream, ft -} - -func (b *StorageDiff) FrameTable() *storage.FrameTable { - return b.source.Load().authoritativeFullFrameTable -} - func (b *StorageDiff) Close() error { if b.chunker == nil { return nil @@ -157,58 +127,21 @@ func (b *StorageDiff) Close() error { } func (b *StorageDiff) ReadAt(ctx context.Context, p []byte, off int64, callerFT *storage.FrameTable) (int, error) { - up, ft := b.resolveSource(callerFT) + up, ft, err := b.resolve(ctx, callerFT) + if err != nil { + return 0, err + } return b.chunker.ReadAt(ctx, p, off, up, ft) } func (b *StorageDiff) Slice(ctx context.Context, off, length int64, callerFT *storage.FrameTable) ([]byte, error) { - up, ft := b.resolveSource(callerFT) - - return b.chunker.Slice(ctx, off, length, up, ft) -} - -// RefreshFrameTable binds the authoritative full self-FT from this build's -// own header sidecar and reopens the upstream at the resulting CT. Idempotent -// one-shot latch; post-Init only. Returns the loaded header on success so -// the caller can promote it (e.g. swap into File.header on a self-uuid -// match). Returns (nil, nil) when the latch short-circuits. -func (b *StorageDiff) RefreshFrameTable(ctx context.Context) (*header.Header, error) { - b.refreshMu.Lock() - defer b.refreshMu.Unlock() - - if b.source.Load().authoritativeFullFrameTable != nil { - return nil, nil - } - - headerPath := storage.Paths{BuildID: b.buildID}.HeaderFile(string(b.diffType)) - h, err := header.LoadHeader(ctx, b.persistence, headerPath) - if err != nil { - return nil, fmt.Errorf("load own header for build %s: %w", b.buildID, err) - } - - bid, err := uuid.Parse(b.buildID) - if err != nil { - return nil, fmt.Errorf("parse build id %s: %w", b.buildID, err) - } - ft := h.GetBuildFrameData(bid) - if ft == nil { - // Own header missing its own entry — shouldn't happen, but fall - // through to an empty FT so the latch fires. - ft = &storage.FrameTable{} - } - - newPath := storage.Paths{BuildID: b.buildID}.DataFile(string(b.diffType), ft.CompressionType()) - newObj, err := b.persistence.OpenSeekable(ctx, newPath, b.storageObjectType) + up, ft, err := b.resolve(ctx, callerFT) if err != nil { - return nil, fmt.Errorf("open upstream at %s: %w", newPath, err) + return nil, err } - // ft was just extracted from the header (so is authoritative) for self-UUID - // (so is full, covers the entire diff). - b.source.Store(&source{upstream: newObj, authoritativeFullFrameTable: ft}) - - return h, nil + return b.chunker.Slice(ctx, off, length, up, ft) } // The local file might not be synced. @@ -246,3 +179,138 @@ func (b *StorageDiff) IsCached(ctx context.Context, off, length int64) bool { return b.chunker.IsCached(ctx, off, length) } + +func refreshBuildHeader(ctx context.Context, persistence storage.StorageProvider, buildID uuid.UUID, diffType DiffType, cause string) (*header.Header, error) { + timer := frameTableRefreshTimer.Begin( + attribute.String("cause", cause), + attribute.String("file_type", string(diffType)), + ) + + headerPath := storage.Paths{BuildID: buildID.String()}.HeaderFile(string(diffType)) + h, bytesLoaded, err := header.LoadHeader(ctx, persistence, headerPath) + if err != nil { + timer.Failure(ctx, int64(bytesLoaded)) + + return nil, fmt.Errorf("load header for build %s: %w", buildID, err) + } + timer.Success(ctx, int64(bytesLoaded)) + + return h, nil +} + +// resolve picks the (upstream, ft) the next read should use, given the +// caller's per-mapping FT hint. The contract: if there is no authoritative FT +// latched AND no peer currently serving this build, we MUST refresh before +// reading. The latched upstream was opened at the bootstrap-guessed CT path +// and may return wrong bytes once the peer is gone; only the authoritative +// header tells us where to read from. +func (b *StorageDiff) resolve(ctx context.Context, callerFT *storage.FrameTable) (storage.RangeOpener, *storage.FrameTable, error) { + cur := b.source.Load() + if cur.fullDiffFrameTable != nil { + return cur.upstream, cur.fullDiffFrameTable, nil + } + if callerFT != nil { + return cur.upstream, callerFT, nil + } + if b.isActivePeer != nil && b.isActivePeer(b.buildID) { + // Peer-active regime: upstream is peer-routed and serves uncompressed + // by basic name. We deliberately do NOT refresh here — the storage + // header may not exist yet and we do not handle ErrNotFound. + return cur.upstream, storage.UncompressedFrameTable, nil + } + if err := b.reloadSource(ctx, refreshCauseProactive); err != nil { + return nil, nil, fmt.Errorf("resolve: %w", err) + } + cur = b.source.Load() + + return cur.upstream, cur.fullDiffFrameTable, nil +} + +// RefreshSource reloads the build's header, latches the authoritative FT, and +// reopens upstream at the resulting CT path. Called by readSegment after a +// PeerTransitionedError. Idempotent: once the source latch is populated, the +// post-refresh upstream is base-routed (no peer wrapper) and cannot emit +// further PeerTransitionedErrors, so a second call is a no-op. +func (b *StorageDiff) RefreshSource(ctx context.Context) error { + return b.reloadSource(ctx, refreshCausePeerTransitioned) +} + +// reloadSource is the idempotent ensure-latched entry. Both RefreshSource +// (PeerTransitionedError) and resolve (read-time peer-left fallback) funnel +// through it; the cause attribute distinguishes them in telemetry. A concurrent +// caller that wins the mutex short-circuits when the latch is already +// populated, so parallel segment reads on a fresh StorageDiff pay only one +// header fetch. +func (b *StorageDiff) reloadSource(ctx context.Context, cause string) error { + b.refreshMu.Lock() + defer b.refreshMu.Unlock() + if b.source.Load().fullDiffFrameTable != nil { + return nil + } + + return b.reloadSourceLocked(ctx, cause) +} + +// loadAuthoritativeHeader fetches the build's storage header and extracts the +// canonical (size, ft). Returns an error if the header lacks a self entry — +// that only happens with peer-served incomplete headers, never with +// storage-uploaded ones (build_upload_v4 always populates self before publish). +// Shared by reloadSourceLocked (read-time refresh) and createDiff +// (construction-time recovery when the peer-active bootstrap finds itself +// routed to base). +func loadAuthoritativeHeader( + ctx context.Context, + persistence storage.StorageProvider, + buildID uuid.UUID, + diffType DiffType, + cause string, +) (int64, *storage.FrameTable, error) { + loaded, err := refreshBuildHeader(ctx, persistence, buildID, diffType, cause) + if err != nil { + return 0, nil, err + } + bd, hasSelf := loaded.Builds[buildID] + if !hasSelf { + return 0, nil, fmt.Errorf("authoritative header for build %s has no self entry (peer-served incomplete?)", buildID) + } + ft := bd.FrameData + if ft == nil { + ft = storage.UncompressedFrameTable // V4+ authoritatively uncompressed + } + + return bd.Size, ft, nil +} + +// reloadSourceLocked re-fetches the header and reopens upstream. Caller must +// hold refreshMu. +// +// V4+ headers on storage always carry a self entry — set unconditionally +// before publish. A missing self entry here can only come from P2P routing +// returning a still-uploading peer's incomplete header. Treating it as "no +// FrameData = uncompressed" would silently corrupt reads of a compressed file. +// Fail loudly; the read path will retry when the peer transitions and the +// storage-authoritative header is available. +// +// V3 never reaches reloadSourceLocked: getBuild's V3 branch latches an +// authoritative empty &{} FT at construction, so resolve short-circuits and +// reloadSource is never called; V3 builds aren't peer-routed so +// PeerTransitionedError never fires against them either. +func (b *StorageDiff) reloadSourceLocked(ctx context.Context, cause string) error { + bid, err := uuid.Parse(b.buildID) + if err != nil { + return fmt.Errorf("parse build id %s: %w", b.buildID, err) + } + _, ft, err := loadAuthoritativeHeader(ctx, b.persistence, bid, b.diffType, cause) + if err != nil { + return fmt.Errorf("reloadSourceLocked: load authoritative header for build %s (cause=%s): %w", b.buildID, cause, err) + } + newPath := storage.Paths{BuildID: b.buildID}.DataFile(string(b.diffType), ft.CompressionType()) + newObj, err := b.persistence.OpenSeekable(ctx, newPath, b.storageObjectType) + if err != nil { + return fmt.Errorf("reloadSourceLocked: reopen upstream for build %s at %s (cause=%s): %w", b.buildID, newPath, cause, err) + } + + b.source.Store(&source{upstream: newObj, fullDiffFrameTable: ft}) + + return nil +} diff --git a/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go b/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go index fc43828373..fb97e2dd62 100644 --- a/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go @@ -5,7 +5,6 @@ package build import ( "bytes" "context" - "errors" "io" "testing" "time" @@ -22,17 +21,17 @@ import ( "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" ) -// B mid-upload + no Builds[A] entry: Init opens at the hint's (uncompressed) -// path, the chunker's first read hits ErrObjectNotExist, the readSegment -// recovery fetches A's own header sidecar and reopens at the compressed path, -// and the retry succeeds with the decompressed payload. -func TestStorageDiff_RefreshesStaleHintFromAncestorHeader(t *testing.T) { +// Parent had no Builds[A] entry (uncompressedSize = 0). Init must load A's +// own header before opening upstream, learn the CT from there, and +// open only the compressed path. The read serves decompressed bytes on the +// first try — no retry loop, no failed open at the uncompressed path. +func TestStorageDiff_LoadsOwnHeaderWhenParentHasNoEntry(t *testing.T) { t.Parallel() const ( blockSize = int64(4 << 10) frameSizeKB = 256 - payloadSize = 256 * 1024 // single zstd frame + payloadSize = 256 * 1024 readLen = blockSize ) @@ -62,18 +61,7 @@ func TestStorageDiff_RefreshesStaleHintFromAncestorHeader(t *testing.T) { aPaths := storage.Paths{BuildID: aID.String()} provider := storage.NewMockStorageProvider(t) - // Init opens the hint (uncompressed) path. The chunker's first range - // fetch surfaces ErrObjectNotExist, triggering the recovery hook. - uncompressedSeekable := storage.NewMockSeekable(t) - uncompressedSeekable.EXPECT().Size(mock.Anything).Return(int64(payloadSize), nil) - uncompressedSeekable.EXPECT(). - OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything). - Return(nil, storage.ErrObjectNotExist).Once() - provider.EXPECT(). - OpenSeekable(mock.Anything, aPaths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything). - Return(uncompressedSeekable, nil).Once() - - // Recovery fetches A's own header from the sidecar. + // Init loads A's own header — no OpenSeekable at the uncompressed path. headerBlob := storage.NewMockBlob(t) headerBlob.EXPECT(). WriteTo(mock.Anything, mock.Anything). @@ -84,7 +72,7 @@ func TestStorageDiff_RefreshesStaleHintFromAncestorHeader(t *testing.T) { OpenBlob(mock.Anything, aPaths.HeaderFile(storage.MemfileName), mock.Anything). Return(headerBlob, nil).Once() - // Reopen at the compressed path; the retried read decompresses cleanly. + // Then open upstream at the compressed path; the read decompresses cleanly. compressedSeekable := storage.NewMockSeekable(t) compressedSeekable.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything). @@ -96,9 +84,135 @@ func TestStorageDiff_RefreshesStaleHintFromAncestorHeader(t *testing.T) { require.Equal(t, payload[:readLen], runRead(t, bHeader, provider, readLen)) } +// When the proactive header load returns a header whose BuildId matches the +// File's own current header, getBuild promotes it via SwapHeader. Subsequent +// reads then pick up the loaded header's authoritative Builds map directly. +func TestStorageDiff_SwapsHeaderOnSelfMatch(t *testing.T) { + t.Parallel() + + const ( + blockSize = int64(4 << 10) + frameSizeKB = 256 + payloadSize = 256 * 1024 + readLen = blockSize + ) + + ctx := t.Context() + selfID := uuid.New() + payload := bytes.Repeat([]byte("self-payload-"), payloadSize/len("self-payload-")+1)[:payloadSize] + + frameTable, compressed, _, err := storage.CompressBytes(ctx, payload, storage.CompressConfig{ + Enabled: true, + Type: storage.CompressionZstd.String(), + Level: 2, + EncoderConcurrency: 1, + FrameEncodeWorkers: 1, + FrameSizeKB: frameSizeKB, + MinPartSizeMB: 50, + }) + require.NoError(t, err) + + // Sidecar header: V4, self-identified, with a populated Builds map. + fullHeader := buildHeader(t, selfID, payloadSize, selfID) + fullHeader.SetBuild(selfID, header.BuildData{Size: int64(payloadSize), FrameData: frameTable}) + fullHeaderBytes, err := header.SerializeHeader(fullHeader) + require.NoError(t, err) + + // File's current header: same BuildId, but Builds map is empty — forces + // the proactive header load when reading the self mapping. + staleHeader := buildHeader(t, selfID, payloadSize, selfID) + + paths := storage.Paths{BuildID: selfID.String()} + provider := storage.NewMockStorageProvider(t) + + headerBlob := storage.NewMockBlob(t) + headerBlob.EXPECT(). + WriteTo(mock.Anything, mock.Anything). + RunAndReturn(func(_ context.Context, w io.Writer) (int64, error) { + return io.Copy(w, bytes.NewReader(fullHeaderBytes)) + }).Once() + provider.EXPECT(). + OpenBlob(mock.Anything, paths.HeaderFile(storage.MemfileName), mock.Anything). + Return(headerBlob, nil).Once() + + compressedSeekable := storage.NewMockSeekable(t) + compressedSeekable.EXPECT(). + OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything). + RunAndReturn(decompressingRangeReader(compressed)) + provider.EXPECT(). + OpenSeekable(mock.Anything, paths.DataFile(storage.MemfileName, storage.CompressionZstd), mock.Anything). + Return(compressedSeekable, nil).Once() + + got, f := runReadOnFile(t, staleHeader, provider, readLen) + require.Equal(t, payload[:readLen], got) + require.NotSame(t, staleHeader, f.Header(), "File.header should have been swapped") + swapped := f.Header() + require.Equal(t, selfID, swapped.Metadata.BuildId, "swapped header should retain self BuildId") + require.Contains(t, swapped.Builds, selfID, "swapped header should carry the loaded Builds entry") + require.Empty(t, staleHeader.Builds, "stale header's empty Builds map proves the swap was needed") +} + +// A loaded header that matches self by BuildId but doesn't carry an +// authoritative self entry in its Builds map (mid-P2P upload, V3 with no +// Builds map, or other incomplete state) must NOT be promoted via SwapHeader +// — it would be a downgrade of File's current view. +func TestStorageDiff_DoesNotSwapWhenLoadedLacksSelfEntry(t *testing.T) { + t.Parallel() + + const ( + blockSize = int64(4 << 10) + payloadSize = 8 << 10 + readLen = blockSize + ) + + selfID := uuid.New() + payload := bytes.Repeat([]byte("incomplete-self-"), payloadSize/len("incomplete-self-")+1)[:payloadSize] + + // Loaded header: self BuildId matches, but Builds map is empty — + // e.g. mid-P2P state where the peer hasn't populated its self entry. + incompleteHeader := buildHeader(t, selfID, payloadSize, selfID) + incompleteBytes, err := header.SerializeHeader(incompleteHeader) + require.NoError(t, err) + + staleHeader := buildHeader(t, selfID, payloadSize, selfID) + paths := storage.Paths{BuildID: selfID.String()} + provider := storage.NewMockStorageProvider(t) + + headerBlob := storage.NewMockBlob(t) + headerBlob.EXPECT(). + WriteTo(mock.Anything, mock.Anything). + RunAndReturn(func(_ context.Context, w io.Writer) (int64, error) { + return io.Copy(w, bytes.NewReader(incompleteBytes)) + }).Once() + provider.EXPECT(). + OpenBlob(mock.Anything, paths.HeaderFile(storage.MemfileName), mock.Anything). + Return(headerBlob, nil).Once() + + // Loaded header has empty Builds → bd is zero → size=0, initCT=None → + // newStorageDiff opens uncompressed path and falls back to obj.Size(). + uncompressedSeekable := storage.NewMockSeekable(t) + uncompressedSeekable.EXPECT(). + Size(mock.Anything). + 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) { + end := min(off+length, int64(len(payload))) + + return io.NopCloser(bytes.NewReader(payload[off:end])), nil + }) + provider.EXPECT(). + OpenSeekable(mock.Anything, paths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything). + Return(uncompressedSeekable, nil).Once() + + got, f := runReadOnFile(t, staleHeader, provider, readLen) + require.Equal(t, payload[:readLen], got) + require.Same(t, staleHeader, f.Header(), "File.header must not be swapped to an incomplete loaded one") +} + // Finalized B + nil FrameData on A means "A is uncompressed and that's the -// authoritative answer." The read must serve uncompressed bytes and must not -// fetch A's header sidecar. +// authoritative answer." Init must take the parent's word for it and not +// fetch A's header. func TestStorageDiff_NoRefreshOnFinalizedHeader(t *testing.T) { t.Parallel() @@ -133,11 +247,75 @@ func TestStorageDiff_NoRefreshOnFinalizedHeader(t *testing.T) { require.Equal(t, payload[:readLen], runRead(t, bHeader, provider, readLen)) } +// When a peer is known to be P2P-serving the ancestor and our parent header +// lacks its Builds entry, create must NOT fetch the ancestor's header — +// that refresh would just return the same in-progress header over P2P. +// Instead, open at the uncompressed path (which the peer-routed seekable +// serves natively) and let per-read callerFT / peer-transition refresh +// handle CT/FT learning. +func TestStorageDiff_SkipsHeaderRefreshWhenPeerActive(t *testing.T) { + t.Parallel() + + const ( + blockSize = int64(4 << 10) + payloadSize = 8 << 10 + readLen = blockSize + ) + + aID := uuid.New() + payload := bytes.Repeat([]byte("p2p-served-"), payloadSize/len("p2p-served-")+1)[:payloadSize] + + // Parent has no Builds[aID] entry — would normally trigger the proactive + // refresh path. + bHeader := buildHeader(t, uuid.New(), payloadSize, aID) + + aPaths := storage.Paths{BuildID: aID.String()} + provider := storage.NewMockStorageProvider(t) + + // Critical: no EXPECT for OpenBlob. peerActiveBootstrapSize discriminates on + // the PeerRouted marker — we wrap the mock so it counts as peer-routed, + // otherwise the bootstrap would refresh and the test would panic. + uncompressedSeekable := storage.NewMockSeekable(t) + uncompressedSeekable.EXPECT(). + Size(mock.Anything). + 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) { + end := min(off+length, int64(len(payload))) + + return io.NopCloser(bytes.NewReader(payload[off:end])), nil + }) + provider.EXPECT(). + OpenSeekable(mock.Anything, aPaths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything). + Return(peerRoutedSeekable{Seekable: uncompressedSeekable}, nil).Once() + + got, _ := runReadWithPeers(t, bHeader, provider, readLen, alwaysActivePeer) + require.Equal(t, payload[:readLen], got) +} + // runRead wires up a File over a fresh DiffStore + noop metrics and returns // the first `n` bytes read from offset 0. func runRead(t *testing.T, h *header.Header, provider storage.StorageProvider, n int64) []byte { t.Helper() - store, err := NewDiffStore(cfg.Config{}, &featureflags.Client{}, t.TempDir(), time.Hour, time.Minute) + buf, _ := runReadOnFile(t, h, provider, n) + + return buf +} + +// runReadOnFile is runRead's variant that exposes the File so callers can +// inspect post-read state (e.g. SwapHeader fired). +func runReadOnFile(t *testing.T, h *header.Header, provider storage.StorageProvider, n int64) ([]byte, *File) { + t.Helper() + + return runReadWithPeers(t, h, provider, n, nil) +} + +// runReadWithPeers wires up a File with the given IsActivePeer hook so tests +// can drive the IsActive-true branch. +func runReadWithPeers(t *testing.T, h *header.Header, provider storage.StorageProvider, n int64, isActivePeer IsActivePeer) ([]byte, *File) { + t.Helper() + store, err := NewDiffStore(cfg.Config{}, &featureflags.Client{}, t.TempDir(), time.Hour, time.Minute, isActivePeer) require.NoError(t, err) m, err := blockmetrics.NewMetrics(noop.NewMeterProvider()) require.NoError(t, err) @@ -148,9 +326,22 @@ func runRead(t *testing.T, h *header.Header, provider storage.StorageProvider, n require.NoError(t, err) require.Equal(t, int(n), got) - return buf + return buf, f +} + +// alwaysActivePeer is an IsActivePeer that reports every build as P2P-active. +func alwaysActivePeer(string) bool { return true } + +// peerRoutedSeekable wraps a Seekable to satisfy the peerclient.PeerRouted +// marker — production peer routing wraps the base Seekable in peerSeekable, +// which advertises IsPeerRouted() == true. Tests that need to take the +// peer-routed branch in newStorageDiff use this shim. +type peerRoutedSeekable struct { + storage.Seekable } +func (peerRoutedSeekable) IsPeerRouted() {} + // testBlockSize is the block size used by buildHeader. Every header-shape // test in this file uses the same value; if a future test legitimately // needs a different one, hoist this back into a parameter. @@ -184,139 +375,3 @@ func decompressingRangeReader(compressed []byte) func(context.Context, int64, in return storage.NewDecompressingReader(bytes.NewReader(compressed[r.Offset:end]), ft.CompressionType()) } } - -// fakeRefresher is a stand-in for *StorageDiff in tests that exercise -// specialErrorHandler's decision logic without spinning up the chunker. -// loadedHeader is what RefreshFrameTable returns (nil simulates the -// idempotency-latch short-circuit; non-nil simulates a successful load). -type fakeRefresher struct { - ft *storage.FrameTable - loadedHeader *header.Header - refreshCalls int -} - -func (f *fakeRefresher) FrameTable() *storage.FrameTable { return f.ft } - -func (f *fakeRefresher) RefreshFrameTable(context.Context) (*header.Header, error) { - f.refreshCalls++ - - return f.loadedHeader, nil -} - -func TestFile_SpecialErrorDecisionMatrix(t *testing.T) { - t.Parallel() - - cases := []struct { - name string - err error - callerFT *storage.FrameTable - latchedFT *storage.FrameTable - wantRecovery bool - }{ - { - name: "ObjectNotExist, no caller hint, no latched FT triggers refresh", - err: storage.ErrObjectNotExist, - wantRecovery: true, - }, - { - name: "ObjectNotExist, no caller hint, latched FT still triggers refresh (C1)", - err: storage.ErrObjectNotExist, - latchedFT: &storage.FrameTable{}, - wantRecovery: true, - }, - { - name: "ObjectNotExist with caller-authoritative hint does not refresh", - err: storage.ErrObjectNotExist, - callerFT: &storage.FrameTable{}, - wantRecovery: false, - }, - { - name: "PeerTransitionedError triggers refresh after backoff", - err: &storage.PeerTransitionedError{}, - wantRecovery: true, - }, - { - name: "Unrelated errors do not trigger refresh", - err: errors.New("boom"), - wantRecovery: false, - }, - } - - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - selfID := uuid.New() - f := NewFile(buildHeader(t, selfID, 4096, selfID), nil, Memfile, nil, blockmetrics.Metrics{}) - refresher := &fakeRefresher{ft: tc.latchedFT} - seg := readSegment{ft: tc.callerFT} - - herr := f.specialErrorHandler(t.Context(), refresher, seg, tc.err) - if tc.wantRecovery { - require.NoError(t, herr, "expected recovery to succeed and signal retry") - require.Equal(t, 1, refresher.refreshCalls) - } else { - require.ErrorIs(t, herr, tc.err, "expected original error to propagate without refresh") - require.Zero(t, refresher.refreshCalls) - } - }) - } -} - -// TestFile_SpecialErrorHandlerPromotesSelfHeader covers the post-refresh -// SwapHeader hook. When the refreshed header's buildID matches the File's -// own, recovery promotes it; an ancestor header (different buildID) is -// dropped, leaving File.header untouched. -func TestFile_SpecialErrorHandlerPromotesSelfHeader(t *testing.T) { - t.Parallel() - - selfID := uuid.New() - ancestorID := uuid.New() - - cases := []struct { - name string - loadedBuildID uuid.UUID - wantSwapped bool - }{ - { - name: "self uuid match promotes loaded header", - loadedBuildID: selfID, - wantSwapped: true, - }, - { - name: "ancestor uuid leaves File.header untouched", - loadedBuildID: ancestorID, - wantSwapped: false, - }, - { - name: "idempotency-latch nil header is a no-op", - loadedBuildID: uuid.Nil, // sentinel meaning "fakeRefresher returns nil" - wantSwapped: false, - }, - } - - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - initial := buildHeader(t, selfID, 4096, selfID) - f := NewFile(initial, nil, Memfile, nil, blockmetrics.Metrics{}) - - var loaded *header.Header - if tc.loadedBuildID != uuid.Nil { - loaded = buildHeader(t, tc.loadedBuildID, 4096, tc.loadedBuildID) - } - refresher := &fakeRefresher{loadedHeader: loaded} - seg := readSegment{} - - require.NoError(t, f.specialErrorHandler(t.Context(), refresher, seg, storage.ErrObjectNotExist)) - require.Equal(t, 1, refresher.refreshCalls) - - if tc.wantSwapped { - require.Same(t, loaded, f.Header(), "expected File.header to be swapped to the loaded header") - } else { - require.Same(t, initial, f.Header(), "expected File.header to remain the initial one") - } - }) - } -} diff --git a/packages/orchestrator/pkg/sandbox/nbd/testutils/template_rootfs.go b/packages/orchestrator/pkg/sandbox/nbd/testutils/template_rootfs.go index d11e45dc98..97debb9270 100644 --- a/packages/orchestrator/pkg/sandbox/nbd/testutils/template_rootfs.go +++ b/packages/orchestrator/pkg/sandbox/nbd/testutils/template_rootfs.go @@ -89,6 +89,7 @@ func TemplateRootfs(ctx context.Context, buildID string) (*BuildDevice, *Cleaner diffCacheDir, 24*time.Hour, 24*time.Hour, + nil, ) if err != nil { return nil, &cleaner, fmt.Errorf("failed to create diff store: %w", err) diff --git a/packages/orchestrator/pkg/sandbox/template/cache.go b/packages/orchestrator/pkg/sandbox/template/cache.go index 2296542fb1..728686aad5 100644 --- a/packages/orchestrator/pkg/sandbox/template/cache.go +++ b/packages/orchestrator/pkg/sandbox/template/cache.go @@ -97,6 +97,7 @@ func NewCache( config.DefaultCacheDir, buildCacheTTL, buildCacheDelayEviction, + peers.IsActive, ) if err != nil { return nil, fmt.Errorf("failed to create build store: %w", err) diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go index 4b9230da73..144c627a7b 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go @@ -7,7 +7,6 @@ import ( "io" "sync" "sync/atomic" - "time" "go.uber.org/zap" @@ -16,11 +15,6 @@ import ( "github.com/e2b-dev/infra/packages/shared/pkg/storage" ) -const ( - postTransitionRetryWindow = 30 * time.Second - postTransitionRetryDelay = 250 * time.Millisecond -) - var _ storage.Seekable = (*peerSeekable)(nil) // peerSeekable reads from the peer orchestrator first. @@ -38,12 +32,6 @@ type peerSeekable struct { base storage.Seekable baseCT storage.CompressionType loaded bool - - // transitionAt is set on first PeerTransitionedError emission after - // uploaded flips to true; nil otherwise. Subsequent base 404s within - // postTransitionRetryWindow re-emit PeerTransitionedError so concurrent - // readers retry against the post-upload header. - transitionAt atomic.Pointer[time.Time] } // getBase returns a base Seekable opened against the storage path composed @@ -71,6 +59,15 @@ func (s *peerSeekable) getBase(ctx context.Context, ct storage.CompressionType) return base, nil } +// Post-tryPeer fall-through rule shared by Size and OpenRangeReader: if +// uploaded has flipped (peer signaled UseStorage), return PeerTransitionedError +// so the caller refreshes + reopens. We never serve base from a post-transition +// peerSeekable — this wrapper's base path was captured pre-finalization (basic +// name); the actual GCS object lives at the CT-qualified path the refreshed +// header reveals. Routing back through the resolver after refresh returns base +// directly (attrResolveUploaded), so the retry hits GCS at the right path with +// no wrapper involved. + func (s *peerSeekable) Size(ctx context.Context) (int64, error) { res, err := tryPeer(ctx, &s.peerHandle, "size peer-seekable", attrOpSize, func(ctx context.Context) (peerAttempt[int64], error) { @@ -91,10 +88,10 @@ func (s *peerSeekable) Size(ctx context.Context) (int64, error) { if res.hit { return res.value, err } + if s.uploaded.Load() { + return 0, &storage.PeerTransitionedError{} + } - // Size only reaches base for V3 builds (uncompressedSize unknown); - // V4 builds carry the size in the header so the chunker never calls Size. - // V3 implies CompressionNone, matching reality. base, err := s.getBase(ctx, storage.CompressionNone) if err != nil { return 0, err @@ -129,12 +126,8 @@ func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length in if res.hit { return res.value, err } - - if s.uploaded != nil && s.uploaded.Load() { - now := time.Now() - if s.transitionAt.CompareAndSwap(nil, &now) { - return nil, &storage.PeerTransitionedError{} - } + if s.uploaded.Load() { + return nil, &storage.PeerTransitionedError{} } base, err := s.getBase(ctx, frameTable.CompressionType()) @@ -142,16 +135,7 @@ func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length in return nil, err } - rc, err := base.OpenRangeReader(ctx, off, length, frameTable) - // GCS can briefly 404 a just-finalized object; within the retry window - // re-emit so build.File reloads the header and retries with backoff. - if errors.Is(err, storage.ErrObjectNotExist) { - if at := s.transitionAt.Load(); at != nil && time.Since(*at) < postTransitionRetryWindow { - return nil, &storage.PeerTransitionedError{RetryAfter: postTransitionRetryDelay} - } - } - - return rc, err + return base.OpenRangeReader(ctx, off, length, frameTable) } func (s *peerSeekable) StoreFile(context.Context, string, ...storage.PutOption) (*storage.FrameTable, [32]byte, error) { diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go index b5266e9094..bf7e5288bb 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go @@ -143,31 +143,20 @@ func TestPeerSeekable_OpenRangeReader_Uploaded_ReturnsPeerTransitionedError(t *t require.ErrorAs(t, err, &transErr) } -// TestPeerStorageProvider_FullTransitionFlow walks the whole peerclient -// surface across a peer→storage transition with a header swap from V3 (basic -// path) to V4 (zstd-compressed path). Regression cover for the bug where the -// post-transition read kept hitting the original uncompressed path. -// -// Sequence: -// 1. Pre-transition: caller passes ft={ct=None}; peer answers; bytes flow. -// 2. Peer signals UseStorage; uploaded flips to true. -// 3. First post-transition call: peerSeekable returns PeerTransitionedError -// immediately (no peer call, no base open). -// 4. Caller (build.File.retryOnTransition, simulated here) reloads the V4 -// header and retries with ft={ct=Zstd}. -// 5. peerSeekable falls through to base, which opens "build-1/memfile.zstd" -// (not "build-1/memfile") and serves the compressed bytes. -func TestPeerStorageProvider_FullTransitionFlow(t *testing.T) { +// TestPeerStorageProvider_TransitionEmitsError covers the peer→storage +// transition contract: while uploaded is false, peerSeekable serves from the +// peer; the call that observes UseStorage flips uploaded and returns its +// bytes; the NEXT call on the same wrapper returns PeerTransitionedError +// without touching peer or base. The wrapper never falls through to base +// post-transition — that's the resolver's job (attrResolveUploaded) once the +// caller reopens. Catching a peerSeekable falling through here would mean we +// regressed to the 404-driven recovery design. +func TestPeerStorageProvider_TransitionEmitsError(t *testing.T) { t.Parallel() uploaded := &atomic.Bool{} - prePeerBytes := []byte("pre-transition peer payload") - postBaseBytes := []byte("post-transition compressed payload") - // Pre-transition peer stream: serves bytes once, then EOF. uploaded is - // flipped via UseStorage on the EOF response so subsequent calls skip - // the peer. preStream := orchestratormocks.NewMockChunkService_ReadAtBuildSeekableClient(t) preStream.EXPECT().Recv().Return(&orchestrator.ReadAtBuildSeekableResponse{Data: prePeerBytes}, nil).Once() preStream.EXPECT().Recv().RunAndReturn(func() (*orchestrator.ReadAtBuildSeekableResponse, error) { @@ -178,27 +167,18 @@ func TestPeerStorageProvider_FullTransitionFlow(t *testing.T) { client := orchestratormocks.NewMockChunkServiceClient(t) client.EXPECT().ReadAtBuildSeekable(mock.Anything, mock.MatchedBy(func(req *orchestrator.ReadAtBuildSeekableRequest) bool { - // Peer is asked by basic name only. return req.GetBuildId() == "build-1" && req.GetName() == storage.MemfileName })).Return(preStream, nil).Once() - // Base is only consulted post-transition, and only against the compressed - // path. If the bug regresses (uncompressed path), this expectation fails. - postBaseSeekable := storage.NewMockSeekable(t) - postBaseSeekable.EXPECT(). - OpenRangeReader(mock.Anything, int64(0), int64(len(postBaseBytes)), mock.Anything). - Return(io.NopCloser(bytes.NewReader(postBaseBytes)), nil).Once() - + // base must NOT be touched: any post-transition fall-through would be a + // regression to 404-driven recovery. NewMockStorageProvider(t) auto-asserts + // no unexpected calls on cleanup. base := storage.NewMockStorageProvider(t) - base.EXPECT(). - OpenSeekable(mock.Anything, "build-1/memfile.zstd", storage.MemfileObjectType). - Return(postBaseSeekable, nil).Once() p := newPeerStorageProvider(base, client, uploaded) seekable, err := p.OpenSeekable(t.Context(), "build-1/memfile", storage.MemfileObjectType) require.NoError(t, err) - // 1. Pre-transition read via peer. ft={ct=None} (V3 header). rc, err := seekable.OpenRangeReader(t.Context(), 0, int64(len(prePeerBytes)), storage.NewFrameTable(storage.CompressionNone, nil)) require.NoError(t, err) @@ -208,19 +188,8 @@ func TestPeerStorageProvider_FullTransitionFlow(t *testing.T) { assert.Equal(t, prePeerBytes, got) require.True(t, uploaded.Load(), "uploaded flag should be set after peer EOF with UseStorage") - // 2. First post-transition call: retriable error, no peer/base contact. _, err = seekable.OpenRangeReader(t.Context(), 0, 1, storage.NewFrameTable(storage.CompressionNone, nil)) var transErr *storage.PeerTransitionedError require.ErrorAs(t, err, &transErr) - - // 3. Caller reloads V4 header and retries with ct=Zstd. This must hit the - // compressed path on base. - rc, err = seekable.OpenRangeReader(t.Context(), 0, int64(len(postBaseBytes)), - storage.NewFrameTable(storage.CompressionZstd, nil)) - require.NoError(t, err) - got, err = io.ReadAll(rc) - require.NoError(t, err) - require.NoError(t, rc.Close()) - assert.Equal(t, postBaseBytes, got) } diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/storage.go b/packages/orchestrator/pkg/sandbox/template/peerclient/storage.go index 9e7b7b6883..eb463e53d0 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/storage.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/storage.go @@ -46,6 +46,19 @@ var ( attrPeerHitFalse = attribute.Bool("peer_hit", false) ) +// PeerRouted marks a Seekable that resolveProvider actually routed through a +// peer at open time. Callers that need to distinguish "the routing provider +// gave me a peer wrapper" from "the routing provider fell through to base" +// type-assert against this marker; presence is the signal — the method body +// is intentionally empty. +type PeerRouted interface { + IsPeerRouted() +} + +func (*peerSeekable) IsPeerRouted() {} + +var _ PeerRouted = (*peerSeekable)(nil) + var _ storage.StorageProvider = (*routingProvider)(nil) // routingProvider wraps a base StorageProvider and, for each Open call, diff --git a/packages/orchestrator/pkg/sandbox/template/storage.go b/packages/orchestrator/pkg/sandbox/template/storage.go index 2b2500ef7d..3fdbcf4d20 100644 --- a/packages/orchestrator/pkg/sandbox/template/storage.go +++ b/packages/orchestrator/pkg/sandbox/template/storage.go @@ -58,7 +58,7 @@ func NewStorage( } var err error - h, err = header.LoadHeader(ctx, persistence, hdrPath) + h, _, err = header.LoadHeader(ctx, persistence, hdrPath) if err != nil && !errors.Is(err, storage.ErrObjectNotExist) { return nil, err } diff --git a/packages/shared/pkg/storage/compress_frame_table.go b/packages/shared/pkg/storage/compress_frame_table.go index 9b8079ae7c..9dc9c73de3 100644 --- a/packages/shared/pkg/storage/compress_frame_table.go +++ b/packages/shared/pkg/storage/compress_frame_table.go @@ -70,6 +70,12 @@ type FrameTable struct { entries []frameEntry // sorted by StartU } +// UncompressedFrameTable is the canonical sentinel for "data is stored +// uncompressed" — an empty FrameTable has no frames to decompress, so a reader +// handed this back skips the U→C translation entirely. Shared across callers +// to avoid per-call allocation; read-only and must never be mutated. +var UncompressedFrameTable = &FrameTable{} + // newFrameTableFromEntries creates a FrameTable from pre-computed absolute-offset entries. func newFrameTableFromEntries(ct CompressionType, entries []frameEntry) *FrameTable { return &FrameTable{compressionType: ct, entries: entries} diff --git a/packages/shared/pkg/storage/header/header.go b/packages/shared/pkg/storage/header/header.go index c50969c317..2b866482fd 100644 --- a/packages/shared/pkg/storage/header/header.go +++ b/packages/shared/pkg/storage/header/header.go @@ -179,14 +179,15 @@ func (t *Header) GetShiftedMapping(ctx context.Context, offset int64) (BuildMap, } // GetBuildFrameData returns the FrameTable for a build: nil = no entry, -// empty &FrameTable{} = authoritatively uncompressed, else compressed. +// storage.UncompressedFrameTable = authoritatively uncompressed, else +// compressed. func (t *Header) GetBuildFrameData(buildID uuid.UUID) *storage.FrameTable { bd, ok := t.Builds[buildID] if !ok { return nil } if bd.FrameData == nil { - return &storage.FrameTable{} + return storage.UncompressedFrameTable } return bd.FrameData diff --git a/packages/shared/pkg/storage/header/serialization.go b/packages/shared/pkg/storage/header/serialization.go index 115acec85d..9f481b7c1b 100644 --- a/packages/shared/pkg/storage/header/serialization.go +++ b/packages/shared/pkg/storage/header/serialization.go @@ -62,19 +62,26 @@ func DeserializeBytes(data []byte) (*Header, error) { } // LoadHeader fetches a serialized header from storage and deserializes it. -// Errors (including storage.ErrObjectNotExist) are returned as-is. -func LoadHeader(ctx context.Context, s storage.StorageProvider, path string) (*Header, error) { +// Returns the on-wire byte count alongside the header so callers can attribute +// it to throughput telemetry. Errors (including storage.ErrObjectNotExist) are +// returned as-is. +func LoadHeader(ctx context.Context, s storage.StorageProvider, path string) (*Header, int, error) { blob, err := s.OpenBlob(ctx, path, storage.MetadataObjectType) if err != nil { - return nil, fmt.Errorf("open blob %s: %w", path, err) + return nil, 0, fmt.Errorf("open blob %s: %w", path, err) } data, err := storage.GetBlob(ctx, blob) if err != nil { - return nil, err + return nil, 0, err } - return DeserializeBytes(data) + h, err := DeserializeBytes(data) + if err != nil { + return nil, len(data), err + } + + return h, len(data), nil } // StoreHeader serializes a header, uploads it, and returns the effective diff --git a/packages/shared/pkg/storage/header/serialization_test.go b/packages/shared/pkg/storage/header/serialization_test.go index db664a3050..6a2fad2305 100644 --- a/packages/shared/pkg/storage/header/serialization_test.go +++ b/packages/shared/pkg/storage/header/serialization_test.go @@ -839,7 +839,7 @@ func TestSerializeDeserialize_V4_Uncompressed_SelfEntry(t *testing.T) { require.Equal(t, uint64(MetadataVersionV4), got.Metadata.Version) require.Len(t, got.Builds, 1) require.Contains(t, got.Builds, buildID) - require.Equal(t, &storage.FrameTable{}, got.GetBuildFrameData(buildID)) + require.Equal(t, storage.UncompressedFrameTable, got.GetBuildFrameData(buildID)) } // Layered chain V4-uncompressed (self) → V4-compressed (mid) → V4-uncompressed @@ -893,8 +893,8 @@ func TestSerializeDeserialize_V4_MixedChain(t *testing.T) { require.Equal(t, 5, got.Mapping.Len()) require.Len(t, got.Builds, 3) - require.Equal(t, &storage.FrameTable{}, got.GetBuildFrameData(selfID)) - require.Equal(t, &storage.FrameTable{}, got.GetBuildFrameData(olderID)) + require.Equal(t, storage.UncompressedFrameTable, got.GetBuildFrameData(selfID)) + require.Equal(t, storage.UncompressedFrameTable, got.GetBuildFrameData(olderID)) gotMidFT := got.GetBuildFrameData(midID) require.NotNil(t, gotMidFT) From 4d2309386f729af003dfa9f18f24b982139515b1 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Mon, 8 Jun 2026 07:04:32 -0700 Subject: [PATCH 3/6] fixed integ test --- .../tests/api/sandboxes/sandbox_rapid_pause_resume_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/internal/tests/api/sandboxes/sandbox_rapid_pause_resume_test.go b/tests/integration/internal/tests/api/sandboxes/sandbox_rapid_pause_resume_test.go index 0cdbd75eb9..474477dcd3 100644 --- a/tests/integration/internal/tests/api/sandboxes/sandbox_rapid_pause_resume_test.go +++ b/tests/integration/internal/tests/api/sandboxes/sandbox_rapid_pause_resume_test.go @@ -170,7 +170,7 @@ func loadHeaderWithPolling(t *testing.T, ctx context.Context, persistence storag var h *header.Header require.Eventually(t, func() bool { var err error - h, err = header.LoadHeader(ctx, persistence, path) + h, _, err = header.LoadHeader(ctx, persistence, path) return err == nil && h != nil }, 2*time.Minute, 500*time.Millisecond, "%s/%s: %s never appeared in storage", name, fileLabel, path) From 64b6f64cc4f141487f067768925650be1b6c0c97 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Mon, 8 Jun 2026 10:30:57 -0700 Subject: [PATCH 4/6] PR feedback: fixes --- .../orchestrator/pkg/sandbox/build/build.go | 12 +++++- .../orchestrator/pkg/sandbox/build/cache.go | 3 +- .../pkg/sandbox/build/storage_diff_test.go | 37 ++++++++----------- .../sandbox/template/peerclient/seekable.go | 16 ++++---- .../template/peerclient/seekable_test.go | 17 +++++---- 5 files changed, 43 insertions(+), 42 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/build/build.go b/packages/orchestrator/pkg/sandbox/build/build.go index c9fc8c6387..2aa2d0da35 100644 --- a/packages/orchestrator/pkg/sandbox/build/build.go +++ b/packages/orchestrator/pkg/sandbox/build/build.go @@ -459,7 +459,17 @@ func (b *File) refreshAncestorAndOpenUpstream(ctx context.Context, buildID uuid. } } - bd := loaded.Builds[buildID] + bd, hasSelf := loaded.Builds[buildID] + if !hasSelf { + // A finalized V4+ storage header always carries a self entry + // (build_upload_v4 populates it before publish). Reaching here means a + // routed OpenBlob hit a peer's in-flight header — which shouldn't be + // possible on this code path (we entered after !peerActive). Surface + // loudly rather than silently latching the zero-value bd as an + // authoritative uncompressed FT. + return nil, 0, nil, fmt.Errorf("createDiff: loaded header for build %s has no self entry (peer-served incomplete?)", buildID) + } + ft := bd.FrameData if ft == nil { ft = storage.UncompressedFrameTable diff --git a/packages/orchestrator/pkg/sandbox/build/cache.go b/packages/orchestrator/pkg/sandbox/build/cache.go index 0490e120ca..b30bd06ec6 100644 --- a/packages/orchestrator/pkg/sandbox/build/cache.go +++ b/packages/orchestrator/pkg/sandbox/build/cache.go @@ -122,8 +122,7 @@ func (s *DiffStore) Close() { } // Get returns the cached Diff for key, refreshing TTL and cancelling any -// pending eviction. Returns (nil, false) if the key isn't present. Use -// Lookup for a side-effect-free probe (e.g. from IsCached). +// pending eviction. Returns (nil, false) if the key isn't present. func (s *DiffStore) Get(key DiffStoreKey) (Diff, bool) { s.resetDelete(key) item := s.cache.Get(key) diff --git a/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go b/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go index fb97e2dd62..215188dcce 100644 --- a/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go @@ -153,10 +153,11 @@ func TestStorageDiff_SwapsHeaderOnSelfMatch(t *testing.T) { } // A loaded header that matches self by BuildId but doesn't carry an -// authoritative self entry in its Builds map (mid-P2P upload, V3 with no -// Builds map, or other incomplete state) must NOT be promoted via SwapHeader -// — it would be a downgrade of File's current view. -func TestStorageDiff_DoesNotSwapWhenLoadedLacksSelfEntry(t *testing.T) { +// authoritative self entry (mid-P2P upload, V3 with no Builds map, or other +// incomplete state) must NOT be promoted via SwapHeader, AND createDiff must +// fail loudly rather than silently latching the zero-value bd as +// authoritatively uncompressed. +func TestStorageDiff_RejectsLoadedHeaderLackingSelfEntry(t *testing.T) { t.Parallel() const ( @@ -166,7 +167,6 @@ func TestStorageDiff_DoesNotSwapWhenLoadedLacksSelfEntry(t *testing.T) { ) selfID := uuid.New() - payload := bytes.Repeat([]byte("incomplete-self-"), payloadSize/len("incomplete-self-")+1)[:payloadSize] // Loaded header: self BuildId matches, but Builds map is empty — // e.g. mid-P2P state where the peer hasn't populated its self entry. @@ -188,25 +188,18 @@ func TestStorageDiff_DoesNotSwapWhenLoadedLacksSelfEntry(t *testing.T) { OpenBlob(mock.Anything, paths.HeaderFile(storage.MemfileName), mock.Anything). Return(headerBlob, nil).Once() - // Loaded header has empty Builds → bd is zero → size=0, initCT=None → - // newStorageDiff opens uncompressed path and falls back to obj.Size(). - uncompressedSeekable := storage.NewMockSeekable(t) - uncompressedSeekable.EXPECT(). - Size(mock.Anything). - 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) { - end := min(off+length, int64(len(payload))) + // No OpenSeekable / Size / OpenRangeReader expectations: refresh must error + // out before any upstream is opened. If any of these fire the mock fails. - return io.NopCloser(bytes.NewReader(payload[off:end])), nil - }) - provider.EXPECT(). - OpenSeekable(mock.Anything, paths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything). - Return(uncompressedSeekable, nil).Once() + store, err := NewDiffStore(cfg.Config{}, &featureflags.Client{}, t.TempDir(), time.Hour, time.Minute, nil) + require.NoError(t, err) + m, err := blockmetrics.NewMetrics(noop.NewMeterProvider()) + require.NoError(t, err) + f := NewFile(staleHeader, store, Memfile, provider, m) - got, f := runReadOnFile(t, staleHeader, provider, readLen) - require.Equal(t, payload[:readLen], got) + _, err = f.ReadAt(t.Context(), make([]byte, readLen), 0) + require.Error(t, err) + require.Contains(t, err.Error(), "no self entry") require.Same(t, staleHeader, f.Header(), "File.header must not be swapped to an incomplete loaded one") } diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go index 144c627a7b..665b3cda9f 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go @@ -88,16 +88,14 @@ func (s *peerSeekable) Size(ctx context.Context) (int64, error) { if res.hit { return res.value, err } - if s.uploaded.Load() { - return 0, &storage.PeerTransitionedError{} - } - - base, err := s.getBase(ctx, storage.CompressionNone) - if err != nil { - return 0, err - } - return base.Size(ctx) + // Size has no caller-provided frame table to source the compression type + // from, and the basic-name fall-through would 404 on compressed V4 builds + // (data lives at .zstd). Surface PeerTransitionedError unconditionally on + // miss so the caller refreshes against the authoritative header — which + // knows the compression type — and either recovers or surfaces a clean "not + // yet on storage" error. + return 0, &storage.PeerTransitionedError{} } func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *storage.FrameTable) (io.ReadCloser, error) { diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go index bf7e5288bb..61b967aa3a 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go @@ -30,17 +30,18 @@ func TestPeerSeekable_Size_PeerSucceeds(t *testing.T) { assert.Equal(t, int64(4096), size) } -func TestPeerSeekable_Size_PeerNotAvailable_FallsBackToBase(t *testing.T) { +// Size has no caller-provided FT and cannot guess the CT, so a peer miss must +// emit PeerTransitionedError so the caller refreshes against the authoritative +// header rather than 404ing against a basic-name fall-through on compressed +// builds. base must NOT be touched. +func TestPeerSeekable_Size_PeerNotAvailable_EmitsPeerTransitionedError(t *testing.T) { t.Parallel() client := orchestratormocks.NewMockChunkServiceClient(t) client.EXPECT().GetBuildFileSize(mock.Anything, mock.Anything).Return(&orchestrator.GetBuildFileSizeResponse{Availability: &orchestrator.PeerAvailability{NotAvailable: true}}, nil) - baseSeekable := storage.NewMockSeekable(t) - baseSeekable.EXPECT().Size(mock.Anything).Return(int64(8192), nil) - + // NewMockStorageProvider auto-fails on any unexpected call. base := storage.NewMockStorageProvider(t) - base.EXPECT().OpenSeekable(mock.Anything, "build-1/memfile", storage.MemfileObjectType).Return(baseSeekable, nil) s := &peerSeekable{ peerHandle: peerHandle{ @@ -52,9 +53,9 @@ func TestPeerSeekable_Size_PeerNotAvailable_FallsBackToBase(t *testing.T) { basePersistence: base, objType: storage.MemfileObjectType, } - size, err := s.Size(t.Context()) - require.NoError(t, err) - assert.Equal(t, int64(8192), size) + _, err := s.Size(t.Context()) + var transErr *storage.PeerTransitionedError + require.ErrorAs(t, err, &transErr) } func TestPeerSeekable_OpenRangeReader_PeerSucceeds(t *testing.T) { From 2365947aa680f8c2ce2617be683a1d7abece7958 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Tue, 9 Jun 2026 14:31:06 -0700 Subject: [PATCH 5/6] self-review --- .../orchestrator/pkg/sandbox/build/build.go | 32 +++++------ .../pkg/sandbox/build/storage_diff.go | 55 +++++-------------- packages/orchestrator/pkg/server/chunks.go | 8 +-- 3 files changed, 32 insertions(+), 63 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/build/build.go b/packages/orchestrator/pkg/sandbox/build/build.go index 2aa2d0da35..55816b6850 100644 --- a/packages/orchestrator/pkg/sandbox/build/build.go +++ b/packages/orchestrator/pkg/sandbox/build/build.go @@ -156,8 +156,8 @@ func (b *File) readSegment(ctx context.Context, p []byte, s readSegment) error { if !errors.As(err, &transitionErr) { return err } - if waitErr := waitTransitionBackoff(ctx, b.fileType, transitionErr); waitErr != nil { - return waitErr + if err = waitTransitionBackoff(ctx, transitionErr); err != nil { + return err } if refreshErr := s.diff.RefreshSource(ctx); refreshErr != nil { return fmt.Errorf("refresh after peer transition: %w", refreshErr) @@ -316,7 +316,7 @@ func (b *File) IsCached(ctx context.Context, off, length int64) bool { // waitTransitionBackoff honors the peer's RetryAfter hint before the caller // retries against base storage. Returns ctx.Err() if cancelled during sleep. -func waitTransitionBackoff(ctx context.Context, fileType DiffType, transErr *storage.PeerTransitionedError) error { +func waitTransitionBackoff(ctx context.Context, transErr *storage.PeerTransitionedError) error { if transErr.RetryAfter > 0 { timer := time.NewTimer(transErr.RetryAfter) defer timer.Stop() @@ -326,7 +326,6 @@ func waitTransitionBackoff(ctx context.Context, fileType DiffType, transErr *sto return ctx.Err() } } - logger.L().Info(ctx, "peer transition detected, recovering", zap.String("file_type", string(fileType))) return nil } @@ -459,20 +458,15 @@ func (b *File) refreshAncestorAndOpenUpstream(ctx context.Context, buildID uuid. } } - bd, hasSelf := loaded.Builds[buildID] - if !hasSelf { - // A finalized V4+ storage header always carries a self entry - // (build_upload_v4 populates it before publish). Reaching here means a - // routed OpenBlob hit a peer's in-flight header — which shouldn't be - // possible on this code path (we entered after !peerActive). Surface - // loudly rather than silently latching the zero-value bd as an - // authoritative uncompressed FT. - return nil, 0, nil, fmt.Errorf("createDiff: loaded header for build %s has no self entry (peer-served incomplete?)", buildID) - } - - ft := bd.FrameData - if ft == nil { - ft = storage.UncompressedFrameTable + // A finalized V4+ storage header always carries a self entry + // (build_upload_v4 populates it before publish). A missing self entry here + // means a routed OpenBlob hit a peer's in-flight header — which shouldn't + // be possible on this code path (we entered after !peerActive). Surface + // loudly rather than silently latching a zero-value bd as an authoritative + // uncompressed FT. + size, ft, err := extractSelfBuildData(loaded, buildID) + if err != nil { + return nil, 0, nil, fmt.Errorf("createDiff: %w", err) } upstream, err := b.openUpstream(ctx, buildID, objType, ft.CompressionType()) @@ -480,7 +474,7 @@ func (b *File) refreshAncestorAndOpenUpstream(ctx context.Context, buildID uuid. return nil, 0, nil, err } - return upstream, bd.Size, ft, nil + return upstream, size, ft, nil } // initialSize is THE only production code path that calls Size on diff --git a/packages/orchestrator/pkg/sandbox/build/storage_diff.go b/packages/orchestrator/pkg/sandbox/build/storage_diff.go index 8212484a33..88165f3ce1 100644 --- a/packages/orchestrator/pkg/sandbox/build/storage_diff.go +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff.go @@ -119,10 +119,6 @@ func (b *StorageDiff) CacheKey() DiffStoreKey { } func (b *StorageDiff) Close() error { - if b.chunker == nil { - return nil - } - return b.chunker.Close() } @@ -150,18 +146,10 @@ func (b *StorageDiff) CachePath(context.Context) (string, error) { } func (b *StorageDiff) FileSize(ctx context.Context) (int64, error) { - if b.chunker == nil { - return 0, nil - } - return b.chunker.FileSize(ctx) } func (b *StorageDiff) Size(_ context.Context) (int64, error) { - if b.chunker == nil { - return 0, nil - } - return b.chunker.Size(), nil } @@ -170,13 +158,8 @@ func (b *StorageDiff) BlockSize() int64 { } // IsCached reports whether [off, off+length) is in the chunker's local cache. -// Returns false if the chunker hasn't been Init'd yet (would otherwise trigger -// OpenSeekable). Side-effect-free. +// Side-effect-free. func (b *StorageDiff) IsCached(ctx context.Context, off, length int64) bool { - if b.chunker == nil { - return false - } - return b.chunker.IsCached(ctx, off, length) } @@ -251,31 +234,19 @@ func (b *StorageDiff) reloadSource(ctx context.Context, cause string) error { return b.reloadSourceLocked(ctx, cause) } -// loadAuthoritativeHeader fetches the build's storage header and extracts the -// canonical (size, ft). Returns an error if the header lacks a self entry — -// that only happens with peer-served incomplete headers, never with -// storage-uploaded ones (build_upload_v4 always populates self before publish). -// Shared by reloadSourceLocked (read-time refresh) and createDiff -// (construction-time recovery when the peer-active bootstrap finds itself -// routed to base). -func loadAuthoritativeHeader( - ctx context.Context, - persistence storage.StorageProvider, - buildID uuid.UUID, - diffType DiffType, - cause string, -) (int64, *storage.FrameTable, error) { - loaded, err := refreshBuildHeader(ctx, persistence, buildID, diffType, cause) - if err != nil { - return 0, nil, err - } +// extractSelfBuildData returns (size, ft) from loaded.Builds[buildID]. Errors +// if the header lacks a self entry — that only happens with peer-served +// incomplete headers, never with storage-uploaded ones (build_upload_v4 always +// populates self before publish). A nil FrameData normalizes to +// UncompressedFrameTable (V4+ authoritatively uncompressed). +func extractSelfBuildData(loaded *header.Header, buildID uuid.UUID) (int64, *storage.FrameTable, error) { bd, hasSelf := loaded.Builds[buildID] if !hasSelf { - return 0, nil, fmt.Errorf("authoritative header for build %s has no self entry (peer-served incomplete?)", buildID) + return 0, nil, fmt.Errorf("header for build %s has no self entry (peer-served incomplete?)", buildID) } ft := bd.FrameData if ft == nil { - ft = storage.UncompressedFrameTable // V4+ authoritatively uncompressed + ft = storage.UncompressedFrameTable } return bd.Size, ft, nil @@ -300,9 +271,13 @@ func (b *StorageDiff) reloadSourceLocked(ctx context.Context, cause string) erro if err != nil { return fmt.Errorf("parse build id %s: %w", b.buildID, err) } - _, ft, err := loadAuthoritativeHeader(ctx, b.persistence, bid, b.diffType, cause) + loaded, err := refreshBuildHeader(ctx, b.persistence, bid, b.diffType, cause) + if err != nil { + return fmt.Errorf("reloadSourceLocked: load header for build %s (cause=%s): %w", b.buildID, cause, err) + } + _, ft, err := extractSelfBuildData(loaded, bid) if err != nil { - return fmt.Errorf("reloadSourceLocked: load authoritative header for build %s (cause=%s): %w", b.buildID, cause, err) + return fmt.Errorf("reloadSourceLocked: build %s (cause=%s): %w", b.buildID, cause, err) } newPath := storage.Paths{BuildID: b.buildID}.DataFile(string(b.diffType), ft.CompressionType()) newObj, err := b.persistence.OpenSeekable(ctx, newPath, b.storageObjectType) diff --git a/packages/orchestrator/pkg/server/chunks.go b/packages/orchestrator/pkg/server/chunks.go index 77c963e14a..0e7c3a2085 100644 --- a/packages/orchestrator/pkg/server/chunks.go +++ b/packages/orchestrator/pkg/server/chunks.go @@ -52,7 +52,7 @@ func toGRPCError(err error) error { func (s *Server) GetBuildFileSize(ctx context.Context, req *orchestrator.GetBuildFileSizeRequest) (*orchestrator.GetBuildFileSizeResponse, error) { telemetry.SetAttributes(ctx, telemetry.WithBuildID(req.GetBuildId()), attribute.String("file_name", req.GetName())) - if s.uploadedBuilds.Has(req.GetBuildId()) { + if s.uploadedBuilds.Get(req.GetBuildId()) != nil { telemetry.SetAttributes(ctx, attribute.Bool("uploaded", true)) return &orchestrator.GetBuildFileSizeResponse{Availability: peerUseStorage}, nil @@ -80,7 +80,7 @@ func (s *Server) GetBuildFileSize(ctx context.Context, req *orchestrator.GetBuil func (s *Server) GetBuildFileExists(ctx context.Context, req *orchestrator.GetBuildFileExistsRequest) (*orchestrator.GetBuildFileExistsResponse, error) { telemetry.SetAttributes(ctx, telemetry.WithBuildID(req.GetBuildId()), attribute.String("file_name", req.GetName())) - if s.uploadedBuilds.Has(req.GetBuildId()) { + if s.uploadedBuilds.Get(req.GetBuildId()) != nil { telemetry.SetAttributes(ctx, attribute.Bool("uploaded", true)) return &orchestrator.GetBuildFileExistsResponse{Availability: peerUseStorage}, nil @@ -124,7 +124,7 @@ func (s *Server) ReadAtBuildSeekable(req *orchestrator.ReadAtBuildSeekableReques attribute.Int64("length", length), ) - if s.uploadedBuilds.Has(req.GetBuildId()) { + if s.uploadedBuilds.Get(req.GetBuildId()) != nil { telemetry.SetAttributes(ctx, attribute.Bool("uploaded", true)) return stream.Send(&orchestrator.ReadAtBuildSeekableResponse{Availability: peerUseStorage}) @@ -159,7 +159,7 @@ func (s *Server) GetBuildBlob(req *orchestrator.GetBuildBlobRequest, stream orch attribute.String("file_name", req.GetName()), ) - if s.uploadedBuilds.Has(req.GetBuildId()) { + if s.uploadedBuilds.Get(req.GetBuildId()) != nil { telemetry.SetAttributes(ctx, attribute.Bool("uploaded", true)) return stream.Send(&orchestrator.GetBuildBlobResponse{Availability: peerUseStorage}) From 898f1dd9ddf824230db758411bf24ee15ff3ed2a Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Tue, 9 Jun 2026 16:16:09 -0700 Subject: [PATCH 6/6] nits --- packages/orchestrator/pkg/sandbox/build/build.go | 5 +++-- packages/orchestrator/pkg/sandbox/build/storage_diff.go | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/build/build.go b/packages/orchestrator/pkg/sandbox/build/build.go index 55816b6850..98452ab000 100644 --- a/packages/orchestrator/pkg/sandbox/build/build.go +++ b/packages/orchestrator/pkg/sandbox/build/build.go @@ -117,9 +117,10 @@ type readSegment struct { srcOff int64 length int64 diff Diff - // ft uses the nil-vs-empty convention: nil = no entry (hint unknown), + + // ft uses the nil-vs-empty convention: nil = no entry, // storage.UncompressedFrameTable = authoritatively uncompressed, - // non-empty = compressed. + // non-empty = see ft.compressionType. ft *storage.FrameTable } diff --git a/packages/orchestrator/pkg/sandbox/build/storage_diff.go b/packages/orchestrator/pkg/sandbox/build/storage_diff.go index 88165f3ce1..0fa61f9bee 100644 --- a/packages/orchestrator/pkg/sandbox/build/storage_diff.go +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff.go @@ -77,7 +77,7 @@ func newStorageDiff( isActivePeer IsActivePeer, upstream storage.Seekable, uncompressedSize int64, - initFT *storage.FrameTable, + initialFT *storage.FrameTable, ff *featureflags.Client, ) (*StorageDiff, error) { cachePath := GenerateDiffCachePath(basePath, buildID, diffType) @@ -98,7 +98,7 @@ func newStorageDiff( chunker: c, cacheKey: GetDiffStoreKey(buildID, diffType), } - d.source.Store(&source{upstream: upstream, fullDiffFrameTable: initFT}) + d.source.Store(&source{upstream: upstream, fullDiffFrameTable: initialFT}) return d, nil }