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 b9784d576f..7b8db7b570 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 826e0e6db4..14c6668b8d 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 { @@ -87,12 +111,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 } @@ -103,7 +121,11 @@ type readSegment struct { srcOff int64 length int64 diff Diff - ft *storage.FrameTable + + // ft uses the nil-vs-empty convention: nil = no entry, + // storage.UncompressedFrameTable = authoritatively uncompressed, + // non-empty = see ft.compressionType. + ft *storage.FrameTable } func (b *File) readSegments(ctx context.Context, p []byte, segments []readSegment, maxParallel int) error { @@ -127,10 +149,28 @@ func (b *File) readSegments(ctx context.Context, p []byte, segments []readSegmen return nil } +// 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 { - n, err := s.diff.ReadAt(ctx, p[s.dstOff:s.dstOff+int(s.length)], s.srcOff, s.ft) + 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 + var transitionErr *storage.PeerTransitionedError + if !errors.As(err, &transitionErr) { + return err + } + 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) + } + n, err = s.diff.ReadAt(ctx, dst, s.srcOff, s.ft) + if err != nil { + return err + } } if int64(n) != s.length { return io.ErrUnexpectedEOF @@ -174,7 +214,7 @@ func (b *File) planRead(ctx context.Context, p []byte, off int64) (segments []re 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, 0, err } @@ -191,16 +231,14 @@ func (b *File) planRead(ctx context.Context, p []byte, off int64) (segments []re return segments, n, len(cacheIDs), 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 } } - // 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) if err != nil { return nil, fmt.Errorf("failed to get build: %w", err) } @@ -223,9 +261,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 { @@ -235,7 +272,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. } } } @@ -283,73 +320,194 @@ 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, 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, swapping header", - zap.String("file_type", string(b.fileType)), - ) + return nil +} - 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) +// 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 true, nil + return b.store.GetOrCreate(ctx, key, func(ctx context.Context) (Diff, error) { + return b.createDiff(ctx, buildID) + }) } -// 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 +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} } - return 0 -} + var ( + upstream storage.Seekable + size int64 + initialCT storage.CompressionType + initialFT *storage.FullFrameTable + ) + + 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 + } -func (b *File) getBuild(ctx context.Context, buildID uuid.UUID, uncompressedSize int64, ct storage.CompressionType) (Diff, error) { - storageDiff, err := newStorageDiff( + default: + initialFT = storage.UncompressedFullFrameTable + } + + 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, ct, + 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("createDiff: open upstream for build %s at %s: %w", buildID, path, err) + } + + return upstream, nil +} + +func (b *File) refreshAncestorAndOpenUpstream(ctx context.Context, buildID uuid.UUID, objType storage.SeekableObjectType) (storage.Seekable, int64, *storage.FullFrameTable, error) { + loaded, err := refreshBuildHeader(ctx, b.persistence, buildID, b.fileType, refreshCauseProactive) if err != nil { - return nil, fmt.Errorf("failed to create storage diff: %w", err) + 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) + } } - source, err := b.store.Get(ctx, storageDiff) + // 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 := loaded.SelfBuildData() if err != nil { - return nil, fmt.Errorf("failed to get build from store: %w", err) + return nil, 0, nil, fmt.Errorf("createDiff: %w", err) + } + + upstream, err := b.openUpstream(ctx, buildID, objType, ft.Table().CompressionType()) + if err != nil { + return nil, 0, nil, err + } + + return upstream, 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..b30bd06ec6 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,23 @@ 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. +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 +152,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 +163,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 0e3761a4f7..aadb0b3ce0 100644 --- a/packages/orchestrator/pkg/sandbox/build/diff.go +++ b/packages/orchestrator/pkg/sandbox/build/diff.go @@ -15,6 +15,8 @@ import ( type DiffType string +type IsActivePeer func(buildID string) bool + type NoDiffError struct{} func (NoDiffError) Error() string { @@ -28,13 +30,21 @@ 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 + // 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{} @@ -69,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 88a18debca..8bb106f1f8 100644 --- a/packages/orchestrator/pkg/sandbox/build/mocks/mockdiff.go +++ b/packages/orchestrator/pkg/sandbox/build/mocks/mockdiff.go @@ -293,60 +293,9 @@ 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, 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 +304,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 +326,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 +366,58 @@ 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 +} + +// 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 } diff --git a/packages/orchestrator/pkg/sandbox/build/read_metrics_test.go b/packages/orchestrator/pkg/sandbox/build/read_metrics_test.go index eb35aa315a..a2ff76a694 100644 --- a/packages/orchestrator/pkg/sandbox/build/read_metrics_test.go +++ b/packages/orchestrator/pkg/sandbox/build/read_metrics_test.go @@ -100,6 +100,7 @@ func TestFileReadAt_RecordsZeroFanoutForNilMapping(t *testing.T) { t.TempDir(), time.Hour, time.Minute, + nil, ) require.NoError(t, err) diff --git a/packages/orchestrator/pkg/sandbox/build/storage_diff.go b/packages/orchestrator/pkg/sandbox/build/storage_diff.go index eea5c839f6..6809077e1b 100644 --- a/packages/orchestrator/pkg/sandbox/build/storage_diff.go +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff.go @@ -5,25 +5,58 @@ package build import ( "context" "fmt" + "sync" + "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" "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" +) + +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.UncompressedFullFrameTable authoritatively uncompressed (only set by refresh) +// ft non-empty authoritatively compressed with the bound full-file FT +// +// fullDiffFrameTable is *FullFrameTable rather than *FrameTable: this is the +// one place in the read path where we hold an upcasted full table. The +// invariant — builds[self] for an ancestor we just refreshed is a complete +// table, never a trimmed one — is documented at (*header.Header).SelfBuildData. +// Everywhere else, FrameTables are treated as potentially partial +// (per-mapping, trimmed). +type source struct { + upstream storage.RangeOpener + fullDiffFrameTable *storage.FullFrameTable +} + type StorageDiff struct { chunker *block.Chunker cachePath string cacheKey DiffStoreKey - storagePath string + buildID string + diffType DiffType storageObjectType storage.SeekableObjectType - blockSize int64 - metrics blockmetrics.Metrics - persistence storage.StorageProvider - featureFlags *featureflags.Client - uncompressedSize int64 + blockSize int64 + metrics blockmetrics.Metrics + persistence storage.StorageProvider + isActivePeer IsActivePeer + + source atomic.Pointer[source] + refreshMu sync.Mutex } var _ Diff = (*StorageDiff)(nil) @@ -36,35 +69,45 @@ 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, - ct storage.CompressionType, + initialFT *storage.FullFrameTable, 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) - - return &StorageDiff{ - storagePath: storage.Paths{BuildID: buildId}.DataFile(string(diffType), ct), + d := &StorageDiff{ + buildID: buildID, + diffType: diffType, storageObjectType: storageObjectType, cachePath: cachePath, blockSize: blockSize, metrics: metrics, persistence: persistence, - featureFlags: ff, - uncompressedSize: uncompressedSize, - cacheKey: GetDiffStoreKey(buildId, diffType), - }, nil + isActivePeer: isActivePeer, + chunker: c, + cacheKey: GetDiffStoreKey(buildID, diffType), + } + d.source.Store(&source{upstream: upstream, fullDiffFrameTable: initialFT}) + + return d, nil } func storageObjectType(diffType DiffType) (storage.SeekableObjectType, bool) { @@ -82,44 +125,26 @@ func (b *StorageDiff) CacheKey() DiffStoreKey { return b.cacheKey } -func (b *StorageDiff) Init(ctx context.Context) error { - obj, err := b.persistence.OpenSeekable(ctx, b.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) - } - } +func (b *StorageDiff) Close() error { + return b.chunker.Close() +} - c, err := block.NewChunker(b.featureFlags, size, b.blockSize, obj, b.cachePath, b.metrics) +func (b *StorageDiff) ReadAt(ctx context.Context, p []byte, off int64, callerFT *storage.FrameTable) (int, error) { + up, ft, err := b.resolve(ctx, callerFT) if err != nil { - return fmt.Errorf("failed to create chunker: %w", err) + return 0, err } - b.chunker = c - - return nil + return b.chunker.ReadAt(ctx, p, off, up, ft) } -func (b *StorageDiff) Close() error { - if b.chunker == nil { - return nil +func (b *StorageDiff) Slice(ctx context.Context, off, length int64, callerFT *storage.FrameTable) ([]byte, error) { + up, ft, err := b.resolve(ctx, callerFT) + if err != nil { + return nil, err } - 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) Slice(ctx context.Context, off, length int64, ft *storage.FrameTable) ([]byte, error) { - return b.chunker.Slice(ctx, off, length, ft) + return b.chunker.Slice(ctx, off, length, up, ft) } // The local file might not be synced. @@ -128,15 +153,11 @@ 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(ctx context.Context) (int64, error) { - return b.FileSize(ctx) +func (b *StorageDiff) Size(_ context.Context) (int64, error) { + return b.chunker.Size(), nil } func (b *StorageDiff) BlockSize() int64 { @@ -144,12 +165,116 @@ 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) +} + +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 b.chunker.IsCached(ctx, off, length) + 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.Table(), 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.Table(), 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) +} + +// 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) + } + 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 := loaded.SelfBuildData() + if err != nil { + return fmt.Errorf("reloadSourceLocked: build %s (cause=%s): %w", b.buildID, cause, err) + } + newPath := storage.Paths{BuildID: b.buildID}.DataFile(string(b.diffType), ft.Table().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 new file mode 100644 index 0000000000..c9a281d151 --- /dev/null +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go @@ -0,0 +1,370 @@ +//go:build linux + +package build + +import ( + "bytes" + "context" + "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" +) + +// 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 + 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.Table()}) + 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 loads A's own header — no OpenSeekable at the uncompressed path. + 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() + + // 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). + 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)) +} + +// 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.Table()}) + 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 (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 ( + blockSize = int64(4 << 10) + payloadSize = 8 << 10 + readLen = blockSize + ) + + selfID := uuid.New() + + // 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() + + // No OpenSeekable / Size / OpenRangeReader expectations: refresh must error + // out before any upstream is opened. If any of these fire the mock fails. + + 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) + + _, 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") +} + +// Finalized B + nil FrameData on A means "A is uncompressed and that's the +// 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() + + 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)) +} + +// 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() + 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) + 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, 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. +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()) + } +} 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 ad337225fe..792a18c792 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) { @@ -92,15 +89,13 @@ func (s *peerSeekable) Size(ctx context.Context) (int64, error) { return res.value, err } - // 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 - } - - 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) { @@ -129,12 +124,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 +133,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.FullFrameTable, [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 8173ba376a..fd8e8a1cfe 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) { @@ -143,31 +144,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 +168,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.NewFullFrameTable(storage.CompressionNone, nil).Table()) require.NoError(t, err) @@ -208,19 +189,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.NewFullFrameTable(storage.CompressionNone, nil).Table()) 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.NewFullFrameTable(storage.CompressionZstd, nil).Table()) - 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 3afb1da16d..c7170947b8 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 ea937927cb..aa37f68326 100644 --- a/packages/shared/pkg/storage/compress_frame_table.go +++ b/packages/shared/pkg/storage/compress_frame_table.go @@ -70,12 +70,23 @@ 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{} + // FullFrameTable marks a FrameTable that covers an entire file with no gaps. // Produced by compressStream / CompressBytes / StoreFile / UploadFramed. // The inner FrameTable is unexported and unembedded so methods are not // promoted; consumers must reach functionality through nil-safe Table(). type FullFrameTable struct{ ft FrameTable } +// UncompressedFullFrameTable is the Full counterpart of UncompressedFrameTable: +// the canonical sentinel an authoritative source latches when a build is known +// to be stored uncompressed end-to-end. Shared across callers; read-only. +var UncompressedFullFrameTable = &FullFrameTable{} + // Table returns the underlying *FrameTable, nil-safely. func (ft *FullFrameTable) Table() *FrameTable { if ft == nil { @@ -85,6 +96,19 @@ func (ft *FullFrameTable) Table() *FrameTable { return &ft.ft } +// FullFromTable promotes a *FrameTable to a *FullFrameTable. The caller MUST +// guarantee ft describes the entire file with no gaps — the only legitimate +// promotion is from a self-build's loaded.Builds[self].FrameData entry, which +// build_upload_v4 unconditionally populates from the full upload. Returns the +// UncompressedFullFrameTable sentinel for nil or empty inputs. +func FullFromTable(ft *FrameTable) *FullFrameTable { + if ft == nil || len(ft.entries) == 0 { + return UncompressedFullFrameTable + } + + return &FullFrameTable{ft: *ft} +} + // 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 b1852e2ff9..5dee6c9679 100644 --- a/packages/shared/pkg/storage/header/header.go +++ b/packages/shared/pkg/storage/header/header.go @@ -178,14 +178,38 @@ 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, +// storage.UncompressedFrameTable = 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.UncompressedFrameTable + } + + return bd.FrameData +} + +// SelfBuildData returns the size and full FrameTable for the header's own +// build (t.Metadata.BuildId). Errors when the self entry is missing — only +// possible with peer-served incomplete headers, never with storage-uploaded +// ones (build_upload_v4 always populates self before publish). +// +// This is the *only* place the FullFrameTable upcast happens in production: +// Builds[id].FrameData is typed *FrameTable to match the trimmed-FT case +// where our header carries only the frames we mapped from an ancestor. For +// the self entry, build_upload_v4 always stores the complete table, so the +// upcast via storage.FullFromTable is sound. +func (t *Header) SelfBuildData() (int64, *storage.FullFrameTable, error) { + bd, hasSelf := t.Builds[t.Metadata.BuildId] + if !hasSelf { + return 0, nil, fmt.Errorf("header for build %s has no self entry (peer-served incomplete?)", t.Metadata.BuildId) + } - return t.Builds[buildID].FrameData + return bd.Size, storage.FullFromTable(bd.FrameData), nil } func (t *Header) getMapping(ctx context.Context, offset int64) (BuildMap, int64, error) { 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 022e0edecb..93418a5fef 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.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.Nil(t, got.GetBuildFrameData(selfID)) - require.Nil(t, 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) @@ -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 FrameTable", 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 FrameTable", 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 1515444d0e..9a2c463c3b 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 a7965fd636..c24bf5fbef 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 828da55639..ea064ef46a 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 1bfb689729..9f1b81c1d1 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) { 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)