From 97809d33ee43c2e5dc9f7e5a185627b086de769d Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Wed, 27 May 2026 13:02:59 -0700 Subject: [PATCH 01/11] refactor(storage): consolidate read-path io wrappers Replace the OpenRangeReader return type io.ReadCloser with a context-aware RangeReader and collapse scattered reader wrappers (offset/section/capture/ observable/decompress) into a single io_wrappers.go. Drops the cacheWriteThroughReader in favor of a reusable captureReader with optional drain-on-close for codecs that stop short of EOF on their source. --- .vscode/settings.json | 2 +- .../pkg/sandbox/block/streaming_chunk.go | 2 +- .../pkg/sandbox/block/streaming_chunk_test.go | 16 +- .../pkg/sandbox/template/peerclient/blob.go | 2 +- .../sandbox/template/peerclient/seekable.go | 9 +- .../template/peerclient/seekable_test.go | 12 +- .../sandbox/template/peerclient/storage.go | 6 +- .../shared/pkg/storage/compress_decode.go | 86 ++++------ packages/shared/pkg/storage/io_wrappers.go | 160 ++++++++++++++++++ ...set_reader_test.go => io_wrappers_test.go} | 0 packages/shared/pkg/storage/mock_seekable.go | 17 +- packages/shared/pkg/storage/offset_reader.go | 23 --- packages/shared/pkg/storage/storage.go | 7 +- packages/shared/pkg/storage/storage_aws.go | 4 +- packages/shared/pkg/storage/storage_cache.go | 4 + .../storage/storage_cache_compressed_test.go | 35 ++-- .../pkg/storage/storage_cache_seekable.go | 136 ++++----------- .../storage_cache_seekable_compressed.go | 134 ++++----------- .../storage/storage_cache_seekable_test.go | 98 +++++------ packages/shared/pkg/storage/storage_fs.go | 25 +-- packages/shared/pkg/storage/storage_google.go | 59 ++----- .../sandbox_rapid_pause_resume_test.go | 2 +- 22 files changed, 378 insertions(+), 461 deletions(-) create mode 100644 packages/shared/pkg/storage/io_wrappers.go rename packages/shared/pkg/storage/{offset_reader_test.go => io_wrappers_test.go} (100%) delete mode 100644 packages/shared/pkg/storage/offset_reader.go diff --git a/.vscode/settings.json b/.vscode/settings.json index 42900fa357..f1b96e2a42 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -52,7 +52,7 @@ }, }, "editor.tabCompletion": "on", - "go.lintTool": "golangci-lint", + "go.lintTool": "golangci-lint-v2", "go.lintFlags": [ "--path-mode=abs", "--fast-only", diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go index f920cdccd0..839cbab5fc 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go @@ -246,7 +246,7 @@ func (c *Chunker) progressiveRead(ctx context.Context, s *fetchSession, mmapSlic return 0, fmt.Errorf("failed to open range reader at %d: %w", s.chunkOff, err) } defer func() { - if closeErr := reader.Close(); closeErr != nil && err == nil { + if closeErr := reader.Close(ctx); closeErr != nil && err == nil { err = closeErr } }() diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go index 2263dc7c4e..ae47c70d65 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go @@ -82,7 +82,7 @@ func (s *fakeSeekable) StoreFile(context.Context, string, ...storage.PutOption) panic("not used") } -func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length int64, frameTable *storage.FrameTable) (io.ReadCloser, error) { +func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length int64, frameTable *storage.FrameTable) (storage.RangeReader, error) { s.fetchCount.Add(1) if s.ctrl != nil { @@ -97,13 +97,13 @@ func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length end := min(offsetU+length, int64(len(s.data))) - return &controlledReader{ + return storage.NewRangeReader(&controlledReader{ data: s.data[offsetU:end], step: max(16*1024, testBlockSize), advance: s.ctrl.advance, consumed: s.ctrl.consumed, closed: s.ctrl.closed, - }, nil + }), nil } var fetchOff, fetchLen int64 @@ -127,10 +127,10 @@ func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length r := io.Reader(bytes.NewReader(s.data[fetchOff:end])) if frameTable.IsCompressed() { - return storage.NewDecompressingReader(r, frameTable.CompressionType()) + return storage.NewDecompressingReader(storage.NewRangeReader(io.NopCloser(r)), frameTable.CompressionType()) } - return io.NopCloser(r), nil + return storage.NewRangeReader(io.NopCloser(r)), nil } func makeCompressedTestData(tb testing.TB, data []byte) (*storage.FrameTable, *fakeSeekable) { @@ -428,13 +428,13 @@ func (s *panicSeekable) StoreFile(context.Context, string, ...storage.PutOption) panic("not used") } -func (s *panicSeekable) OpenRangeReader(_ context.Context, off int64, length int64, _ *storage.FrameTable) (io.ReadCloser, error) { +func (s *panicSeekable) OpenRangeReader(_ context.Context, off int64, length int64, _ *storage.FrameTable) (storage.RangeReader, error) { end := min(off+length, int64(len(s.data))) - return &panicReader{ + return storage.NewRangeReader(&panicReader{ data: s.data[off:end], panicAfter: int(s.panicAfter - off), - }, nil + }), nil } type panicReader struct { diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/blob.go b/packages/orchestrator/pkg/sandbox/template/peerclient/blob.go index 65891b2f70..c546f1467c 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/blob.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/blob.go @@ -65,7 +65,7 @@ func (b *peerBlob) WriteTo(ctx context.Context, dst io.Writer) (int64, error) { } reader := newPeerStreamReader(recv, cancel) - defer reader.Close() + defer reader.Close(ctx) n, err := io.Copy(dst, reader) if err != nil { diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go index 4b9230da73..540850205d 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go @@ -4,7 +4,6 @@ import ( "context" "errors" "fmt" - "io" "sync" "sync/atomic" "time" @@ -103,9 +102,9 @@ func (s *peerSeekable) Size(ctx context.Context) (int64, error) { return base.Size(ctx) } -func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *storage.FrameTable) (io.ReadCloser, error) { +func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *storage.FrameTable) (storage.RangeReader, error) { res, err := tryPeer(ctx, &s.peerHandle, "peer-seekable-open-range-reader", attrOpRangeReader, - func(ctx context.Context) (peerAttempt[io.ReadCloser], error) { + func(ctx context.Context) (peerAttempt[storage.RangeReader], error) { streamCtx, cancel := context.WithCancel(ctx) recv, err := openPeerSeekableStream(streamCtx, s.client, &orchestrator.ReadAtBuildSeekableRequest{ @@ -118,10 +117,10 @@ func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length in logger.L().Warn(ctx, "failed to open range reader from peer", logger.WithBuildID(s.buildID), zap.Int64("off", off), zap.Int64("length", length), zap.Error(err)) cancel() - return peerAttempt[io.ReadCloser]{}, nil + return peerAttempt[storage.RangeReader]{}, nil } - return peerAttempt[io.ReadCloser]{ + return peerAttempt[storage.RangeReader]{ value: newPeerStreamReader(recv, cancel), hit: true, }, nil diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go index b5266e9094..e90187d971 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go @@ -74,7 +74,7 @@ func TestPeerSeekable_OpenRangeReader_PeerSucceeds(t *testing.T) { s := &peerSeekable{peerHandle: peerHandle{client: client, buildID: "build-1", name: storage.MemfileName, uploaded: &atomic.Bool{}}} rc, err := s.OpenRangeReader(t.Context(), 10, int64(len(data)), nil) require.NoError(t, err) - defer rc.Close() + defer rc.Close(t.Context()) got, err := io.ReadAll(rc) require.NoError(t, err) @@ -89,7 +89,7 @@ func TestPeerSeekable_OpenRangeReader_PeerError_FallsBackToBase(t *testing.T) { client.EXPECT().ReadAtBuildSeekable(mock.Anything, mock.Anything).Return(nil, errors.New("peer unavailable")) baseSeekable := storage.NewMockSeekable(t) - baseSeekable.EXPECT().OpenRangeReader(mock.Anything, int64(0), int64(len(baseData)), (*storage.FrameTable)(nil)).Return(io.NopCloser(bytes.NewReader(baseData)), nil) + baseSeekable.EXPECT().OpenRangeReader(mock.Anything, int64(0), int64(len(baseData)), (*storage.FrameTable)(nil)).Return(storage.NewRangeReader(io.NopCloser(bytes.NewReader(baseData))), nil) base := storage.NewMockStorageProvider(t) base.EXPECT().OpenSeekable(mock.Anything, "build-1/memfile", storage.MemfileObjectType).Return(baseSeekable, nil) @@ -106,7 +106,7 @@ func TestPeerSeekable_OpenRangeReader_PeerError_FallsBackToBase(t *testing.T) { } rc, err := s.OpenRangeReader(t.Context(), 0, int64(len(baseData)), nil) require.NoError(t, err) - defer rc.Close() + defer rc.Close(t.Context()) got, err := io.ReadAll(rc) require.NoError(t, err) @@ -187,7 +187,7 @@ func TestPeerStorageProvider_FullTransitionFlow(t *testing.T) { postBaseSeekable := storage.NewMockSeekable(t) postBaseSeekable.EXPECT(). OpenRangeReader(mock.Anything, int64(0), int64(len(postBaseBytes)), mock.Anything). - Return(io.NopCloser(bytes.NewReader(postBaseBytes)), nil).Once() + Return(storage.NewRangeReader(io.NopCloser(bytes.NewReader(postBaseBytes))), nil).Once() base := storage.NewMockStorageProvider(t) base.EXPECT(). @@ -204,7 +204,7 @@ func TestPeerStorageProvider_FullTransitionFlow(t *testing.T) { require.NoError(t, err) got, err := io.ReadAll(rc) require.NoError(t, err) - require.NoError(t, rc.Close()) + require.NoError(t, rc.Close(t.Context())) assert.Equal(t, prePeerBytes, got) require.True(t, uploaded.Load(), "uploaded flag should be set after peer EOF with UseStorage") @@ -221,6 +221,6 @@ func TestPeerStorageProvider_FullTransitionFlow(t *testing.T) { require.NoError(t, err) got, err = io.ReadAll(rc) require.NoError(t, err) - require.NoError(t, rc.Close()) + require.NoError(t, rc.Close(t.Context())) 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..c6313149af 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/storage.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/storage.go @@ -260,9 +260,9 @@ func tryPeer[T any]( return peerAttempt[T]{}, nil } -var _ io.ReadCloser = (*peerStreamReader)(nil) +var _ storage.RangeReader = (*peerStreamReader)(nil) -// peerStreamReader wraps a gRPC streaming recv function as an io.ReadCloser. +// peerStreamReader wraps a gRPC streaming recv function as a storage.RangeReader. // cancel is called on Close to signal the server to terminate the stream. type peerStreamReader struct { recv func() ([]byte, error) @@ -304,7 +304,7 @@ func (r *peerStreamReader) Read(p []byte) (int, error) { } } -func (r *peerStreamReader) Close() error { +func (r *peerStreamReader) Close(context.Context) error { r.cancel() return nil diff --git a/packages/shared/pkg/storage/compress_decode.go b/packages/shared/pkg/storage/compress_decode.go index b8a6d414a7..a0a706bab3 100644 --- a/packages/shared/pkg/storage/compress_decode.go +++ b/packages/shared/pkg/storage/compress_decode.go @@ -1,6 +1,7 @@ package storage import ( + "context" "fmt" "io" "sync" @@ -9,6 +10,8 @@ import ( lz4 "github.com/pierrec/lz4/v4" ) +var _ RangeReader = (*decompressReader)(nil) + var lz4DecoderPool sync.Pool func getLZ4Decoder(r io.Reader) *lz4.Reader { @@ -53,78 +56,47 @@ func putZstdDecoder(dec *zstd.Decoder) { zstdDecoderPool.Put(dec) } -// NewDecompressingReader wraps a reader with the appropriate decompressor. -// Close releases the decompressor back to its pool but does NOT close the -// underlying reader — the caller is responsible for closing it. -func NewDecompressingReader(raw io.Reader, ct CompressionType) (io.ReadCloser, error) { +// decompressReader decompresses inner on Read; Close releases the codec back +// to its pool and closes inner. +type decompressReader struct { + inner RangeReader + dec io.Reader + releaseCodec func() +} + +func NewDecompressingReader(inner RangeReader, ct CompressionType) (RangeReader, error) { + var dec io.Reader + var releaseCodec func() + switch ct { case CompressionLZ4: - dec := getLZ4Decoder(raw) - - return &pooledDecoder{ - Reader: dec, - close: func() { putLZ4Decoder(dec) }, - }, nil + d := getLZ4Decoder(inner) + dec, releaseCodec = d, func() { putLZ4Decoder(d) } case CompressionZstd: - dec, err := getZstdDecoder(raw) + d, err := getZstdDecoder(inner) if err != nil { return nil, fmt.Errorf("failed to create zstd decoder: %w", err) } - - return &pooledDecoder{ - Reader: dec, - close: func() { putZstdDecoder(dec) }, - }, nil + dec, releaseCodec = d, func() { putZstdDecoder(d) } default: return nil, fmt.Errorf("unsupported compression type: %s", ct) } -} - -// pooledDecoder wraps a decompressor from a sync.Pool. -// Close returns the decompressor to the pool. -type pooledDecoder struct { - io.Reader - close func() + return &decompressReader{ + inner: inner, + dec: dec, + releaseCodec: releaseCodec, + }, nil } -func (r *pooledDecoder) Close() error { - r.close() - - return nil +func (r *decompressReader) Read(p []byte) (int, error) { + return r.dec.Read(p) } -// newDecompressingReadCloser wraps raw with the appropriate decompressor and -// takes ownership: Close releases the decompressor back to the pool AND closes raw. -func newDecompressingReadCloser(raw io.ReadCloser, ct CompressionType) (io.ReadCloser, error) { - dec, err := NewDecompressingReader(raw, ct) - if err != nil { - return nil, err - } - - return &decompressingReadCloser{dec: dec, raw: raw}, nil -} - -// decompressingReadCloser reads from the decompressor and closes both the -// decompressor (returning it to the pool) and the underlying raw stream. -type decompressingReadCloser struct { - dec io.ReadCloser // decompressor — reads from raw - raw io.Closer // underlying stream -} - -func (c *decompressingReadCloser) Read(p []byte) (int, error) { - return c.dec.Read(p) -} - -func (c *decompressingReadCloser) Close() error { - decErr := c.dec.Close() - rawErr := c.raw.Close() - - if decErr != nil { - return decErr - } +func (r *decompressReader) Close(ctx context.Context) error { + r.releaseCodec() - return rawErr + return r.inner.Close(ctx) } diff --git a/packages/shared/pkg/storage/io_wrappers.go b/packages/shared/pkg/storage/io_wrappers.go new file mode 100644 index 0000000000..1e79763b89 --- /dev/null +++ b/packages/shared/pkg/storage/io_wrappers.go @@ -0,0 +1,160 @@ +package storage + +import ( + "bytes" + "context" + "errors" + "io" + "os" + + "go.opentelemetry.io/otel/trace" + + "github.com/e2b-dev/infra/packages/shared/pkg/telemetry" +) + +var ( + _ io.Reader = (*offsetReader)(nil) + _ RangeReader = (*sectionReader)(nil) + _ RangeReader = (*observableReader)(nil) + _ RangeReader = (*rangeReader)(nil) + _ RangeReader = (*captureReader)(nil) +) + +// offsetReader adapts an io.ReaderAt into a sequential io.Reader +// starting at the given offset. +type offsetReader struct { + wrapped io.ReaderAt + offset int64 +} + +func (r *offsetReader) Read(p []byte) (n int, err error) { + n, err = r.wrapped.ReadAt(p, r.offset) + r.offset += int64(n) + + return +} + +func newOffsetReader(reader io.ReaderAt, offset int64) *offsetReader { + return &offsetReader{reader, offset} +} + +// rangeReader adapts an io.ReadCloser into a RangeReader by ignoring the +// Close context. +type rangeReader struct { + io.ReadCloser +} + +func NewRangeReader(rc io.ReadCloser) RangeReader { return &rangeReader{ReadCloser: rc} } + +func (p *rangeReader) Close(context.Context) error { + return p.ReadCloser.Close() +} + +type sectionReader struct { + *io.SectionReader + + file *os.File +} + +func newSectionReader(f *os.File, off, length int64) *sectionReader { + return §ionReader{ + SectionReader: io.NewSectionReader(f, off, length), + file: f, + } +} + +func (r *sectionReader) Close(context.Context) error { + return r.file.Close() +} + +// captureReader tees every read byte into a buffer and hands the captured +// bytes to onClose on Close. Used by the cache writeback paths. +// +// drainOnClose=true reads inner to EOF on Close even if the caller above hasn't +// consumed everything. Needed when capturing under a decoder that stops short +// of EOF on its source (e.g. lz4.Reader skips the 4-byte EndMark). +type captureReader struct { + inner RangeReader + buf *bytes.Buffer + onClose func(ctx context.Context, captured []byte) + drainOnClose bool +} + +func newCaptureReader(inner RangeReader, capHint int, drainOnClose bool, onClose func(context.Context, []byte)) *captureReader { + return &captureReader{ + inner: inner, + buf: bytes.NewBuffer(make([]byte, 0, capHint)), + onClose: onClose, + drainOnClose: drainOnClose, + } +} + +func (r *captureReader) Read(p []byte) (int, error) { + n, err := r.inner.Read(p) + if n > 0 { + r.buf.Write(p[:n]) + } + + return n, err +} + +func (r *captureReader) Close(ctx context.Context) error { + if r.drainOnClose { + _, _ = io.Copy(io.Discard, r) + } + err := r.inner.Close(ctx) + r.onClose(ctx, r.buf.Bytes()) + + return err +} + +// observableReader layers OTEL observability (legacy per-backend timer + span) +// onto an inner RangeReader, all applied on Close. timer and span are optional; +// pass nil if unused. +type observableReader struct { + inner RangeReader + timer *telemetry.Stopwatch + span trace.Span + + bytes int64 + readErr error +} + +func newObservableReader(inner RangeReader, timer *telemetry.Stopwatch, span trace.Span) *observableReader { + return &observableReader{inner: inner, timer: timer, span: span} +} + +func (r *observableReader) Read(p []byte) (int, error) { + n, err := r.inner.Read(p) + r.bytes += int64(n) + + if err != nil && !errors.Is(err, io.EOF) { + r.readErr = err + } + + return n, err +} + +func (r *observableReader) Close(ctx context.Context) error { + closeErr := r.inner.Close(ctx) + + if r.timer != nil { + if r.readErr != nil || closeErr != nil { + r.timer.Failure(ctx, r.bytes) + } else { + r.timer.Success(ctx, r.bytes) + } + } + + if r.span != nil { + if closeErr != nil { + recordError(r.span, closeErr) + } else if r.readErr != nil { + recordError(r.span, r.readErr) + } + + r.span.End() + } + + return closeErr +} diff --git a/packages/shared/pkg/storage/offset_reader_test.go b/packages/shared/pkg/storage/io_wrappers_test.go similarity index 100% rename from packages/shared/pkg/storage/offset_reader_test.go rename to packages/shared/pkg/storage/io_wrappers_test.go diff --git a/packages/shared/pkg/storage/mock_seekable.go b/packages/shared/pkg/storage/mock_seekable.go index 440836bc54..e33c5d8141 100644 --- a/packages/shared/pkg/storage/mock_seekable.go +++ b/packages/shared/pkg/storage/mock_seekable.go @@ -6,7 +6,6 @@ package storage import ( "context" - "io" mock "github.com/stretchr/testify/mock" ) @@ -39,23 +38,23 @@ func (_m *MockSeekable) EXPECT() *MockSeekable_Expecter { } // OpenRangeReader provides a mock function for the type MockSeekable -func (_mock *MockSeekable) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (io.ReadCloser, error) { +func (_mock *MockSeekable) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, error) { ret := _mock.Called(ctx, offsetU, length, frameTable) if len(ret) == 0 { panic("no return value specified for OpenRangeReader") } - var r0 io.ReadCloser + var r0 RangeReader var r1 error - if returnFunc, ok := ret.Get(0).(func(context.Context, int64, int64, *FrameTable) (io.ReadCloser, error)); ok { + if returnFunc, ok := ret.Get(0).(func(context.Context, int64, int64, *FrameTable) (RangeReader, error)); ok { return returnFunc(ctx, offsetU, length, frameTable) } - if returnFunc, ok := ret.Get(0).(func(context.Context, int64, int64, *FrameTable) io.ReadCloser); ok { + if returnFunc, ok := ret.Get(0).(func(context.Context, int64, int64, *FrameTable) RangeReader); ok { r0 = returnFunc(ctx, offsetU, length, frameTable) } else { if ret.Get(0) != nil { - r0 = ret.Get(0).(io.ReadCloser) + r0 = ret.Get(0).(RangeReader) } } if returnFunc, ok := ret.Get(1).(func(context.Context, int64, int64, *FrameTable) error); ok { @@ -108,12 +107,12 @@ func (_c *MockSeekable_OpenRangeReader_Call) Run(run func(ctx context.Context, o return _c } -func (_c *MockSeekable_OpenRangeReader_Call) Return(readCloser io.ReadCloser, err error) *MockSeekable_OpenRangeReader_Call { - _c.Call.Return(readCloser, err) +func (_c *MockSeekable_OpenRangeReader_Call) Return(rangeReader RangeReader, err error) *MockSeekable_OpenRangeReader_Call { + _c.Call.Return(rangeReader, err) return _c } -func (_c *MockSeekable_OpenRangeReader_Call) RunAndReturn(run func(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (io.ReadCloser, error)) *MockSeekable_OpenRangeReader_Call { +func (_c *MockSeekable_OpenRangeReader_Call) RunAndReturn(run func(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, error)) *MockSeekable_OpenRangeReader_Call { _c.Call.Return(run) return _c } diff --git a/packages/shared/pkg/storage/offset_reader.go b/packages/shared/pkg/storage/offset_reader.go deleted file mode 100644 index 29d9048d6c..0000000000 --- a/packages/shared/pkg/storage/offset_reader.go +++ /dev/null @@ -1,23 +0,0 @@ -package storage - -import ( - "io" -) - -type offsetReader struct { - wrapped io.ReaderAt - offset int64 -} - -var _ io.Reader = (*offsetReader)(nil) - -func (r *offsetReader) Read(p []byte) (n int, err error) { - n, err = r.wrapped.ReadAt(p, r.offset) - r.offset += int64(n) - - return -} - -func newOffsetReader(reader io.ReaderAt, offset int64) *offsetReader { - return &offsetReader{reader, offset} -} diff --git a/packages/shared/pkg/storage/storage.go b/packages/shared/pkg/storage/storage.go index 668ca49fb4..d1cf6a8825 100644 --- a/packages/shared/pkg/storage/storage.go +++ b/packages/shared/pkg/storage/storage.go @@ -146,9 +146,14 @@ type SeekableReader interface { Size(ctx context.Context) (int64, error) } +type RangeReader interface { + io.Reader + Close(ctx context.Context) error +} + // StreamingReader supports progressive reads via a streaming range reader. type StreamingReader interface { - OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (io.ReadCloser, error) + OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, error) } type SeekableWriter interface { diff --git a/packages/shared/pkg/storage/storage_aws.go b/packages/shared/pkg/storage/storage_aws.go index 5e62a784b9..8429d33e61 100644 --- a/packages/shared/pkg/storage/storage_aws.go +++ b/packages/shared/pkg/storage/storage_aws.go @@ -233,7 +233,7 @@ func (o *awsObject) Put(ctx context.Context, data []byte, opts ...PutOption) err return nil } -func (o *awsObject) OpenRangeReader(ctx context.Context, off, length int64, frameTable *FrameTable) (io.ReadCloser, error) { +func (o *awsObject) OpenRangeReader(ctx context.Context, off, length int64, frameTable *FrameTable) (RangeReader, error) { if frameTable.IsCompressed() { return nil, errors.New("compressed reads are not supported on AWS") } @@ -253,7 +253,7 @@ func (o *awsObject) OpenRangeReader(ctx context.Context, off, length int64, fram return nil, fmt.Errorf("failed to create S3 range reader for %q: %w", o.path, err) } - return resp.Body, nil + return NewRangeReader(resp.Body), nil } func (o *awsObject) Size(ctx context.Context) (int64, error) { diff --git a/packages/shared/pkg/storage/storage_cache.go b/packages/shared/pkg/storage/storage_cache.go index 9f93559662..7d5838758c 100644 --- a/packages/shared/pkg/storage/storage_cache.go +++ b/packages/shared/pkg/storage/storage_cache.go @@ -151,6 +151,10 @@ func ignoreEOF(err error) error { // isCompleteRead reports whether a read of n bytes into a buffer of expected // size represents a valid, cacheable result. A read is complete when either // the full buffer was filled or io.EOF explains a non-empty short read (last chunk). +// +// Writeback callers pass err=nil: a streaming reader always ends in io.EOF +// regardless of whether the upstream was truncated, so the byte count is the +// only reliable signal that the captured bytes are safe to cache. func isCompleteRead(n, expected int, err error) bool { return n == expected || (n > 0 && errors.Is(err, io.EOF)) } diff --git a/packages/shared/pkg/storage/storage_cache_compressed_test.go b/packages/shared/pkg/storage/storage_cache_compressed_test.go index 8314abb0a1..d19b396c15 100644 --- a/packages/shared/pkg/storage/storage_cache_compressed_test.go +++ b/packages/shared/pkg/storage/storage_cache_compressed_test.go @@ -68,19 +68,16 @@ func TestDecompressingCacheReader(t *testing.T) { c := newTestCache(t) framePath := makeFrameFilename(c.path, Range{Offset: 0, Length: len(compressed)}) - rc, err := newDecompressingCacheReader( - io.NopCloser(bytes.NewReader(compressed)), - CompressionLZ4, - len(compressed), - &c, t.Context(), framePath, 0, - ) + capturing := newCaptureReader(bytesRangeReader(compressed), len(compressed), true, + c.compressedFrameWriteback(framePath, 0, len(compressed))) + rc, err := NewDecompressingReader(capturing, CompressionLZ4) require.NoError(t, err) got, err := io.ReadAll(rc) require.NoError(t, err) require.Equal(t, original, got) - require.NoError(t, rc.Close()) + mustClose(t, rc) c.wg.Wait() cached, err := os.ReadFile(framePath) @@ -104,12 +101,9 @@ func TestDecompressingCacheReader(t *testing.T) { compressedProd := lz4CompressProd(t, original) framePath := makeFrameFilename(c.path, Range{Offset: 0, Length: len(compressedProd)}) - rc, err := newDecompressingCacheReader( - io.NopCloser(bytes.NewReader(compressedProd)), - CompressionLZ4, - len(compressedProd), - &c, t.Context(), framePath, 0, - ) + capturing := newCaptureReader(bytesRangeReader(compressedProd), len(compressedProd), true, + c.compressedFrameWriteback(framePath, 0, len(compressedProd))) + rc, err := NewDecompressingReader(capturing, CompressionLZ4) require.NoError(t, err) out := make([]byte, len(original)) @@ -118,7 +112,8 @@ func TestDecompressingCacheReader(t *testing.T) { require.Equal(t, len(original), n) require.Equal(t, original, out) - require.NoError(t, rc.Close(), "writeback failure must not surface as a read error") + closeErr := rc.Close(t.Context()) + require.NoError(t, closeErr, "writeback failure must not surface as a read error") c.wg.Wait() _, err = os.Stat(framePath) @@ -131,19 +126,17 @@ func TestDecompressingCacheReader(t *testing.T) { c := newTestCache(t) framePath := makeFrameFilename(c.path, Range{Offset: 0, Length: len(compressed)}) - rc, err := newDecompressingCacheReader( - io.NopCloser(bytes.NewReader(compressed)), - CompressionLZ4, - len(compressed)+100, // wrong size - &c, t.Context(), framePath, 0, - ) + capturing := newCaptureReader(bytesRangeReader(compressed), len(compressed)+100, true, + c.compressedFrameWriteback(framePath, 0, len(compressed)+100)) // wrong expected size + rc, err := NewDecompressingReader(capturing, CompressionLZ4) require.NoError(t, err) got, err := io.ReadAll(rc) require.NoError(t, err) require.Equal(t, original, got, "decompressed data should be correct regardless") - require.NoError(t, rc.Close(), "writeback failure must not surface as a read error") + closeErr := rc.Close(t.Context()) + require.NoError(t, closeErr, "writeback failure must not surface as a read error") c.wg.Wait() diff --git a/packages/shared/pkg/storage/storage_cache_seekable.go b/packages/shared/pkg/storage/storage_cache_seekable.go index a3b20c1118..86835bceb7 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable.go +++ b/packages/shared/pkg/storage/storage_cache_seekable.go @@ -1,7 +1,6 @@ package storage import ( - "bytes" "context" "errors" "fmt" @@ -80,7 +79,7 @@ var ( _ StreamingReader = (*cachedSeekable)(nil) ) -func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *FrameTable) (io.ReadCloser, error) { +func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *FrameTable) (RangeReader, error) { compressed := frameTable.IsCompressed() ctx, span := c.tracer.Start(ctx, "read", trace.WithAttributes( @@ -89,27 +88,25 @@ func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length attribute.Bool("compressed", compressed), )) + var rc RangeReader + var err error if compressed { - rc, err := c.openReaderCompressed(ctx, off, frameTable) - if err != nil { - recordError(span, err) - span.End() - - return nil, err - } - - rc = withSpan(rc, span) - - return rc, nil + rc, err = c.openReaderCompressed(ctx, off, frameTable) + } else if err = c.validateReadParams(length, off); err == nil { + rc, err = c.openReaderUncompressed(ctx, off, length) } - if err := c.validateReadParams(length, off); err != nil { + if err != nil { recordError(span, err) span.End() return nil, err } + return newObservableReader(rc, nil, span), nil +} + +func (c *cachedSeekable) openReaderUncompressed(ctx context.Context, off, length int64) (RangeReader, error) { timer := cacheSlabReadTimerFactory.Begin( attribute.String(nfsCacheOperationAttr, nfsCacheOperationAttrReadAt), attribute.Bool("compressed", false), @@ -122,10 +119,7 @@ func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length recordCacheRead(ctx, true, length, cacheTypeSeekable, cacheOpOpenRangeReader) timer.Success(ctx, length) - rc := io.ReadCloser(&fsRangeReadCloser{Reader: io.NewSectionReader(fp, 0, length), file: fp}) - rc = withSpan(rc, span) - - return withNFSGauge(ctx, rc), nil + return withNFSGauge(ctx, newSectionReader(fp, 0, length)), nil } if !os.IsNotExist(err) { @@ -136,122 +130,58 @@ func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length rc, err := c.inner.OpenRangeReader(ctx, off, length, nil) if err != nil { - recordError(span, err) - span.End() - return nil, fmt.Errorf("failed to open inner range reader: %w", err) } recordCacheRead(ctx, false, length, cacheTypeSeekable, cacheOpOpenRangeReader) if !skipCacheWriteback(ctx) { - rc = newCacheWriteThroughReader(rc, c, ctx, off, length, chunkPath) + rc = newCaptureReader(rc, int(length), false, + c.uncompressedChunkWriteback(chunkPath, off, length)) } - rc = withSpan(rc, span) - return rc, nil } -// withSpan wraps a reader with an OTEL span that ends on Close. -func withSpan(rc io.ReadCloser, span trace.Span) io.ReadCloser { - return &spanReadCloser{inner: rc, span: span} -} - -type spanReadCloser struct { - inner io.ReadCloser - span trace.Span -} - -func (r *spanReadCloser) Read(p []byte) (int, error) { - return r.inner.Read(p) -} - -func (r *spanReadCloser) Close() error { - err := r.inner.Close() - recordError(r.span, err) - r.span.End() - - return err -} - // nfsGaugeReadCloser wraps a reader and decrements the NFS concurrent reads // gauge on Close. type nfsGaugeReadCloser struct { - io.ReadCloser - - ctx context.Context //nolint:containedctx // needed for gauge decrement in Close + RangeReader } -func (r *nfsGaugeReadCloser) Close() error { - nfsCacheConcurrentReads.Add(r.ctx, -1) +func (r *nfsGaugeReadCloser) Close(ctx context.Context) error { + nfsCacheConcurrentReads.Add(ctx, -1) - return r.ReadCloser.Close() + return r.RangeReader.Close(ctx) } -func withNFSGauge(ctx context.Context, rc io.ReadCloser) io.ReadCloser { +func withNFSGauge(ctx context.Context, rc RangeReader) RangeReader { nfsCacheConcurrentReads.Add(ctx, 1) - return &nfsGaugeReadCloser{ReadCloser: rc, ctx: ctx} -} - -// newCacheWriteThroughReader wraps a reader, buffering all data read through it. -// On Close, it asynchronously writes the buffered data to the NFS cache only -// if the total bytes read match the expected length (to avoid caching truncated data). -func newCacheWriteThroughReader(inner io.ReadCloser, cache *cachedSeekable, ctx context.Context, off, expectedLen int64, chunkPath string) io.ReadCloser { - return &cacheWriteThroughReader{ - inner: inner, - buf: bytes.NewBuffer(make([]byte, 0, expectedLen)), - cache: cache, - ctx: ctx, - off: off, - expectedLen: expectedLen, - chunkPath: chunkPath, - } + return &nfsGaugeReadCloser{RangeReader: rc} } -type cacheWriteThroughReader struct { - inner io.ReadCloser - buf *bytes.Buffer - cache *cachedSeekable - ctx context.Context //nolint:containedctx // needed for async cache write-back in Close - off int64 - expectedLen int64 - chunkPath string -} - -func (r *cacheWriteThroughReader) Read(p []byte) (int, error) { - n, err := r.inner.Read(p) - if n > 0 { - r.buf.Write(p[:n]) - } - - return n, err -} - -func (r *cacheWriteThroughReader) Close() error { - closeErr := r.inner.Close() - - // Only cache when the total bytes read match the expected length. - // Unlike ReadAt where io.EOF can justify a short read (last chunk), - // a streaming reader always ends with EOF regardless of whether the - // data was truncated, so the byte count is the only reliable check. - if isCompleteRead(r.buf.Len(), int(r.expectedLen), nil) { - data := make([]byte, r.buf.Len()) - copy(data, r.buf.Bytes()) +// uncompressedChunkWriteback returns a captureReader callback that persists +// the captured chunk to the NFS cache in a detached goroutine. Best-effort: +// a short capture (e.g. upstream truncation) is dropped silently — a streaming +// reader always ends in EOF, so byte count is the only reliable signal. +func (c *cachedSeekable) uncompressedChunkWriteback(chunkPath string, off, expectedLen int64) func(context.Context, []byte) { + return func(ctx context.Context, captured []byte) { + if !isCompleteRead(len(captured), int(expectedLen), nil) { + return + } - r.cache.goCtx(r.ctx, func(ctx context.Context) { - ctx, span := r.cache.tracer.Start(ctx, "write range reader chunk back to cache") + c.goCtx(ctx, func(ctx context.Context) { + ctx, span := c.tracer.Start(ctx, "write range reader chunk back to cache") defer span.End() - if err := r.cache.writeToCache(ctx, r.off, r.chunkPath, data); err != nil { + err := c.writeToCache(ctx, off, chunkPath, captured) + if err != nil { recordError(span, err) recordCacheWriteError(ctx, cacheTypeSeekable, cacheOpOpenRangeReader, err) } }) } - - return closeErr } func (c *cachedSeekable) Size(ctx context.Context) (n int64, e error) { diff --git a/packages/shared/pkg/storage/storage_cache_seekable_compressed.go b/packages/shared/pkg/storage/storage_cache_seekable_compressed.go index 4886b449d4..d2e2996776 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable_compressed.go +++ b/packages/shared/pkg/storage/storage_cache_seekable_compressed.go @@ -1,10 +1,8 @@ package storage import ( - "bytes" "context" "fmt" - "io" "os" "go.opentelemetry.io/otel/attribute" @@ -19,13 +17,14 @@ var compressedCacheReadAttrs = []attribute.KeyValue{ // openReaderCompressed handles the compressed cache path for OpenRangeReader. // NFS stores compressed frames (.frm); on hit we decompress, on miss we fetch // raw compressed bytes and tee them to NFS on Close. -func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64, frameTable *FrameTable) (io.ReadCloser, error) { +func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64, frameTable *FrameTable) (RangeReader, error) { r, err := frameTable.LocateCompressed(offsetU) if err != nil { return nil, fmt.Errorf("frame lookup for offset %d: %w", offsetU, err) } path := makeFrameFilename(c.path, r) + ct := frameTable.CompressionType() timer := cacheSlabReadTimerFactory.Begin(compressedCacheReadAttrs...) @@ -37,14 +36,14 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64 recordCacheRead(ctx, true, int64(r.Length), cacheTypeSeekable, cacheOpOpenRangeReader) timer.Success(ctx, int64(r.Length)) - decompressed, err := newDecompressingReadCloser(f, frameTable.CompressionType()) + dec, err := NewDecompressingReader(NewRangeReader(f), ct) if err != nil { f.Close() return nil, fmt.Errorf("decompress cached frame: %w", err) } - return withNFSGauge(ctx, decompressed), nil + return withNFSGauge(ctx, dec), nil case statErr == nil: // Confirmed size mismatch: drop the file so the miss path rewrites it. f.Close() @@ -70,111 +69,46 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64 recordCacheRead(ctx, false, int64(r.Length), cacheTypeSeekable, cacheOpOpenRangeReader) - rc, err := newDecompressingCacheReader(raw, frameTable.CompressionType(), r.Length, c, ctx, path, offsetU) - if err != nil { - raw.Close() - - return nil, fmt.Errorf("create decompressor: %w", err) + in := raw + if !skipCacheWriteback(ctx) { + in = newCaptureReader(raw, r.Length, true, + c.compressedFrameWriteback(path, offsetU, r.Length)) } - return rc, nil -} - -// newDecompressingCacheReader creates a reader that decompresses on Read and -// writes the accumulated compressed bytes to the NFS cache on Close. -func newDecompressingCacheReader( - raw io.ReadCloser, - ct CompressionType, - expectedSize int, - cache *cachedSeekable, - ctx context.Context, //nolint:revive // ctx after other params for readability at call site - framePath string, - offset int64, -) (io.ReadCloser, error) { - var compressedBuf bytes.Buffer - compressedBuf.Grow(expectedSize) - - tee := io.TeeReader(raw, &compressedBuf) - - dec, err := NewDecompressingReader(tee, ct) + dec, err := NewDecompressingReader(in, ct) if err != nil { - return nil, err - } + in.Close(ctx) - return &decompressingCacheReader{ - decompressor: dec, - raw: raw, - compressedBuf: &compressedBuf, - expectedSize: expectedSize, - cache: cache, - ctx: ctx, - framePath: framePath, - offset: offset, - }, nil -} - -type decompressingCacheReader struct { - decompressor io.ReadCloser // decompresses on Read - raw io.ReadCloser // underlying compressed stream (must be closed) - compressedBuf *bytes.Buffer - expectedSize int - cache *cachedSeekable - ctx context.Context //nolint:containedctx // needed for async cache write-back in Close - framePath string - offset int64 -} + return nil, fmt.Errorf("create decompressor: %w", err) + } -func (r *decompressingCacheReader) Read(p []byte) (int, error) { - return r.decompressor.Read(p) + return dec, nil } -func (r *decompressingCacheReader) Close() error { - // Drive the decompressor to EOF before closing it. With io.ReadFull bounded - // by the uncompressed size, an LZ4 frame written with BlockChecksum=true / - // Checksum=false leaves the 4-byte EndMark unread — the next Read on the - // decoder pulls the EndMark (block-size = 0 → io.EOF) from raw through the - // tee, populating compressedBuf with the full encoded frame for cache writeback. - _, _ = io.Copy(io.Discard, r.decompressor) - - decErr := r.decompressor.Close() - rawErr := r.raw.Close() - - if decErr != nil { - return decErr - } - if rawErr != nil { - return rawErr - } - - got := r.compressedBuf.Len() - if skipCacheWriteback(r.ctx) { - return nil - } +// compressedFrameWriteback returns a captureReader callback that +// persists the captured frame to the NFS cache in a detached goroutine. +// Best-effort: a short capture is logged and skipped — the caller already +// got valid decompressed bytes. +func (c *cachedSeekable) compressedFrameWriteback(framePath string, offset int64, expectedSize int) func(context.Context, []byte) { + return func(ctx context.Context, frame []byte) { + if !isCompleteRead(len(frame), expectedSize, nil) { + recordCacheWriteError(ctx, cacheTypeSeekable, cacheOpOpenRangeReader, + fmt.Errorf("compressed frame cache writeback short: got %d bytes, expected %d for %s", len(frame), expectedSize, framePath)) + + return + } - // Cache writeback is best-effort. After draining above, a remaining shortfall - // implies upstream truncation — log/metric and skip writeback rather than - // poison the read (the caller already received valid decompressed bytes). - if !isCompleteRead(got, r.expectedSize, nil) { - recordCacheWriteError(r.ctx, cacheTypeSeekable, cacheOpOpenRangeReader, - fmt.Errorf("compressed frame cache writeback short: got %d bytes, expected %d for %s", got, r.expectedSize, r.framePath)) + c.goCtx(ctx, func(ctx context.Context) { + ctx, span := c.tracer.Start(ctx, "write compressed frame back to cache") + defer span.End() - return nil + err := c.writeToCache(ctx, offset, framePath, frame) + if err != nil { + recordError(span, err) + recordCacheWriteError(ctx, cacheTypeSeekable, cacheOpOpenRangeReader, err) + } + }) } - - data := r.compressedBuf.Bytes() - r.compressedBuf = nil - - r.cache.goCtx(r.ctx, func(ctx context.Context) { - ctx, span := r.cache.tracer.Start(ctx, "write compressed frame back to cache") - defer span.End() - - if err := r.cache.writeToCache(ctx, r.offset, r.framePath, data); err != nil { - recordError(span, err) - recordCacheWriteError(ctx, cacheTypeSeekable, cacheOpOpenRangeReader, err) - } - }) - - return nil } // makeFrameFilename returns the NFS cache path for a compressed frame. diff --git a/packages/shared/pkg/storage/storage_cache_seekable_test.go b/packages/shared/pkg/storage/storage_cache_seekable_test.go index d989eb084f..c242ce6ed5 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable_test.go +++ b/packages/shared/pkg/storage/storage_cache_seekable_test.go @@ -14,6 +14,17 @@ import ( "github.com/stretchr/testify/require" ) +// mustClose closes a RangeReader and asserts no error. +func mustClose(t *testing.T, rc RangeReader) { + t.Helper() + require.NoError(t, rc.Close(t.Context())) +} + +// bytesRangeReader wraps an in-memory byte slice as a RangeReader for tests. +func bytesRangeReader(b []byte) RangeReader { + return NewRangeReader(io.NopCloser(bytes.NewReader(b))) +} + // testReadAt emulates the removed cachedSeekable.ReadAt via OpenRangeReader. // This preserves the base test structure after ReadAt was removed from the Seekable interface. func testReadAt(ctx context.Context, c *cachedSeekable, buff []byte, off int64) (int, error) { @@ -24,7 +35,7 @@ func testReadAt(ctx context.Context, c *cachedSeekable, buff []byte, off int64) n, err := io.ReadFull(rc, buff) - closeErr := rc.Close() + closeErr := rc.Close(ctx) if errors.Is(err, io.ErrUnexpectedEOF) { err = io.EOF } @@ -181,10 +192,10 @@ func TestCachedFileObjectProvider_WriteTo(t *testing.T) { inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - RunAndReturn(func(_ context.Context, off int64, length int64, _ *FrameTable) (io.ReadCloser, error) { + RunAndReturn(func(_ context.Context, off int64, length int64, _ *FrameTable) (RangeReader, error) { end := min(int(off)+int(length), len(fakeData)) - return io.NopCloser(bytes.NewReader(fakeData[off:end])), nil + return NewRangeReader(io.NopCloser(bytes.NewReader(fakeData[off:end]))), nil }) tempDir := t.TempDir() @@ -316,7 +327,7 @@ func TestCachedSeekableObjectProvider_ReadAt(t *testing.T) { inner := NewMockSeekable(t) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - Return(io.NopCloser(bytes.NewReader(nil)), nil) + Return(NewRangeReader(io.NopCloser(bytes.NewReader(nil))), nil) c := cachedSeekable{ path: tempDir, @@ -345,7 +356,7 @@ func TestCachedSeekableObjectProvider_ReadAt(t *testing.T) { inner := NewMockSeekable(t) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - Return(io.NopCloser(bytes.NewReader(data)), nil) + Return(NewRangeReader(io.NopCloser(bytes.NewReader(data))), nil) c := cachedSeekable{ path: tempDir, @@ -407,7 +418,7 @@ func TestCachedSeekable_ReadAt_PreservesEOF(t *testing.T) { inner := NewMockSeekable(t) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - Return(io.NopCloser(bytes.NewReader([]byte{1, 2, 3})), nil) + Return(NewRangeReader(io.NopCloser(bytes.NewReader([]byte{1, 2, 3}))), nil) c := cachedSeekable{ path: tempDir, @@ -431,7 +442,7 @@ func TestCachedSeekable_ReadAt_PreservesEOF(t *testing.T) { inner := NewMockSeekable(t) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - Return(io.NopCloser(bytes.NewReader([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})), nil) + Return(NewRangeReader(io.NopCloser(bytes.NewReader([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}))), nil) c := cachedSeekable{ path: tempDir, @@ -457,8 +468,8 @@ func TestCachedSeekable_ReadAt_SkipCacheWriteback(t *testing.T) { inner := NewMockSeekable(t) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - RunAndReturn(func(_ context.Context, _ int64, _ int64, _ *FrameTable) (io.ReadCloser, error) { - return io.NopCloser(bytes.NewReader(data)), nil + RunAndReturn(func(_ context.Context, _ int64, _ int64, _ *FrameTable) (RangeReader, error) { + return NewRangeReader(io.NopCloser(bytes.NewReader(data))), nil }) c := cachedSeekable{ @@ -493,7 +504,7 @@ func TestCachedSeekable_OpenRangeReader(t *testing.T) { inner := NewMockSeekable(t) inner.EXPECT(). OpenRangeReader(mock.Anything, int64(0), int64(len(data)), (*FrameTable)(nil)). - Return(io.NopCloser(bytes.NewReader(data)), nil). + Return(NewRangeReader(io.NopCloser(bytes.NewReader(data))), nil). Once() c := cachedSeekable{ @@ -510,7 +521,7 @@ func TestCachedSeekable_OpenRangeReader(t *testing.T) { got, err := io.ReadAll(rc) require.NoError(t, err) assert.Equal(t, data, got) - require.NoError(t, rc.Close()) + require.NoError(t, rc.Close(t.Context())) c.wg.Wait() @@ -522,7 +533,7 @@ func TestCachedSeekable_OpenRangeReader(t *testing.T) { got2, err := io.ReadAll(rc2) require.NoError(t, err) assert.Equal(t, data, got2) - require.NoError(t, rc2.Close()) + require.NoError(t, rc2.Close(t.Context())) }) t.Run("skip cache writeback returns inner directly", func(t *testing.T) { @@ -534,8 +545,8 @@ func TestCachedSeekable_OpenRangeReader(t *testing.T) { inner := NewMockSeekable(t) inner.EXPECT(). OpenRangeReader(mock.Anything, int64(0), int64(len(data)), (*FrameTable)(nil)). - RunAndReturn(func(_ context.Context, _ int64, _ int64, _ *FrameTable) (io.ReadCloser, error) { - return io.NopCloser(bytes.NewReader(data)), nil + RunAndReturn(func(_ context.Context, _ int64, _ int64, _ *FrameTable) (RangeReader, error) { + return NewRangeReader(io.NopCloser(bytes.NewReader(data))), nil }). Times(2) @@ -554,7 +565,7 @@ func TestCachedSeekable_OpenRangeReader(t *testing.T) { got, err := io.ReadAll(rc) require.NoError(t, err) assert.Equal(t, data, got) - require.NoError(t, rc.Close()) + require.NoError(t, rc.Close(ctx)) c.wg.Wait() @@ -569,7 +580,7 @@ func TestCachedSeekable_OpenRangeReader(t *testing.T) { got2, err := io.ReadAll(rc2) require.NoError(t, err) assert.Equal(t, data, got2) - require.NoError(t, rc2.Close()) + require.NoError(t, rc2.Close(ctx)) }) t.Run("truncated inner read does not populate cache", func(t *testing.T) { @@ -580,7 +591,7 @@ func TestCachedSeekable_OpenRangeReader(t *testing.T) { inner := NewMockSeekable(t) inner.EXPECT(). OpenRangeReader(mock.Anything, int64(0), int64(5), (*FrameTable)(nil)). - Return(io.NopCloser(bytes.NewReader([]byte{0xAA, 0xBB})), nil) + Return(NewRangeReader(io.NopCloser(bytes.NewReader([]byte{0xAA, 0xBB}))), nil) c := cachedSeekable{ path: tempDir, @@ -595,7 +606,7 @@ func TestCachedSeekable_OpenRangeReader(t *testing.T) { got, err := io.ReadAll(rc) require.NoError(t, err) assert.Equal(t, []byte{0xAA, 0xBB}, got) - require.NoError(t, rc.Close()) + require.NoError(t, rc.Close(t.Context())) c.wg.Wait() @@ -680,23 +691,16 @@ func TestCacheWriteThroughReader(t *testing.T) { c := newTestCache(t) data := []byte("hello") - inner := io.NopCloser(bytes.NewReader(data)) - - r := &cacheWriteThroughReader{ - inner: inner, - buf: bytes.NewBuffer(make([]byte, 0, len(data))), - cache: &c, - ctx: t.Context(), - off: 0, - expectedLen: int64(len(data)), - chunkPath: c.makeChunkFilename(0), - } + inner := NewRangeReader(io.NopCloser(bytes.NewReader(data))) + + r := newCaptureReader(inner, len(data), false, + c.uncompressedChunkWriteback(c.makeChunkFilename(0), 0, int64(len(data)))) got, err := io.ReadAll(r) require.NoError(t, err) assert.Equal(t, data, got) - require.NoError(t, r.Close()) + require.NoError(t, r.Close(t.Context())) c.wg.Wait() cached, err := os.ReadFile(c.makeChunkFilename(0)) @@ -711,23 +715,16 @@ func TestCacheWriteThroughReader(t *testing.T) { // Inner has only 2 bytes but expectedLen is 5. The reader is // fully consumed (EOF is reached), yet the total doesn't match // the expected length so it must not be cached. - inner := io.NopCloser(bytes.NewReader([]byte{0xAA, 0xBB})) - - r := &cacheWriteThroughReader{ - inner: inner, - buf: bytes.NewBuffer(make([]byte, 0, 5)), - cache: &c, - ctx: t.Context(), - off: 0, - expectedLen: 5, - chunkPath: c.makeChunkFilename(0), - } + inner := NewRangeReader(io.NopCloser(bytes.NewReader([]byte{0xAA, 0xBB}))) + + r := newCaptureReader(inner, 5, false, + c.uncompressedChunkWriteback(c.makeChunkFilename(0), 0, 5)) got, err := io.ReadAll(r) require.NoError(t, err) assert.Equal(t, []byte{0xAA, 0xBB}, got) - require.NoError(t, r.Close()) + require.NoError(t, r.Close(t.Context())) c.wg.Wait() _, err = os.Stat(c.makeChunkFilename(0)) @@ -739,17 +736,10 @@ func TestCacheWriteThroughReader(t *testing.T) { c := newTestCache(t) data := []byte("hello") - inner := io.NopCloser(bytes.NewReader(data)) - - r := &cacheWriteThroughReader{ - inner: inner, - buf: bytes.NewBuffer(make([]byte, 0, len(data))), - cache: &c, - ctx: t.Context(), - off: 0, - expectedLen: int64(len(data)), - chunkPath: c.makeChunkFilename(0), - } + inner := NewRangeReader(io.NopCloser(bytes.NewReader(data))) + + r := newCaptureReader(inner, len(data), false, + c.uncompressedChunkWriteback(c.makeChunkFilename(0), 0, int64(len(data)))) // Read only 2 of 5 bytes, then close without reaching EOF. buf := make([]byte, 2) @@ -757,7 +747,7 @@ func TestCacheWriteThroughReader(t *testing.T) { require.NoError(t, err) assert.Equal(t, 2, n) - require.NoError(t, r.Close()) + require.NoError(t, r.Close(t.Context())) c.wg.Wait() _, err = os.Stat(c.makeChunkFilename(0)) diff --git a/packages/shared/pkg/storage/storage_fs.go b/packages/shared/pkg/storage/storage_fs.go index 25b08f7b10..52e6c66a6c 100644 --- a/packages/shared/pkg/storage/storage_fs.go +++ b/packages/shared/pkg/storage/storage_fs.go @@ -39,16 +39,6 @@ var ( _ StreamingReader = (*fsObject)(nil) ) -type fsRangeReadCloser struct { - io.Reader - - file *os.File -} - -func (r *fsRangeReadCloser) Close() error { - return r.file.Close() -} - func newFileSystemStorage(cfg StorageConfig) *fsStorage { return &fsStorage{ basePath: cfg.GetLocalBasePath(), @@ -205,16 +195,13 @@ func (o *fsObject) storeFileCompressed(ctx context.Context, localPath string, cf return ft, checksum, nil } -func (o *fsObject) openRangeReader(_ context.Context, off, length int64) (io.ReadCloser, error) { +func (o *fsObject) openRangeReader(_ context.Context, off, length int64) (RangeReader, error) { f, err := o.getHandle(true) if err != nil { return nil, err } - return &fsRangeReadCloser{ - Reader: io.NewSectionReader(f, off, length), - file: f, - }, nil + return newSectionReader(f, off, length), nil } func (o *fsObject) Exists(_ context.Context) (bool, error) { @@ -315,7 +302,7 @@ func (u *fsPartUploader) Complete(_ context.Context) error { return os.WriteFile(u.fullPath, u.Assemble(), 0o644) } -func (o *fsObject) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (io.ReadCloser, error) { +func (o *fsObject) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, error) { if frameTable.IsCompressed() { r, err := frameTable.LocateCompressed(offsetU) if err != nil { @@ -327,14 +314,14 @@ func (o *fsObject) OpenRangeReader(ctx context.Context, offsetU int64, length in return nil, err } - decompressed, err := newDecompressingReadCloser(raw, frameTable.CompressionType()) + dec, err := NewDecompressingReader(raw, frameTable.CompressionType()) if err != nil { - raw.Close() + raw.Close(ctx) return nil, err } - return decompressed, nil + return dec, nil } return o.openRangeReader(ctx, offsetU, length) diff --git a/packages/shared/pkg/storage/storage_google.go b/packages/shared/pkg/storage/storage_google.go index 0556cecb73..eab8e60226 100644 --- a/packages/shared/pkg/storage/storage_google.go +++ b/packages/shared/pkg/storage/storage_google.go @@ -285,20 +285,25 @@ func (o *gcpObject) openRangeReader(ctx context.Context, off, length int64) (io. return nil, fmt.Errorf("failed to create GCS range reader for %q at %d+%d: %w", o.path, off, length, err) } + gcsConcurrentReads.Add(ctx, 1) + return &idleTimeoutReader{ ReadCloser: reader, cancel: cancel, timer: time.AfterFunc(googleReadTimeout, cancel), + gaugeCtx: ctx, }, nil } // idleTimeoutReader fires cancel() after googleReadTimeout with no Read -// activity (in-flight Read with no progress, or no Read called). +// activity (in-flight Read with no progress, or no Read called). It also +// pairs with gcsConcurrentReads: +1 on construction, -1 on Close. type idleTimeoutReader struct { io.ReadCloser - cancel context.CancelFunc - timer *time.Timer + cancel context.CancelFunc + timer *time.Timer + gaugeCtx context.Context //nolint:containedctx // needed to decrement gcsConcurrentReads on Close } func (r *idleTimeoutReader) Read(p []byte) (int, error) { @@ -317,6 +322,7 @@ func (r *idleTimeoutReader) Read(p []byte) (int, error) { func (r *idleTimeoutReader) Close() error { r.timer.Stop() defer r.cancel() + gcsConcurrentReads.Add(r.gaugeCtx, -1) return r.ReadCloser.Close() } @@ -598,7 +604,7 @@ func parseServiceAccountBase64(serviceAccount string) (*gcpServiceToken, error) return &sa, nil } -func (o *gcpObject) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (io.ReadCloser, error) { +func (o *gcpObject) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, error) { timer := googleReadTimerFactory.Begin(attribute.String(gcsOperationAttr, gcsOperationAttrReadAt)) if !frameTable.IsCompressed() { @@ -609,9 +615,7 @@ func (o *gcpObject) OpenRangeReader(ctx context.Context, offsetU int64, length i return nil, err } - gcsConcurrentReads.Add(ctx, 1) - - return &timedReadCloser{inner: rc, timer: timer, ctx: ctx}, nil + return newObservableReader(NewRangeReader(rc), timer, nil), nil } r, err := frameTable.LocateCompressed(offsetU) @@ -628,7 +632,7 @@ func (o *gcpObject) OpenRangeReader(ctx context.Context, offsetU int64, length i return nil, err } - decompressed, err := newDecompressingReadCloser(raw, frameTable.CompressionType()) + dec, err := NewDecompressingReader(NewRangeReader(raw), frameTable.CompressionType()) if err != nil { raw.Close() timer.Failure(ctx, 0) @@ -636,44 +640,7 @@ func (o *gcpObject) OpenRangeReader(ctx context.Context, offsetU int64, length i return nil, err } - gcsConcurrentReads.Add(ctx, 1) - - return &timedReadCloser{inner: decompressed, timer: timer, ctx: ctx}, nil -} - -// timedReadCloser wraps a reader with OTEL timer metrics. -// Close records success (with total bytes read) or failure on the timer. -type timedReadCloser struct { - inner io.ReadCloser - timer *telemetry.Stopwatch - ctx context.Context //nolint:containedctx // needed for timer recording in Close - bytesRead int64 - closeErr error -} - -func (r *timedReadCloser) Read(p []byte) (int, error) { - n, err := r.inner.Read(p) - r.bytesRead += int64(n) - - if err != nil && err != io.EOF { - r.closeErr = err - } - - return n, err -} - -func (r *timedReadCloser) Close() error { - gcsConcurrentReads.Add(r.ctx, -1) - - err := r.inner.Close() - - if r.closeErr != nil || err != nil { - r.timer.Failure(r.ctx, r.bytesRead) - } else { - r.timer.Success(r.ctx, r.bytesRead) - } - - return err + return newObservableReader(dec, timer, nil), nil } func isResourceExhausted(err error) bool { 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..6da9cb9b90 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 @@ -149,7 +149,7 @@ func verifyChecksum(t *testing.T, ctx context.Context, persistence storage.Stora rc, err := obj.OpenRangeReader(ctx, 0, bd.Size, bd.FrameData) require.NoErrorf(t, err, "%s/%s: open range reader", node.name, fileName) - defer rc.Close() + defer rc.Close(ctx) hasher := sha256.New() n, err := io.Copy(hasher, rc) From 082285cddf83ebd56fae4fd936441366fe11a414 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Mon, 8 Jun 2026 14:05:30 -0700 Subject: [PATCH 02/11] Rid of OffsetReader --- packages/shared/pkg/storage/io_wrappers.go | 19 --- .../shared/pkg/storage/io_wrappers_test.go | 128 ------------------ .../pkg/storage/storage_cache_seekable.go | 5 +- 3 files changed, 2 insertions(+), 150 deletions(-) delete mode 100644 packages/shared/pkg/storage/io_wrappers_test.go diff --git a/packages/shared/pkg/storage/io_wrappers.go b/packages/shared/pkg/storage/io_wrappers.go index 1e79763b89..41be954845 100644 --- a/packages/shared/pkg/storage/io_wrappers.go +++ b/packages/shared/pkg/storage/io_wrappers.go @@ -13,31 +13,12 @@ import ( ) var ( - _ io.Reader = (*offsetReader)(nil) _ RangeReader = (*sectionReader)(nil) _ RangeReader = (*observableReader)(nil) _ RangeReader = (*rangeReader)(nil) _ RangeReader = (*captureReader)(nil) ) -// offsetReader adapts an io.ReaderAt into a sequential io.Reader -// starting at the given offset. -type offsetReader struct { - wrapped io.ReaderAt - offset int64 -} - -func (r *offsetReader) Read(p []byte) (n int, err error) { - n, err = r.wrapped.ReadAt(p, r.offset) - r.offset += int64(n) - - return -} - -func newOffsetReader(reader io.ReaderAt, offset int64) *offsetReader { - return &offsetReader{reader, offset} -} - // rangeReader adapts an io.ReadCloser into a RangeReader by ignoring the // Close context. type rangeReader struct { diff --git a/packages/shared/pkg/storage/io_wrappers_test.go b/packages/shared/pkg/storage/io_wrappers_test.go deleted file mode 100644 index 1c59c46cbc..0000000000 --- a/packages/shared/pkg/storage/io_wrappers_test.go +++ /dev/null @@ -1,128 +0,0 @@ -package storage - -import ( - "bytes" - "io" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestOffsetReader_Read(t *testing.T) { - t.Parallel() - - data := []byte("hello world") - readerAt := bytes.NewReader(data) - - tests := []struct { - name string - offset int64 - readSize int - expectedData string - expectedN int - expectedErr error - expectedOffset int64 - }{ - { - name: "read from start", - offset: 0, - readSize: 5, - expectedData: "hello", - expectedN: 5, - expectedErr: nil, - expectedOffset: 5, - }, - { - name: "read from offset", - offset: 6, - readSize: 5, - expectedData: "world", - expectedN: 5, - expectedErr: nil, - expectedOffset: 11, - }, - { - name: "read until EOF", - offset: 0, - readSize: 11, - expectedData: "hello world", - expectedN: 11, - expectedErr: nil, - expectedOffset: 11, - }, - { - name: "read past EOF", - offset: 0, - readSize: 15, - expectedData: "hello world", - expectedN: 11, - expectedErr: io.EOF, - expectedOffset: 11, - }, - { - name: "read exactly at EOF", - offset: 11, - readSize: 5, - expectedData: "", - expectedN: 0, - expectedErr: io.EOF, - expectedOffset: 11, - }, - { - name: "read zero bytes", - offset: 0, - readSize: 0, - expectedData: "", - expectedN: 0, - expectedErr: nil, - expectedOffset: 0, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - r := newOffsetReader(readerAt, tt.offset) - p := make([]byte, tt.readSize) - n, err := r.Read(p) - - require.ErrorIs(t, err, tt.expectedErr) - assert.Equal(t, tt.expectedN, n) - assert.Equal(t, tt.expectedData, string(p[:n])) - assert.Equal(t, tt.expectedOffset, r.offset) - }) - } -} - -func TestOffsetReader_SequentialReads(t *testing.T) { - t.Parallel() - - data := []byte("hello world") - readerAt := bytes.NewReader(data) - r := newOffsetReader(readerAt, 0) - - // First read - p1 := make([]byte, 6) - n1, err1 := r.Read(p1) - require.NoError(t, err1) - assert.Equal(t, 6, n1) - assert.Equal(t, "hello ", string(p1[:n1])) - assert.Equal(t, int64(6), r.offset) - - // Second read - p2 := make([]byte, 5) - n2, err2 := r.Read(p2) - require.NoError(t, err2) - assert.Equal(t, 5, n2) - assert.Equal(t, "world", string(p2[:n2])) - assert.Equal(t, int64(11), r.offset) - - // Third read (EOF) - p3 := make([]byte, 5) - n3, err3 := r.Read(p3) - require.ErrorIs(t, err3, io.EOF) - assert.Equal(t, 0, n3) - assert.Equal(t, int64(11), r.offset) -} diff --git a/packages/shared/pkg/storage/storage_cache_seekable.go b/packages/shared/pkg/storage/storage_cache_seekable.go index 86835bceb7..ae3ab1aa30 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable.go +++ b/packages/shared/pkg/storage/storage_cache_seekable.go @@ -501,9 +501,8 @@ func (c *cachedSeekable) writeChunkFromFile(ctx context.Context, offset int64, i } defer utils.Cleanup(ctx, "failed to close file", output.Close) - offsetReader := newOffsetReader(input, offset) - count, err := io.CopyN(output, offsetReader, c.chunkSize) - if ignoreEOF(err) != nil { + count, err := io.Copy(output, io.NewSectionReader(input, offset, c.chunkSize)) + if err != nil { writeTimer.Failure(ctx, count) safelyRemoveFile(ctx, chunkPath) From 6d9b21d691ef4dfc6c3de552defac18d949eb14d Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Mon, 8 Jun 2026 15:06:03 -0700 Subject: [PATCH 03/11] restored .vscode/settings.json from main --- .vscode/settings.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index f1b96e2a42..42900fa357 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -52,7 +52,7 @@ }, }, "editor.tabCompletion": "on", - "go.lintTool": "golangci-lint-v2", + "go.lintTool": "golangci-lint", "go.lintFlags": [ "--path-mode=abs", "--fast-only", From 5787a00ad07d61f6e6cb9411fdf3b948d29409e3 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Tue, 9 Jun 2026 10:45:30 -0700 Subject: [PATCH 04/11] PR feedback --- packages/shared/pkg/storage/io_wrappers.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/shared/pkg/storage/io_wrappers.go b/packages/shared/pkg/storage/io_wrappers.go index 41be954845..18d5e5ebba 100644 --- a/packages/shared/pkg/storage/io_wrappers.go +++ b/packages/shared/pkg/storage/io_wrappers.go @@ -84,7 +84,9 @@ func (r *captureReader) Close(ctx context.Context) error { _, _ = io.Copy(io.Discard, r) } err := r.inner.Close(ctx) - r.onClose(ctx, r.buf.Bytes()) + if err == nil { + r.onClose(ctx, r.buf.Bytes()) + } return err } From a53795197a720840a6a0b9f2c38352048deddc49 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Thu, 11 Jun 2026 15:46:54 -0700 Subject: [PATCH 05/11] cleanup post main merge --- .../pkg/sandbox/block/streaming_chunk_test.go | 12 ++++++------ packages/shared/pkg/storage/storage_google.go | 18 ++++++++---------- 2 files changed, 14 insertions(+), 16 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go index dac1211783..8295daaa46 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go @@ -97,13 +97,13 @@ func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length end := min(offsetU+length, int64(len(s.data))) - return storage.NewRangeReader(&controlledReader{ + return &controlledReader{ data: s.data[offsetU:end], step: max(16*1024, testBlockSize), advance: s.ctrl.advance, consumed: s.ctrl.consumed, closed: s.ctrl.closed, - }), nil + }, nil } var fetchOff, fetchLen int64 @@ -432,10 +432,10 @@ func (s *panicSeekable) StoreFile(context.Context, string, ...storage.PutOption) func (s *panicSeekable) OpenRangeReader(_ context.Context, off int64, length int64, _ *storage.FrameTable) (storage.RangeReader, error) { end := min(off+length, int64(len(s.data))) - return storage.NewRangeReader(&panicReader{ + return &panicReader{ data: s.data[off:end], panicAfter: int(s.panicAfter - off), - }), nil + }, nil } type panicReader struct { @@ -460,7 +460,7 @@ func (r *panicReader) Read(p []byte) (int, error) { return n, nil } -func (r *panicReader) Close() error { +func (r *panicReader) Close(context.Context) error { return nil } @@ -587,7 +587,7 @@ func (r *controlledReader) Read(p []byte) (int, error) { return n, nil } -func (r *controlledReader) Close() error { +func (r *controlledReader) Close(context.Context) error { select { case r.closed <- struct{}{}: default: diff --git a/packages/shared/pkg/storage/storage_google.go b/packages/shared/pkg/storage/storage_google.go index 779a1bb7cb..486be15b05 100644 --- a/packages/shared/pkg/storage/storage_google.go +++ b/packages/shared/pkg/storage/storage_google.go @@ -264,7 +264,7 @@ func (o *gcpObject) Size(ctx context.Context) (int64, error) { return attrs.Size, nil } -func (o *gcpObject) openRangeReader(ctx context.Context, off, length int64) (io.ReadCloser, error) { +func (o *gcpObject) openRangeReader(ctx context.Context, off, length int64) (RangeReader, error) { readCtx, cancel := context.WithCancel(ctx) openTimer := time.AfterFunc(googleReadTimeout, cancel) @@ -291,7 +291,6 @@ func (o *gcpObject) openRangeReader(ctx context.Context, off, length int64) (io. ReadCloser: reader, cancel: cancel, timer: time.AfterFunc(googleReadTimeout, cancel), - gaugeCtx: ctx, }, nil } @@ -301,9 +300,8 @@ func (o *gcpObject) openRangeReader(ctx context.Context, off, length int64) (io. type idleTimeoutReader struct { io.ReadCloser - cancel context.CancelFunc - timer *time.Timer - gaugeCtx context.Context //nolint:containedctx // needed to decrement gcsConcurrentReads on Close + cancel context.CancelFunc + timer *time.Timer } func (r *idleTimeoutReader) Read(p []byte) (int, error) { @@ -319,10 +317,10 @@ func (r *idleTimeoutReader) Read(p []byte) (int, error) { return n, err } -func (r *idleTimeoutReader) Close() error { +func (r *idleTimeoutReader) Close(ctx context.Context) error { r.timer.Stop() defer r.cancel() - gcsConcurrentReads.Add(r.gaugeCtx, -1) + gcsConcurrentReads.Add(ctx, -1) return r.ReadCloser.Close() } @@ -616,7 +614,7 @@ func (o *gcpObject) OpenRangeReader(ctx context.Context, offsetU int64, length i return nil, err } - return newObservableReader(NewRangeReader(rc), timer, nil), nil + return newObservableReader(rc, timer, nil), nil } r, err := frameTable.LocateCompressed(offsetU) @@ -633,9 +631,9 @@ func (o *gcpObject) OpenRangeReader(ctx context.Context, offsetU int64, length i return nil, err } - dec, err := NewDecompressingReader(NewRangeReader(raw), frameTable.CompressionType()) + dec, err := NewDecompressingReader(raw, frameTable.CompressionType()) if err != nil { - raw.Close() + raw.Close(ctx) timer.Failure(ctx, 0) return nil, err From 2c94bd7fa6e306587fcb5b89e50cf96903b17642 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Mon, 15 Jun 2026 16:12:45 -0700 Subject: [PATCH 06/11] PR feedback: avoid writeback after decompressor creation fails When NewDecompressingReader fails, closing the captureReader drained the raw stream into its buffer and persisted those bytes to NFS, poisoning the .frm cache on a failed miss. Close raw directly on the error path to bypass the drain+writeback side effect. --- .../shared/pkg/storage/storage_cache_seekable_compressed.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/shared/pkg/storage/storage_cache_seekable_compressed.go b/packages/shared/pkg/storage/storage_cache_seekable_compressed.go index d2e2996776..2e8077c4bb 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable_compressed.go +++ b/packages/shared/pkg/storage/storage_cache_seekable_compressed.go @@ -77,7 +77,7 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64 dec, err := NewDecompressingReader(in, ct) if err != nil { - in.Close(ctx) + raw.Close(ctx) return nil, fmt.Errorf("create decompressor: %w", err) } From e4409392e86caad7cea550606b43b8f7b4ddba99 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Mon, 15 Jun 2026 16:39:05 -0700 Subject: [PATCH 07/11] feat(storage): observability for the read path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an orchestrator.read.* metric family with consistent attributes (file_type/source/codec/outcome) covering each stage of a read, plus per-layer chunker and build-file timers, so dashboards can attribute latency end-to-end from sandbox-visible read to backend fetch. Does not remove any of the prior metrics, this will be done separately after the dashboards are updated. Metrics: - orchestrator.file.read_at build.File.ReadAt — per-fault unit, aggregates all underlying mappings into one record - orchestrator.chunk.slice Chunker.Slice — per per-mapping unit, source=mmap on cache hit else the backend that served - orchestrator.read.open OpenRangeReader (open / TTFB) - orchestrator.read.read source-read wall, compressed bytes - orchestrator.read.decompress decompress CPU + uncompressed bytes - orchestrator.read.fetch total fetch wall + bytes delivered - orchestrator.read.writeback NFS cache writeback wall + bytes - orchestrator.read.pipeline.efficiency fetch / (open+read+decompress) - orchestrator.read.cache NFS hit/miss/writeback events - orchestrator.read.inflight concurrent fetches gauge Spans: - chunk.fetch runFetch goroutine span --- .../cmd/inspect-build/validate.go | 2 +- .../pkg/sandbox/block/fetch_session.go | 13 + .../pkg/sandbox/block/metrics/main.go | 14 +- .../pkg/sandbox/block/streaming_chunk.go | 92 +++++- .../pkg/sandbox/block/streaming_chunk_test.go | 26 +- .../orchestrator/pkg/sandbox/build/build.go | 40 ++- .../pkg/sandbox/build/storage_diff.go | 2 +- .../pkg/sandbox/build/storage_diff_test.go | 26 +- .../sandbox/template/peerclient/seekable.go | 8 +- .../template/peerclient/seekable_test.go | 15 +- .../sandbox/template/peerclient/storage.go | 23 +- .../shared/pkg/storage/compress_decode.go | 61 +++- packages/shared/pkg/storage/io_wrappers.go | 83 ++++-- packages/shared/pkg/storage/mock_seekable.go | 24 +- packages/shared/pkg/storage/read_attrs.go | 26 ++ .../pkg/storage/read_attrs_precomputed.go | 145 ++++++++++ .../shared/pkg/storage/read_attrs_test.go | 57 ++++ packages/shared/pkg/storage/read_metrics.go | 151 ++++++++++ packages/shared/pkg/storage/storage.go | 34 ++- packages/shared/pkg/storage/storage_aws.go | 10 +- packages/shared/pkg/storage/storage_cache.go | 1 + .../storage/storage_cache_compressed_test.go | 16 +- .../pkg/storage/storage_cache_seekable.go | 41 ++- .../storage_cache_seekable_compressed.go | 34 ++- .../storage/storage_cache_seekable_test.go | 269 +++++++----------- packages/shared/pkg/storage/storage_fs.go | 27 +- packages/shared/pkg/storage/storage_google.go | 26 +- packages/shared/pkg/telemetry/meters.go | 50 +++- .../sandbox_rapid_pause_resume_test.go | 2 +- 29 files changed, 987 insertions(+), 331 deletions(-) create mode 100644 packages/shared/pkg/storage/read_attrs.go create mode 100644 packages/shared/pkg/storage/read_attrs_precomputed.go create mode 100644 packages/shared/pkg/storage/read_attrs_test.go create mode 100644 packages/shared/pkg/storage/read_metrics.go diff --git a/packages/orchestrator/cmd/inspect-build/validate.go b/packages/orchestrator/cmd/inspect-build/validate.go index 9548c04012..f4bc0936bb 100644 --- a/packages/orchestrator/cmd/inspect-build/validate.go +++ b/packages/orchestrator/cmd/inspect-build/validate.go @@ -190,7 +190,7 @@ func openChunker(ctx context.Context, storagePath, buildID, artifact string, h * return nil, nil, 0, nil, err } - chunker, err := block.NewChunker(flags, size, int64(h.Metadata.BlockSize), filepath.Join(cacheDir, "cache"), m) + chunker, err := block.NewChunker(flags, size, int64(h.Metadata.BlockSize), filepath.Join(cacheDir, "cache"), m, seekableType(artifact)) if err != nil { os.RemoveAll(cacheDir) diff --git a/packages/orchestrator/pkg/sandbox/block/fetch_session.go b/packages/orchestrator/pkg/sandbox/block/fetch_session.go index 8b013afb01..9f80484261 100644 --- a/packages/orchestrator/pkg/sandbox/block/fetch_session.go +++ b/packages/orchestrator/pkg/sandbox/block/fetch_session.go @@ -7,6 +7,8 @@ import ( "fmt" "sync" "sync/atomic" + + "github.com/e2b-dev/infra/packages/shared/pkg/storage" ) type fetchSession struct { @@ -21,12 +23,23 @@ type fetchSession struct { fetchErr error done bool // true once terminated (success or error) + // source is the backend that served the fetch; set by runFetch. + source atomic.Int32 + // bytesReady is the byte count (from chunkOff) up to which all blocks // are fully written and marked cached. Atomic so registerAndWait can // do a lock-free fast-path check: bytesReady only increases. bytesReady atomic.Int64 } +func (s *fetchSession) setSource(src storage.Source) { + s.source.Store(int32(src)) +} + +func (s *fetchSession) Source() storage.Source { + return storage.Source(s.source.Load()) +} + // contains reports whether the session covers the byte range [off, off+length). func (s *fetchSession) contains(off, length int64) bool { return s.chunkOff <= off && s.chunkOff+s.chunkLen >= off+length diff --git a/packages/orchestrator/pkg/sandbox/block/metrics/main.go b/packages/orchestrator/pkg/sandbox/block/metrics/main.go index f1d67440da..25da8b0b71 100644 --- a/packages/orchestrator/pkg/sandbox/block/metrics/main.go +++ b/packages/orchestrator/pkg/sandbox/block/metrics/main.go @@ -12,17 +12,20 @@ const ( orchestratorBlockSlices = "orchestrator.blocks.slices" orchestratorBlockChunksFetch = "orchestrator.blocks.chunks.fetch" orchestratorBlockChunksStore = "orchestrator.blocks.chunks.store" + orchestratorChunkSlice = "orchestrator.chunk.slice" ) type Metrics struct { // SlicesMetric is used to measure page faulting performance. SlicesTimerFactory telemetry.TimerFactory - // WriteChunksMetric is used to measure the time taken to download chunks from remote storage + // RemoteReadsTimerFactory is used to measure the time taken to download chunks from remote storage RemoteReadsTimerFactory telemetry.TimerFactory // WriteChunksMetric is used to measure performance of writing chunks to disk. WriteChunksTimerFactory telemetry.TimerFactory + + ChunkSliceTimerFactory telemetry.FloatTimerFactory } func NewMetrics(meterProvider metric.MeterProvider) (Metrics, error) { @@ -58,5 +61,14 @@ func NewMetrics(meterProvider metric.MeterProvider) (Metrics, error) { return m, fmt.Errorf("failed to get stored chunks metric: %w", err) } + if m.ChunkSliceTimerFactory, err = telemetry.NewFloatTimerFactory( + blocksMeter, orchestratorChunkSlice, + "Time taken by Chunker to serve a Slice() (source=mmap when served from cache)", + "Bytes returned", + "Slice call count", + ); err != nil { + return m, fmt.Errorf("error creating chunk slice timer factory: %w", err) + } + return m, nil } diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go index f7d9213b0c..be5c64da07 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go @@ -30,6 +30,7 @@ type Chunker struct { metrics metrics.Metrics fetchTimeout time.Duration featureFlags *featureflags.Client + objType storage.SeekableObjectType size int64 @@ -42,6 +43,7 @@ func NewChunker( size, blockSize int64, cachePath string, metrics metrics.Metrics, + objType storage.SeekableObjectType, ) (*Chunker, error) { cache, err := NewCache(size, blockSize, cachePath, false) if err != nil { @@ -54,6 +56,7 @@ func NewChunker( metrics: metrics, featureFlags: ff, fetchTimeout: defaultFetchTimeout, + objType: objType, }, nil } @@ -69,42 +72,52 @@ func (c *Chunker) ReadAt(ctx context.Context, b []byte, off int64, upstream stor } func (c *Chunker) Slice(ctx context.Context, off, length int64, upstream storage.RangeOpener, ft *storage.FrameTable) ([]byte, error) { + ct := ft.CompressionType() attrs := chunkerAttrs if ft.IsCompressed() { attrs = chunkerAttrsCompressed } + + sliceStart := time.Now() timer := c.metrics.SlicesTimerFactory.Begin() - // Fast path: already cached + // Fast path: already cached. b, err := c.cache.Slice(off, length) if err == nil { timer.RecordRaw(ctx, length, attrs.successFromCache) + c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), length, storage.OKAttrs(c.objType, storage.SourceMmap, ct)) return b, nil } if !errors.As(err, &BytesNotAvailableError{}) { timer.RecordRaw(ctx, length, attrs.failCacheRead) + c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), 0, storage.ErrAttrs(c.objType, storage.SourceMmap, ct, err)) return nil, fmt.Errorf("failed read from cache at offset %d: %w", off, err) } // Fetch every chunk the range spans (one fetch session per chunk). + var src storage.Source end := off + length for cur := off; cur < end; { chunkOff, chunkLen, lerr := c.locateChunk(cur, ft) if lerr != nil { timer.RecordRaw(ctx, length, attrs.failRemoteFetch) + c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), 0, storage.ErrAttrs(c.objType, src, ct, lerr)) return nil, fmt.Errorf("failed to locate chunk for offset %d: %w", cur, lerr) } chunkEnd := chunkOff + chunkLen rangeEnd := min(end, chunkEnd) - if err := c.fetch(ctx, cur, rangeEnd-cur, upstream, ft); err != nil { + s, err := c.fetch(ctx, cur, rangeEnd-cur, upstream, ft) + if err != nil { timer.RecordRaw(ctx, length, attrs.failRemoteFetch) + c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), 0, storage.ErrAttrs(c.objType, s, ct, err)) return nil, fmt.Errorf("failed to ensure data at %d-%d: %w", cur, rangeEnd, err) } + src = max(src, s) cur = chunkEnd } @@ -112,11 +125,13 @@ func (c *Chunker) Slice(ctx context.Context, off, length int64, upstream storage b, cacheErr := c.cache.sliceDirect(off, length) if cacheErr != nil { timer.RecordRaw(ctx, length, attrs.failLocalReadAgain) + c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), 0, storage.ErrAttrs(c.objType, src, ct, cacheErr)) return nil, fmt.Errorf("failed to read from cache after ensuring data at %d-%d: %w", off, off+length, cacheErr) } timer.RecordRaw(ctx, length, attrs.successFromRemote) + c.metrics.ChunkSliceTimerFactory.Record(ctx, time.Since(sliceStart), length, storage.OKAttrs(c.objType, src, ct)) return b, nil } @@ -156,15 +171,15 @@ func (c *Chunker) getOrCreateSession(ctx context.Context, off, length int64, ups // 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, upstream storage.RangeOpener, ft *storage.FrameTable) error { +func (c *Chunker) fetch(ctx context.Context, off, length int64, upstream storage.RangeOpener, ft *storage.FrameTable) (storage.Source, error) { chunkOff, chunkLen, err := c.locateChunk(off, ft) if err != nil { - return fmt.Errorf("failed to locate chunk for offset %d: %w", off, err) + return storage.UnknownSource, fmt.Errorf("failed to locate chunk for offset %d: %w", off, err) } session, justGotCached := c.getOrCreateSession(ctx, chunkOff, chunkLen, upstream, ft) if justGotCached { - return nil + return storage.SourceMmap, nil } blockSize := c.cache.BlockSize() @@ -176,11 +191,11 @@ func (c *Chunker) fetch(ctx context.Context, off, length int64, upstream storage break // tail belongs to the caller's next chunk fetch. } if err := session.registerAndWait(ctx, b); err != nil { - return err + return session.Source(), err } } - return nil + return session.Source(), nil } // runFetch fetches data from storage into the mmap cache. Runs in a background goroutine. @@ -188,6 +203,13 @@ func (c *Chunker) runFetch(ctx context.Context, s *fetchSession, upstream storag ctx, cancel := context.WithTimeout(ctx, c.fetchTimeout) defer cancel() + ctx, span := tracer.Start(ctx, "chunk.fetch") + defer span.End() + span.SetAttributes( + attribute.Int64("off", s.chunkOff), + attribute.Int64("len", s.chunkLen), + ) + defer c.releaseSession(s) // Unconditionally terminate the session on exit so registerAndWait @@ -211,16 +233,26 @@ func (c *Chunker) runFetch(ctx context.Context, s *fetchSession, upstream storag } defer releaseLock() + ct := ft.CompressionType() + attrs := chunkerAttrs if ft.IsCompressed() { attrs = chunkerAttrsCompressed } fetchTimer := c.metrics.RemoteReadsTimerFactory.Begin() - readBytes, err := c.progressiveRead(ctx, s, mmapSlice, upstream, ft) + fetchStart := time.Now() + + stats, src, open, err := c.progressiveRead(ctx, s, mmapSlice, upstream, ft) + var readBytes int64 + if stats != nil { + readBytes = stats.UncompressedBytes + } if err != nil { fetchTimer.RecordRaw(ctx, readBytes, attrs.remoteFailure) + storage.RecordReadFetch(ctx, time.Since(fetchStart), readBytes, storage.ErrAttrs(c.objType, src, ct, err)) + s.fail(err) return @@ -231,24 +263,56 @@ func (c *Chunker) runFetch(ctx context.Context, s *fetchSession, upstream storag // closing the TOCTOU window in getOrCreateSession. c.cache.setIsCached(s.chunkOff, s.chunkLen) + fetchWall := time.Since(fetchStart) fetchTimer.RecordRaw(ctx, readBytes, attrs.remoteSuccess) + storage.RecordReadFetch(ctx, fetchWall, readBytes, storage.OKAttrs(c.objType, src, ct)) + + // fetch wall / (open + read + decompress); >1 = unaccounted overhead. + if stats != nil { + if work := open + stats.Read + stats.Decompress; work > 0 { + ratio := fetchWall.Seconds() / work.Seconds() + storage.RecordPipelineEfficiency(ctx, ratio, storage.OKAttrs(c.objType, src, ct)) + } + } + s.setDone() } -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) +func (c *Chunker) progressiveRead(ctx context.Context, s *fetchSession, mmapSlice []byte, upstream storage.RangeOpener, ft *storage.FrameTable) (stats *storage.ReadStats, source storage.Source, open time.Duration, err error) { + doneInflight := storage.StartInflight(ctx, storage.InflightFetchAttrs(c.objType)) + defer doneInflight() + + openStart := time.Now() + reader, source, err := upstream.OpenRangeReader(ctx, s.chunkOff, s.chunkLen, ft) + open = time.Since(openStart) + s.setSource(source) if err != nil { - return 0, fmt.Errorf("failed to open range reader at %d: %w", s.chunkOff, err) + return nil, source, open, fmt.Errorf("failed to open range reader at %d: %w", s.chunkOff, err) } + defer func() { - if closeErr := reader.Close(ctx); closeErr != nil && err == nil { + var closeErr error + stats, closeErr = reader.Close(ctx) + if closeErr != nil && err == nil { err = closeErr } + if err != nil { + return + } + + ct := ft.CompressionType() + okAttrs := storage.OKAttrs(c.objType, source, ct) + + storage.RecordReadOpen(ctx, open, 0, okAttrs) + if stats != nil { + storage.RecordReadRead(ctx, stats.Read, stats.CompressedBytes, okAttrs) + } }() blockSize := c.cache.BlockSize() readBatch := max(blockSize, int64(c.featureFlags.IntFlag(ctx, featureflags.MinChunkerReadSizeKB))*1024) + var totalRead int64 for totalRead < s.chunkLen { // Read in batches of max(blockSize, minReadBatchSize) to align notification // granularity with the read size and minimize lock/notify overhead. @@ -273,11 +337,11 @@ func (c *Chunker) progressiveRead(ctx context.Context, s *fetchSession, mmapSlic break // all bytes received; trailing EOF is expected } - return totalRead, fmt.Errorf("failed reading at offset %d after %d bytes: %w", s.chunkOff, totalRead, readErr) + return stats, source, open, fmt.Errorf("failed reading at offset %d after %d bytes: %w", s.chunkOff, totalRead, readErr) } } - return totalRead, nil + return stats, source, open, nil } // releaseSession removes s from the active list (swap-delete). diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go index 8295daaa46..534bf474cf 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go @@ -68,7 +68,7 @@ type testControl struct { func newTestChunker(t *testing.T, size int64) *Chunker { t.Helper() - c, err := NewChunker(&featureflags.Client{}, size, testBlockSize, t.TempDir()+"/cache", newTestMetrics(t)) + c, err := NewChunker(&featureflags.Client{}, size, testBlockSize, t.TempDir()+"/cache", newTestMetrics(t), storage.MemfileObjectType) require.NoError(t, err) return c @@ -82,7 +82,7 @@ func (s *fakeSeekable) StoreFile(context.Context, string, ...storage.PutOption) panic("not used") } -func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length int64, frameTable *storage.FrameTable) (storage.RangeReader, error) { +func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length int64, frameTable *storage.FrameTable) (storage.RangeReader, storage.Source, error) { s.fetchCount.Add(1) if s.ctrl != nil { @@ -103,14 +103,14 @@ func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length advance: s.ctrl.advance, consumed: s.ctrl.consumed, closed: s.ctrl.closed, - }, nil + }, storage.SourceFS, nil } var fetchOff, fetchLen int64 if frameTable.IsCompressed() { r, err := frameTable.LocateCompressed(offsetU) if err != nil { - return nil, fmt.Errorf("frame lookup: %w", err) + return nil, storage.UnknownSource, fmt.Errorf("frame lookup: %w", err) } fetchOff = r.Offset @@ -127,10 +127,12 @@ func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length r := io.Reader(bytes.NewReader(s.data[fetchOff:end])) if frameTable.IsCompressed() { - return storage.NewDecompressingReader(storage.NewRangeReader(io.NopCloser(r)), frameTable.CompressionType()) + dec, err := storage.NewDecompressingReader(storage.NewRangeReader(io.NopCloser(r)), frameTable.CompressionType()) + + return dec, storage.SourceFS, err } - return storage.NewRangeReader(io.NopCloser(r)), nil + return storage.NewRangeReader(io.NopCloser(r)), storage.SourceFS, nil } func makeCompressedTestData(tb testing.TB, data []byte) (*storage.FrameTable, *fakeSeekable) { @@ -429,13 +431,13 @@ func (s *panicSeekable) StoreFile(context.Context, string, ...storage.PutOption) panic("not used") } -func (s *panicSeekable) OpenRangeReader(_ context.Context, off int64, length int64, _ *storage.FrameTable) (storage.RangeReader, error) { +func (s *panicSeekable) OpenRangeReader(_ context.Context, off int64, length int64, _ *storage.FrameTable) (storage.RangeReader, storage.Source, error) { end := min(off+length, int64(len(s.data))) return &panicReader{ data: s.data[off:end], panicAfter: int(s.panicAfter - off), - }, nil + }, storage.SourceFS, nil } type panicReader struct { @@ -460,8 +462,8 @@ func (r *panicReader) Read(p []byte) (int, error) { return n, nil } -func (r *panicReader) Close(context.Context) error { - return nil +func (r *panicReader) Close(context.Context) (*storage.ReadStats, error) { + return nil, nil } func TestChunker_PanicRecovery(t *testing.T) { @@ -587,11 +589,11 @@ func (r *controlledReader) Read(p []byte) (int, error) { return n, nil } -func (r *controlledReader) Close(context.Context) error { +func (r *controlledReader) Close(context.Context) (*storage.ReadStats, error) { select { case r.closed <- struct{}{}: default: } - return nil + return nil, nil } diff --git a/packages/orchestrator/pkg/sandbox/build/build.go b/packages/orchestrator/pkg/sandbox/build/build.go index 4f93d57145..06da10b07f 100644 --- a/packages/orchestrator/pkg/sandbox/build/build.go +++ b/packages/orchestrator/pkg/sandbox/build/build.go @@ -12,6 +12,8 @@ import ( "github.com/google/uuid" "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" "go.uber.org/zap" "golang.org/x/sync/errgroup" @@ -45,12 +47,18 @@ var ( )) ) +var fileReadAtTimer = utils.Must(telemetry.NewFloatTimerFactory(meter, "orchestrator.file.read_at", + "Time to serve a build.File ReadAt across all source builds", + "Bytes read", "ReadAt call count")) + type File struct { header atomic.Pointer[header.Header] store *DiffStore fileType DiffType persistence storage.StorageProvider metrics blockmetrics.Metrics + + okAttrs metric.MeasurementOption } func NewFile( @@ -65,12 +73,23 @@ func NewFile( fileType: fileType, persistence: persistence, metrics: metrics, + okAttrs: metric.WithAttributes( + attribute.String(storage.AttrFileType, string(fileType)), + attribute.String(storage.AttrOutcome, storage.OutcomeOK), + ), } f.header.Store(header) return f } +func (b *File) errAttrs(err error) metric.MeasurementOption { + return metric.WithAttributes( + attribute.String(storage.AttrFileType, string(b.fileType)), + attribute.String(storage.AttrOutcome, storage.Outcome(err)), + ) +} + func (b *File) Header() *header.Header { return b.header.Load() } @@ -79,9 +98,24 @@ func (b *File) SwapHeader(h *header.Header) { b.header.Store(h) } -// ReadAt fills p from the mapped build segments, optionally in parallel. +// ReadAt records file.read_at timing around the readAt worker. Slice's +// compose path goes through readAt directly to avoid double counting. +func (b *File) ReadAt(ctx context.Context, p []byte, off int64) (n int, err error) { + start := time.Now() + n, err = b.readAt(ctx, p, off) + if err != nil { + fileReadAtTimer.Record(ctx, time.Since(start), int64(n), b.errAttrs(err)) + + return n, err + } + fileReadAtTimer.Record(ctx, time.Since(start), int64(n), b.okAttrs) + + return n, nil +} + +// readAt fills p from the mapped build segments, optionally in parallel. // Cache eviction or a peer transition re-resolves and retries. -func (b *File) ReadAt(ctx context.Context, p []byte, off int64) (int, error) { +func (b *File) readAt(ctx context.Context, p []byte, off int64) (int, error) { maxParallel := b.store.flags.IntFlag(ctx, featureflags.MaxParallelBuildReadSegments) for { @@ -259,7 +293,7 @@ func (b *File) Slice(ctx context.Context, off, length int64) ([]byte, error) { } } out := make([]byte, length) - if _, err := b.ReadAt(ctx, out, off); err != nil { + if _, err := b.readAt(ctx, out, off); err != nil { return nil, fmt.Errorf("failed to read at: %w", err) } diff --git a/packages/orchestrator/pkg/sandbox/build/storage_diff.go b/packages/orchestrator/pkg/sandbox/build/storage_diff.go index d3f4081fc7..d468be7bfb 100644 --- a/packages/orchestrator/pkg/sandbox/build/storage_diff.go +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff.go @@ -85,7 +85,7 @@ func newStorageDiff( ff *featureflags.Client, ) (*StorageDiff, error) { cachePath := GenerateDiffCachePath(basePath, buildID, diffType) - c, err := block.NewChunker(ff, uncompressedSize, blockSize, cachePath, metrics) + c, err := block.NewChunker(ff, uncompressedSize, blockSize, cachePath, metrics, storageObjectType) if err != nil { return nil, fmt.Errorf("create chunker for build %s: %w", buildID, err) } diff --git a/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go b/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go index e00fe8021a..42cf8e646a 100644 --- a/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go +++ b/packages/orchestrator/pkg/sandbox/build/storage_diff_test.go @@ -188,10 +188,10 @@ func TestStorageDiff_NoRefreshOnFinalizedHeader(t *testing.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) (storage.RangeReader, error) { + RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (storage.RangeReader, storage.Source, error) { end := min(off+length, int64(len(payload))) - return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), nil + return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), storage.SourceFS, nil }) provider.EXPECT(). OpenSeekable(mock.Anything, uncompressedPath, mock.Anything). @@ -235,10 +235,10 @@ func TestStorageDiff_SkipsHeaderRefreshWhenPeerActive(t *testing.T) { 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) (storage.RangeReader, error) { + RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (storage.RangeReader, storage.Source, error) { end := min(off+length, int64(len(payload))) - return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), nil + return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), storage.SourceFS, nil }) provider.EXPECT(). OpenSeekable(mock.Anything, aPaths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything). @@ -299,10 +299,10 @@ func TestStorageDiff_V3AncestorFallsBackToUncompressed(t *testing.T) { Return(int64(payloadSize), nil).Once() rawSeekable.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything). - RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (storage.RangeReader, error) { + RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (storage.RangeReader, storage.Source, error) { end := min(off+length, int64(len(payload))) - return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), nil + return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), storage.SourceFS, nil }) provider.EXPECT(). OpenSeekable(mock.Anything, aPaths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything). @@ -348,10 +348,10 @@ func TestStorageDiff_ReloadSourceLatchesV3AsUncompressed(t *testing.T) { rawSeekable := storage.NewMockSeekable(t) rawSeekable.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything). - RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (storage.RangeReader, error) { + RunAndReturn(func(_ context.Context, off, length int64, _ *storage.FrameTable) (storage.RangeReader, storage.Source, error) { end := min(off+length, int64(len(payload))) - return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), nil + return storage.NewRangeReader(io.NopCloser(bytes.NewReader(payload[off:end]))), storage.SourceFS, nil }) provider.EXPECT(). OpenSeekable(mock.Anything, aPaths.DataFile(storage.MemfileName, storage.CompressionNone), mock.Anything). @@ -434,14 +434,16 @@ func buildHeader(t *testing.T, selfID uuid.UUID, size int64, mapsTo uuid.UUID) * // 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) (storage.RangeReader, error) { - return func(_ context.Context, offsetU, _ int64, ft *storage.FrameTable) (storage.RangeReader, error) { +func decompressingRangeReader(compressed []byte) func(context.Context, int64, int64, *storage.FrameTable) (storage.RangeReader, storage.Source, error) { + return func(_ context.Context, offsetU, _ int64, ft *storage.FrameTable) (storage.RangeReader, storage.Source, error) { r, err := ft.LocateCompressed(offsetU) if err != nil { - return nil, err + return nil, storage.SourceFS, err } end := min(r.Offset+int64(r.Length), int64(len(compressed))) - return storage.NewDecompressingReader(storage.NewRangeReader(io.NopCloser(bytes.NewReader(compressed[r.Offset:end]))), ft.CompressionType()) + rc, err := storage.NewDecompressingReader(storage.NewRangeReader(io.NopCloser(bytes.NewReader(compressed[r.Offset:end]))), ft.CompressionType()) + + return rc, storage.SourceFS, err } } diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go index ee0c8a972f..ef9209565f 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go @@ -97,7 +97,7 @@ func (s *peerSeekable) Size(ctx context.Context) (int64, error) { return 0, &storage.PeerTransitionedError{} } -func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *storage.FrameTable) (storage.RangeReader, error) { +func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *storage.FrameTable) (storage.RangeReader, storage.Source, error) { res, err := tryPeer(ctx, &s.peerHandle, "peer-seekable-open-range-reader", attrOpRangeReader, func(ctx context.Context) (peerAttempt[storage.RangeReader], error) { streamCtx, cancel := context.WithCancel(ctx) @@ -121,15 +121,15 @@ func (s *peerSeekable) OpenRangeReader(ctx context.Context, off int64, length in }, nil }) if res.hit { - return res.value, err + return res.value, storage.SourcePeer, err } if s.uploaded.Load() { - return nil, &storage.PeerTransitionedError{} + return nil, storage.SourcePeer, &storage.PeerTransitionedError{} } base, err := s.getBase(ctx, frameTable.CompressionType()) if err != nil { - return nil, err + return nil, storage.SourcePeer, err } return base.OpenRangeReader(ctx, off, length, frameTable) diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go index 5c1f8ea9d6..a3f98be5a0 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/seekable_test.go @@ -73,7 +73,7 @@ func TestPeerSeekable_OpenRangeReader_PeerSucceeds(t *testing.T) { })).Return(stream, nil) s := &peerSeekable{peerHandle: peerHandle{client: client, buildID: "build-1", name: storage.MemfileName, uploaded: &atomic.Bool{}}} - rc, err := s.OpenRangeReader(t.Context(), 10, int64(len(data)), nil) + rc, _, err := s.OpenRangeReader(t.Context(), 10, int64(len(data)), nil) require.NoError(t, err) defer rc.Close(t.Context()) @@ -90,7 +90,7 @@ func TestPeerSeekable_OpenRangeReader_PeerError_FallsBackToBase(t *testing.T) { client.EXPECT().ReadAtBuildSeekable(mock.Anything, mock.Anything).Return(nil, errors.New("peer unavailable")) baseSeekable := storage.NewMockSeekable(t) - baseSeekable.EXPECT().OpenRangeReader(mock.Anything, int64(0), int64(len(baseData)), (*storage.FrameTable)(nil)).Return(storage.NewRangeReader(io.NopCloser(bytes.NewReader(baseData))), nil) + baseSeekable.EXPECT().OpenRangeReader(mock.Anything, int64(0), int64(len(baseData)), (*storage.FrameTable)(nil)).Return(storage.NewRangeReader(io.NopCloser(bytes.NewReader(baseData))), storage.SourceFS, nil) base := storage.NewMockStorageProvider(t) base.EXPECT().OpenSeekable(mock.Anything, "build-1/memfile", storage.MemfileObjectType).Return(baseSeekable, nil) @@ -105,7 +105,7 @@ func TestPeerSeekable_OpenRangeReader_PeerError_FallsBackToBase(t *testing.T) { basePersistence: base, objType: storage.MemfileObjectType, } - rc, err := s.OpenRangeReader(t.Context(), 0, int64(len(baseData)), nil) + rc, _, err := s.OpenRangeReader(t.Context(), 0, int64(len(baseData)), nil) require.NoError(t, err) defer rc.Close(t.Context()) @@ -137,7 +137,7 @@ func TestPeerSeekable_OpenRangeReader_Uploaded_ReturnsPeerTransitionedError(t *t objType: storage.MemfileObjectType, } - _, err := s.OpenRangeReader(t.Context(), 0, 100, nil) + _, _, err := s.OpenRangeReader(t.Context(), 0, 100, nil) require.Error(t, err) var transErr *storage.PeerTransitionedError @@ -180,16 +180,17 @@ func TestPeerStorageProvider_TransitionEmitsError(t *testing.T) { seekable, err := p.OpenSeekable(t.Context(), "build-1/memfile", storage.MemfileObjectType) require.NoError(t, err) - rc, err := seekable.OpenRangeReader(t.Context(), 0, int64(len(prePeerBytes)), + rc, _, err := seekable.OpenRangeReader(t.Context(), 0, int64(len(prePeerBytes)), storage.NewFullFrameTable(storage.CompressionNone, nil).Table()) require.NoError(t, err) got, err := io.ReadAll(rc) require.NoError(t, err) - require.NoError(t, rc.Close(t.Context())) + _, err = rc.Close(t.Context()) + require.NoError(t, err) assert.Equal(t, prePeerBytes, got) require.True(t, uploaded.Load(), "uploaded flag should be set after peer EOF with UseStorage") - _, err = seekable.OpenRangeReader(t.Context(), 0, 1, + _, _, err = seekable.OpenRangeReader(t.Context(), 0, 1, storage.NewFullFrameTable(storage.CompressionNone, nil).Table()) var transErr *storage.PeerTransitionedError require.ErrorAs(t, err, &transErr) diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/storage.go b/packages/orchestrator/pkg/sandbox/template/peerclient/storage.go index de51c870b8..55ea1e749c 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/storage.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/storage.go @@ -277,11 +277,17 @@ var _ storage.RangeReader = (*peerStreamReader)(nil) // peerStreamReader wraps a gRPC streaming recv function as a storage.RangeReader. // cancel is called on Close to signal the server to terminate the stream. +// bytes/read accumulate the source-read wall and byte count so Close can +// return non-nil ReadStats — Chunker.runFetch otherwise records this fetch as +// 0 bytes and skips orchestrator.read.read for peer-served reads. type peerStreamReader struct { recv func() ([]byte, error) current *bytes.Reader done bool cancel context.CancelFunc + + bytes int64 + read time.Duration } func newPeerStreamReader(recv func() ([]byte, error), cancel context.CancelFunc) *peerStreamReader { @@ -292,6 +298,15 @@ func newPeerStreamReader(recv func() ([]byte, error), cancel context.CancelFunc) } func (r *peerStreamReader) Read(p []byte) (int, error) { + t0 := time.Now() + n, err := r.read1(p) + r.read += time.Since(t0) + r.bytes += int64(n) + + return n, err +} + +func (r *peerStreamReader) read1(p []byte) (int, error) { for { if r.current != nil && r.current.Len() > 0 { return r.current.Read(p) @@ -317,8 +332,12 @@ func (r *peerStreamReader) Read(p []byte) (int, error) { } } -func (r *peerStreamReader) Close(context.Context) error { +func (r *peerStreamReader) Close(context.Context) (*storage.ReadStats, error) { r.cancel() - return nil + return &storage.ReadStats{ + CompressedBytes: r.bytes, + UncompressedBytes: r.bytes, + Read: r.read, + }, nil } diff --git a/packages/shared/pkg/storage/compress_decode.go b/packages/shared/pkg/storage/compress_decode.go index a0a706bab3..d262fed86c 100644 --- a/packages/shared/pkg/storage/compress_decode.go +++ b/packages/shared/pkg/storage/compress_decode.go @@ -2,9 +2,11 @@ package storage import ( "context" + "errors" "fmt" "io" "sync" + "time" "github.com/klauspost/compress/zstd" lz4 "github.com/pierrec/lz4/v4" @@ -56,25 +58,44 @@ func putZstdDecoder(dec *zstd.Decoder) { zstdDecoderPool.Put(dec) } -// decompressReader decompresses inner on Read; Close releases the codec back -// to its pool and closes inner. +// decompressReader decompresses on Read, metering raw pulls vs decoded output +// separately so source-read wall and decompression CPU are split. On Close it +// emits the orchestrator.read.decompress record (decompress time + +// uncompressed bytes, keyed by file_type/source/codec). No drain: the source +// range is bounded to exactly C bytes by construction, and the cache-writeback +// path's captureReader has its own drainOnClose for the LZ4-BlockChecksum +// 4-byte EndMark tail. type decompressReader struct { - inner RangeReader - dec io.Reader + inner RangeReader // retained to call Close; + meteredIn *meteredReader + meteredOut *meteredReader releaseCodec func() + ct CompressionType + source Source + objType SeekableObjectType + readErr error } +// NewDecompressingReader wraps inner so Read returns decompressed bytes; +// metric attribution falls back to defaults (callers that care provide it via +// newDecompressReader directly). func NewDecompressingReader(inner RangeReader, ct CompressionType) (RangeReader, error) { + return newDecompressReader(inner, ct, UnknownSource, UnknownSeekableObjectType) +} + +func newDecompressReader(inner RangeReader, ct CompressionType, src Source, ot SeekableObjectType) (*decompressReader, error) { + compressed := &meteredReader{inner: inner} + var dec io.Reader var releaseCodec func() switch ct { case CompressionLZ4: - d := getLZ4Decoder(inner) + d := getLZ4Decoder(compressed) dec, releaseCodec = d, func() { putLZ4Decoder(d) } case CompressionZstd: - d, err := getZstdDecoder(inner) + d, err := getZstdDecoder(compressed) if err != nil { return nil, fmt.Errorf("failed to create zstd decoder: %w", err) } @@ -86,17 +107,37 @@ func NewDecompressingReader(inner RangeReader, ct CompressionType) (RangeReader, return &decompressReader{ inner: inner, - dec: dec, + meteredIn: compressed, + meteredOut: &meteredReader{inner: dec}, releaseCodec: releaseCodec, + ct: ct, + source: src, + objType: ot, }, nil } func (r *decompressReader) Read(p []byte) (int, error) { - return r.dec.Read(p) + n, err := r.meteredOut.Read(p) + if err != nil && !errors.Is(err, io.EOF) { + r.readErr = err + } + + return n, err } -func (r *decompressReader) Close(ctx context.Context) error { +func (r *decompressReader) Close(ctx context.Context) (*ReadStats, error) { r.releaseCodec() - return r.inner.Close(ctx) + stats := &ReadStats{ + CompressedBytes: r.meteredIn.bytes, + UncompressedBytes: r.meteredOut.bytes, + Read: time.Duration(r.meteredIn.nanos), + Decompress: max(0, time.Duration(r.meteredOut.nanos)-time.Duration(r.meteredIn.nanos)), + } + + recordDecompressStep(ctx, r, stats, r.readErr) + + _, innerErr := r.inner.Close(ctx) + + return stats, innerErr } diff --git a/packages/shared/pkg/storage/io_wrappers.go b/packages/shared/pkg/storage/io_wrappers.go index 18d5e5ebba..c3fdc14344 100644 --- a/packages/shared/pkg/storage/io_wrappers.go +++ b/packages/shared/pkg/storage/io_wrappers.go @@ -6,6 +6,7 @@ import ( "errors" "io" "os" + "time" "go.opentelemetry.io/otel/trace" @@ -13,6 +14,7 @@ import ( ) var ( + _ io.Reader = (*meteredReader)(nil) _ RangeReader = (*sectionReader)(nil) _ RangeReader = (*observableReader)(nil) _ RangeReader = (*rangeReader)(nil) @@ -20,15 +22,15 @@ var ( ) // rangeReader adapts an io.ReadCloser into a RangeReader by ignoring the -// Close context. +// Close context. It does not meter, so Close returns nil stats. type rangeReader struct { io.ReadCloser } func NewRangeReader(rc io.ReadCloser) RangeReader { return &rangeReader{ReadCloser: rc} } -func (p *rangeReader) Close(context.Context) error { - return p.ReadCloser.Close() +func (p *rangeReader) Close(context.Context) (*ReadStats, error) { + return nil, p.ReadCloser.Close() } type sectionReader struct { @@ -44,16 +46,33 @@ func newSectionReader(f *os.File, off, length int64) *sectionReader { } } -func (r *sectionReader) Close(context.Context) error { - return r.file.Close() +func (r *sectionReader) Close(context.Context) (*ReadStats, error) { + return nil, r.file.Close() +} + +// meteredReader records cumulative time and bytes spent pulling from inner so +// a decoder built on top can separate source-read wall from decompression CPU. +// Single-goroutine: Read is sequential, stats are read after EOF in Close. +type meteredReader struct { + inner io.Reader + nanos int64 + bytes int64 +} + +func (m *meteredReader) Read(p []byte) (int, error) { + t0 := time.Now() + n, err := m.inner.Read(p) + m.nanos += int64(time.Since(t0)) + m.bytes += int64(n) + + return n, err } // captureReader tees every read byte into a buffer and hands the captured -// bytes to onClose on Close. Used by the cache writeback paths. -// +// bytes to onClose on Close. Used by the compressed cache writeback path. // drainOnClose=true reads inner to EOF on Close even if the caller above hasn't -// consumed everything. Needed when capturing under a decoder that stops short -// of EOF on its source (e.g. lz4.Reader skips the 4-byte EndMark). +// consumed everything — the compressed cache needs the full frame regardless +// of how many bytes the decoder happened to demand. type captureReader struct { inner RangeReader buf *bytes.Buffer @@ -79,36 +98,52 @@ func (r *captureReader) Read(p []byte) (int, error) { return n, err } -func (r *captureReader) Close(ctx context.Context) error { +func (r *captureReader) Close(ctx context.Context) (*ReadStats, error) { if r.drainOnClose { _, _ = io.Copy(io.Discard, r) } - err := r.inner.Close(ctx) + stats, err := r.inner.Close(ctx) if err == nil { r.onClose(ctx, r.buf.Bytes()) } - return err + return stats, err } -// observableReader layers OTEL observability (legacy per-backend timer + span) -// onto an inner RangeReader, all applied on Close. timer and span are optional; -// pass nil if unused. +// observableReader layers OTEL observability onto an inner RangeReader, all +// applied on Close. The with* builder methods are optional and chainable. type observableReader struct { inner RangeReader + timer *telemetry.Stopwatch span trace.Span bytes int64 readErr error + + read time.Duration } -func newObservableReader(inner RangeReader, timer *telemetry.Stopwatch, span trace.Span) *observableReader { - return &observableReader{inner: inner, timer: timer, span: span} +func newObservableReader(inner RangeReader) *observableReader { + return &observableReader{inner: inner} +} + +func (r *observableReader) withTimer(t *telemetry.Stopwatch) *observableReader { + r.timer = t + + return r +} + +func (r *observableReader) withSpan(s trace.Span) *observableReader { + r.span = s + + return r } func (r *observableReader) Read(p []byte) (int, error) { + t0 := time.Now() n, err := r.inner.Read(p) + r.read += time.Since(t0) r.bytes += int64(n) if err != nil && !errors.Is(err, io.EOF) { @@ -118,8 +153,16 @@ func (r *observableReader) Read(p []byte) (int, error) { return n, err } -func (r *observableReader) Close(ctx context.Context) error { - closeErr := r.inner.Close(ctx) +func (r *observableReader) Close(ctx context.Context) (*ReadStats, error) { + stats, closeErr := r.inner.Close(ctx) + + if stats == nil { + stats = &ReadStats{ + CompressedBytes: r.bytes, + UncompressedBytes: r.bytes, + Read: r.read, + } + } if r.timer != nil { if r.readErr != nil || closeErr != nil { @@ -139,5 +182,5 @@ func (r *observableReader) Close(ctx context.Context) error { r.span.End() } - return closeErr + return stats, closeErr } diff --git a/packages/shared/pkg/storage/mock_seekable.go b/packages/shared/pkg/storage/mock_seekable.go index 6c3e18a4fc..3fa5665c56 100644 --- a/packages/shared/pkg/storage/mock_seekable.go +++ b/packages/shared/pkg/storage/mock_seekable.go @@ -38,7 +38,7 @@ func (_m *MockSeekable) EXPECT() *MockSeekable_Expecter { } // OpenRangeReader provides a mock function for the type MockSeekable -func (_mock *MockSeekable) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, error) { +func (_mock *MockSeekable) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, Source, error) { ret := _mock.Called(ctx, offsetU, length, frameTable) if len(ret) == 0 { @@ -46,8 +46,9 @@ func (_mock *MockSeekable) OpenRangeReader(ctx context.Context, offsetU int64, l } var r0 RangeReader - var r1 error - if returnFunc, ok := ret.Get(0).(func(context.Context, int64, int64, *FrameTable) (RangeReader, error)); ok { + var r1 Source + var r2 error + if returnFunc, ok := ret.Get(0).(func(context.Context, int64, int64, *FrameTable) (RangeReader, Source, error)); ok { return returnFunc(ctx, offsetU, length, frameTable) } if returnFunc, ok := ret.Get(0).(func(context.Context, int64, int64, *FrameTable) RangeReader); ok { @@ -57,12 +58,17 @@ func (_mock *MockSeekable) OpenRangeReader(ctx context.Context, offsetU int64, l r0 = ret.Get(0).(RangeReader) } } - if returnFunc, ok := ret.Get(1).(func(context.Context, int64, int64, *FrameTable) error); ok { + if returnFunc, ok := ret.Get(1).(func(context.Context, int64, int64, *FrameTable) Source); ok { r1 = returnFunc(ctx, offsetU, length, frameTable) } else { - r1 = ret.Error(1) + r1 = ret.Get(1).(Source) } - return r0, r1 + if returnFunc, ok := ret.Get(2).(func(context.Context, int64, int64, *FrameTable) error); ok { + r2 = returnFunc(ctx, offsetU, length, frameTable) + } else { + r2 = ret.Error(2) + } + return r0, r1, r2 } // MockSeekable_OpenRangeReader_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'OpenRangeReader' @@ -107,12 +113,12 @@ func (_c *MockSeekable_OpenRangeReader_Call) Run(run func(ctx context.Context, o return _c } -func (_c *MockSeekable_OpenRangeReader_Call) Return(rangeReader RangeReader, err error) *MockSeekable_OpenRangeReader_Call { - _c.Call.Return(rangeReader, err) +func (_c *MockSeekable_OpenRangeReader_Call) Return(rangeReader RangeReader, source Source, err error) *MockSeekable_OpenRangeReader_Call { + _c.Call.Return(rangeReader, source, err) return _c } -func (_c *MockSeekable_OpenRangeReader_Call) RunAndReturn(run func(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, error)) *MockSeekable_OpenRangeReader_Call { +func (_c *MockSeekable_OpenRangeReader_Call) RunAndReturn(run func(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, Source, error)) *MockSeekable_OpenRangeReader_Call { _c.Call.Return(run) return _c } diff --git a/packages/shared/pkg/storage/read_attrs.go b/packages/shared/pkg/storage/read_attrs.go new file mode 100644 index 0000000000..2c1559c12d --- /dev/null +++ b/packages/shared/pkg/storage/read_attrs.go @@ -0,0 +1,26 @@ +package storage + +// Closed-enum attribute vocabulary for the orchestrator.read.* / +// orchestrator.chunk.* metric families. + +const ( + AttrSource = "source" + AttrCodec = "codec" + AttrOutcome = "outcome" + AttrEvent = "event" + AttrFileType = "file_type" +) + +const ( + OutcomeOK = "ok" + OutcomeErrCanceled = "err_canceled" + OutcomeErrIO = "err_io" + OutcomeErrTimeout = "err_timeout" +) + +const ( + CacheEventHit = "hit" + CacheEventMiss = "miss" + CacheEventWritebackOK = "writeback_ok" + CacheEventWritebackErr = "writeback_err" +) diff --git a/packages/shared/pkg/storage/read_attrs_precomputed.go b/packages/shared/pkg/storage/read_attrs_precomputed.go new file mode 100644 index 0000000000..08944bd0c8 --- /dev/null +++ b/packages/shared/pkg/storage/read_attrs_precomputed.go @@ -0,0 +1,145 @@ +package storage + +import ( + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +// Precomputed attribute sets for hot-path read.*/chunk.* emissions; cold error +// paths build attrs inline via ErrAttrs. + +// Source identifies the backend that served a read. The zero value +// (UnknownSource) is the default for pre-resolution failures and any state +// before the backend is known. +type Source int8 + +const ( + // Order is latency-ascending and load-bearing: a multi-chunk Slice records + // the slowest source it touched via max() over per-fetch sources. Unknown + // at 0 is lighter than every real source so it gets replaced on first hit. + UnknownSource Source = iota + SourceMmap + SourceFS + SourcePeer + SourceNFS + SourceGCS + SourceAWS + numSources +) + +func (s Source) String() string { return sourceStrings[s] } + +var sourceStrings = [numSources]string{ + UnknownSource: "unknown", + SourceMmap: "mmap", + SourceFS: "fs", + SourcePeer: "peer", + SourceNFS: "nfs", + SourceGCS: "gcs", + SourceAWS: "aws", +} + +const numCodecs = 3 // CompressionNone, Zstd, LZ4 + +var ( + tableOK [numSeekableObjectTypes][numSources][numCodecs]metric.MeasurementOption + + tableCacheHit [numSeekableObjectTypes][numSources][numCodecs]metric.MeasurementOption + tableCacheMiss [numSeekableObjectTypes][numSources][numCodecs]metric.MeasurementOption + tableCacheWritebackOK [numSeekableObjectTypes][numSources][numCodecs]metric.MeasurementOption + tableCacheWritebackErr [numSeekableObjectTypes][numSources][numCodecs]metric.MeasurementOption + + // keyed by file_type only: inflight is incremented before the source is + // known (the OpenRangeReader call itself dominates GCS latency). + tableInflightFetch [numSeekableObjectTypes]metric.MeasurementOption +) + +func init() { + set := func(kvs ...attribute.KeyValue) metric.MeasurementOption { + return metric.WithAttributeSet(attribute.NewSet(kvs...)) + } + + for ot := range numSeekableObjectTypes { + ftAttr := attribute.String(AttrFileType, ot.String()) + + tableInflightFetch[ot] = set(ftAttr) + + for s := range numSources { + srcAttr := attribute.String(AttrSource, sourceStrings[s]) + + for ct := range CompressionType(numCodecs) { + codecAttr := attribute.String(AttrCodec, ct.String()) + outcomeOK := attribute.String(AttrOutcome, OutcomeOK) + + tableOK[ot][s][ct] = set( + ftAttr, srcAttr, codecAttr, outcomeOK, + ) + + tableCacheHit[ot][s][ct] = set( + ftAttr, attribute.String(AttrEvent, CacheEventHit), + srcAttr, codecAttr, + ) + tableCacheMiss[ot][s][ct] = set( + ftAttr, attribute.String(AttrEvent, CacheEventMiss), + srcAttr, codecAttr, + ) + tableCacheWritebackOK[ot][s][ct] = set( + ftAttr, attribute.String(AttrEvent, CacheEventWritebackOK), + srcAttr, codecAttr, + ) + tableCacheWritebackErr[ot][s][ct] = set( + ftAttr, attribute.String(AttrEvent, CacheEventWritebackErr), + srcAttr, codecAttr, + ) + } + } + } +} + +func safeAttrIdx(o SeekableObjectType, s Source, c CompressionType) (SeekableObjectType, Source, CompressionType) { + if uint(o) >= uint(numSeekableObjectTypes) { + o = UnknownSeekableObjectType + } + if uint(s) >= uint(numSources) { + s = UnknownSource + } + if uint(c) >= numCodecs { + c = CompressionNone + } + + return o, s, c +} + +func OKAttrs(o SeekableObjectType, s Source, c CompressionType) metric.MeasurementOption { + o, s, c = safeAttrIdx(o, s, c) + + return tableOK[o][s][c] +} + +func CacheHitAttrs(o SeekableObjectType, s Source, c CompressionType) metric.MeasurementOption { + o, s, c = safeAttrIdx(o, s, c) + + return tableCacheHit[o][s][c] +} + +func CacheMissAttrs(o SeekableObjectType, s Source, c CompressionType) metric.MeasurementOption { + o, s, c = safeAttrIdx(o, s, c) + + return tableCacheMiss[o][s][c] +} + +func CacheWritebackOKAttrs(o SeekableObjectType, s Source, c CompressionType) metric.MeasurementOption { + o, s, c = safeAttrIdx(o, s, c) + + return tableCacheWritebackOK[o][s][c] +} + +func CacheWritebackErrAttrs(o SeekableObjectType, s Source, c CompressionType) metric.MeasurementOption { + o, s, c = safeAttrIdx(o, s, c) + + return tableCacheWritebackErr[o][s][c] +} + +func InflightFetchAttrs(o SeekableObjectType) metric.MeasurementOption { + return tableInflightFetch[o] +} diff --git a/packages/shared/pkg/storage/read_attrs_test.go b/packages/shared/pkg/storage/read_attrs_test.go new file mode 100644 index 0000000000..7d1b549f5f --- /dev/null +++ b/packages/shared/pkg/storage/read_attrs_test.go @@ -0,0 +1,57 @@ +package storage + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestSourceStringsPopulated(t *testing.T) { + t.Parallel() + + for s := range numSources { + require.NotEmptyf(t, s.String(), "source %d has empty string label", s) + } +} + +func TestSeekableObjectTypeStrings(t *testing.T) { + t.Parallel() + + require.Equal(t, "memfile", MemfileObjectType.String()) + require.Equal(t, "rootfs", RootFSObjectType.String()) + require.Equal(t, "unknown", UnknownSeekableObjectType.String()) + for o := range numSeekableObjectTypes { + require.NotEmptyf(t, o.String(), "object type %d has empty file_type label", o) + } +} + +func TestOutcomeMapping(t *testing.T) { + t.Parallel() + + require.Equal(t, OutcomeOK, Outcome(nil)) + require.Equal(t, OutcomeErrCanceled, Outcome(context.Canceled)) + require.Equal(t, OutcomeErrTimeout, Outcome(context.DeadlineExceeded)) + require.Equal(t, OutcomeErrIO, Outcome(errors.New("boom"))) +} + +// TestPrecomputedAttrsPopulated guards the invariant that every emission site +// finds a non-nil precomputed attribute set — no enum combination is missed by +// the init() loops. +func TestPrecomputedAttrsPopulated(t *testing.T) { + t.Parallel() + + for o := range numSeekableObjectTypes { + require.NotNil(t, InflightFetchAttrs(o)) + for s := range numSources { + for c := range CompressionType(numCodecs) { + require.NotNil(t, OKAttrs(o, s, c)) + require.NotNil(t, CacheHitAttrs(o, s, c)) + require.NotNil(t, CacheMissAttrs(o, s, c)) + require.NotNil(t, CacheWritebackOKAttrs(o, s, c)) + require.NotNil(t, CacheWritebackErrAttrs(o, s, c)) + } + } + } +} diff --git a/packages/shared/pkg/storage/read_metrics.go b/packages/shared/pkg/storage/read_metrics.go new file mode 100644 index 0000000000..ae88776442 --- /dev/null +++ b/packages/shared/pkg/storage/read_metrics.go @@ -0,0 +1,151 @@ +package storage + +import ( + "context" + "errors" + "time" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + + "github.com/e2b-dev/infra/packages/shared/pkg/telemetry" + "github.com/e2b-dev/infra/packages/shared/pkg/utils" +) + +// Instruments for the `orchestrator.read.*` family — one metric per stage +// (open/read/decompress/fetch/writeback) keyed by file_type/source/codec/outcome. + +func mustFloatHist(name, desc, unit string) metric.Float64Histogram { + return utils.Must(meter.Float64Histogram(name, + metric.WithDescription(desc), + metric.WithUnit(unit), + )) +} + +var ( + readOpen = utils.Must(telemetry.NewFloatTimerFactory(meter, + "orchestrator.read.open", + "OpenRangeReader (open / TTFB) wall", + "Bytes (always 0 — open transfers no payload)", + "Number of opens", + )) + readRead = utils.Must(telemetry.NewFloatTimerFactory(meter, + "orchestrator.read.read", + "Raw source-read wall (decompression excluded)", + "Compressed/stored bytes read from the source", + "Number of source reads", + )) + readDecompress = utils.Must(telemetry.NewFloatTimerFactory(meter, + "orchestrator.read.decompress", + "Decompression CPU wall (decoder read time minus source transfer)", + "Uncompressed bytes produced", + "Number of decompress records", + )) + readFetch = utils.Must(telemetry.NewFloatTimerFactory(meter, + "orchestrator.read.fetch", + "Total fetch wall — should ≈ open + read + decompress; any excess is overhead (see read.pipeline.efficiency)", + "Bytes delivered to the app", + "Number of fetches", + )) + readWriteback = utils.Must(telemetry.NewFloatTimerFactory(meter, + "orchestrator.read.writeback", + "Async NFS cache writeback wall", + "Bytes written to NFS", + "Number of writebacks", + )) + + // fetch / (open + read + decompress); 1.0 = fully explained, >1 = overhead. + readPipelineEfficiency = mustFloatHist( + "orchestrator.read.pipeline.efficiency", + "fetch / (open + read + decompress) — 1.0 = fetch wall fully explained by work, >1 = overhead", "1", + ) + + readCache = utils.Must(meter.Int64Counter( + "orchestrator.read.cache", + metric.WithDescription("NFS read-cache events (hit / miss / writeback). The mmap tier is orchestrator.chunk.cache."), + metric.WithUnit("1"), + )) + + readInflight = utils.Must(meter.Int64UpDownCounter( + "orchestrator.read.inflight", + metric.WithDescription("In-flight read-path fetches (cache miss → backend), by file_type"), + metric.WithUnit("1"), + )) +) + +func RecordReadOpen(ctx context.Context, dur time.Duration, bytes int64, attrs metric.MeasurementOption) { + readOpen.Record(ctx, dur, bytes, attrs) +} + +func RecordReadRead(ctx context.Context, dur time.Duration, bytes int64, attrs metric.MeasurementOption) { + readRead.Record(ctx, dur, bytes, attrs) +} + +func RecordReadFetch(ctx context.Context, dur time.Duration, bytes int64, attrs metric.MeasurementOption) { + readFetch.Record(ctx, dur, bytes, attrs) +} + +func RecordReadDecompress(ctx context.Context, dur time.Duration, bytes int64, attrs metric.MeasurementOption) { + readDecompress.Record(ctx, dur, bytes, attrs) +} + +func RecordPipelineEfficiency(ctx context.Context, ratio float64, attrs metric.MeasurementOption) { + readPipelineEfficiency.Record(ctx, ratio, attrs) +} + +// StartInflight increments the read.inflight gauge and returns a func that +// decrements it; defer the returned func so the +1/-1 can't drift apart. +func StartInflight(ctx context.Context, attrs metric.MeasurementOption) func() { + readInflight.Add(ctx, 1, attrs) + + return func() { readInflight.Add(ctx, -1, attrs) } +} + +// Outcome maps a read-path error to the closed read.* outcome enum. +func Outcome(err error) string { + switch { + case err == nil: + return OutcomeOK + case errors.Is(err, context.Canceled): + return OutcomeErrCanceled + case errors.Is(err, context.DeadlineExceeded): + return OutcomeErrTimeout + default: + return OutcomeErrIO + } +} + +// ErrAttrs builds the error-path attribute set for read.* records. Hot OK +// paths use the precomputed OKAttrs. +func ErrAttrs(o SeekableObjectType, s Source, c CompressionType, err error) metric.MeasurementOption { + return metric.WithAttributes( + attribute.String(AttrFileType, o.String()), + attribute.String(AttrSource, s.String()), + attribute.String(AttrCodec, c.String()), + attribute.String(AttrOutcome, Outcome(err)), + ) +} + +func recordDecompressStep(ctx context.Context, r *decompressReader, stats *ReadStats, readErr error) { + if readErr == nil { + readDecompress.Record(ctx, stats.Decompress, stats.UncompressedBytes, OKAttrs(r.objType, r.source, r.ct)) + + return + } + + readDecompress.Record(ctx, stats.Decompress, stats.UncompressedBytes, ErrAttrs(r.objType, r.source, r.ct, readErr)) +} + +// recordWriteback emits the read.writeback timer and its read.cache event. +// src is the originating fetch source (kept for cross-correlation); writebacks +// always target NFS. +func recordWriteback(ctx context.Context, dur time.Duration, bytes int64, ot SeekableObjectType, src Source, ct CompressionType, err error) { + if err == nil { + readWriteback.Record(ctx, dur, bytes, OKAttrs(ot, src, ct)) + readCache.Add(ctx, 1, CacheWritebackOKAttrs(ot, SourceNFS, ct)) + + return + } + readWriteback.Record(ctx, dur, bytes, ErrAttrs(ot, src, ct, err)) + readCache.Add(ctx, 1, CacheWritebackErrAttrs(ot, SourceNFS, ct)) +} diff --git a/packages/shared/pkg/storage/storage.go b/packages/shared/pkg/storage/storage.go index 43491661ef..cf145515a2 100644 --- a/packages/shared/pkg/storage/storage.go +++ b/packages/shared/pkg/storage/storage.go @@ -68,8 +68,20 @@ const ( UnknownSeekableObjectType SeekableObjectType = iota MemfileObjectType RootFSObjectType + numSeekableObjectTypes ) +func (t SeekableObjectType) String() string { + switch t { + case MemfileObjectType: + return "memfile" + case RootFSObjectType: + return "rootfs" + default: + return "unknown" + } +} + type ObjectType int const ( @@ -140,14 +152,32 @@ 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) +} + +// ReadStats is what a RangeReader did over its lifetime; returned from Close. +type ReadStats struct { + CompressedBytes int64 + UncompressedBytes int64 + Read time.Duration // source I/O wall, excluding open and decompression + Decompress time.Duration +} + + type RangeReader interface { io.Reader - Close(ctx context.Context) error + // Close returns stats from the reader's lifetime, or nil if the reader did + // not meter (e.g. a pure adapter). Callers should treat nil as "no stats". + Close(ctx context.Context) (*ReadStats, error) } // RangeOpener supports progressive reads via a streaming range reader. +// OpenRangeReader returns the Source that served the bytes. type RangeOpener interface { - OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, error) + OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, Source, error) } type SeekableWriter interface { diff --git a/packages/shared/pkg/storage/storage_aws.go b/packages/shared/pkg/storage/storage_aws.go index b608eeb423..7c15025a45 100644 --- a/packages/shared/pkg/storage/storage_aws.go +++ b/packages/shared/pkg/storage/storage_aws.go @@ -233,9 +233,9 @@ func (o *awsObject) Put(ctx context.Context, data []byte, opts ...PutOption) err return nil } -func (o *awsObject) OpenRangeReader(ctx context.Context, off, length int64, frameTable *FrameTable) (RangeReader, error) { +func (o *awsObject) OpenRangeReader(ctx context.Context, off, length int64, frameTable *FrameTable) (RangeReader, Source, error) { if frameTable.IsCompressed() { - return nil, errors.New("compressed reads are not supported on AWS") + return nil, SourceAWS, errors.New("compressed reads are not supported on AWS") } readRange := aws.String(fmt.Sprintf("bytes=%d-%d", off, off+length-1)) @@ -247,13 +247,13 @@ func (o *awsObject) OpenRangeReader(ctx context.Context, off, length int64, fram if err != nil { var nsk *types.NoSuchKey if errors.As(err, &nsk) { - return nil, ErrObjectNotExist + return nil, SourceAWS, ErrObjectNotExist } - return nil, fmt.Errorf("failed to create S3 range reader for %q: %w", o.path, err) + return nil, SourceAWS, fmt.Errorf("failed to create S3 range reader for %q: %w", o.path, err) } - return NewRangeReader(resp.Body), nil + return newObservableReader(NewRangeReader(resp.Body)), SourceAWS, nil } func (o *awsObject) Size(ctx context.Context) (int64, error) { diff --git a/packages/shared/pkg/storage/storage_cache.go b/packages/shared/pkg/storage/storage_cache.go index 7d5838758c..e53a7110c7 100644 --- a/packages/shared/pkg/storage/storage_cache.go +++ b/packages/shared/pkg/storage/storage_cache.go @@ -122,6 +122,7 @@ func (c cache) OpenSeekable(ctx context.Context, path string, objectType Seekabl inner: innerObject, flags: c.flags, tracer: c.tracer, + objType: objectType, }, nil } diff --git a/packages/shared/pkg/storage/storage_cache_compressed_test.go b/packages/shared/pkg/storage/storage_cache_compressed_test.go index d19b396c15..205747934c 100644 --- a/packages/shared/pkg/storage/storage_cache_compressed_test.go +++ b/packages/shared/pkg/storage/storage_cache_compressed_test.go @@ -69,8 +69,8 @@ func TestDecompressingCacheReader(t *testing.T) { framePath := makeFrameFilename(c.path, Range{Offset: 0, Length: len(compressed)}) capturing := newCaptureReader(bytesRangeReader(compressed), len(compressed), true, - c.compressedFrameWriteback(framePath, 0, len(compressed))) - rc, err := NewDecompressingReader(capturing, CompressionLZ4) + c.compressedFrameWriteback(framePath, 0, len(compressed), SourceFS, CompressionLZ4)) + rc, err := newDecompressReader(capturing, CompressionLZ4, SourceFS, c.objType) require.NoError(t, err) got, err := io.ReadAll(rc) @@ -102,8 +102,8 @@ func TestDecompressingCacheReader(t *testing.T) { framePath := makeFrameFilename(c.path, Range{Offset: 0, Length: len(compressedProd)}) capturing := newCaptureReader(bytesRangeReader(compressedProd), len(compressedProd), true, - c.compressedFrameWriteback(framePath, 0, len(compressedProd))) - rc, err := NewDecompressingReader(capturing, CompressionLZ4) + c.compressedFrameWriteback(framePath, 0, len(compressedProd), SourceFS, CompressionLZ4)) + rc, err := newDecompressReader(capturing, CompressionLZ4, SourceFS, c.objType) require.NoError(t, err) out := make([]byte, len(original)) @@ -112,7 +112,7 @@ func TestDecompressingCacheReader(t *testing.T) { require.Equal(t, len(original), n) require.Equal(t, original, out) - closeErr := rc.Close(t.Context()) + _, closeErr := rc.Close(t.Context()) require.NoError(t, closeErr, "writeback failure must not surface as a read error") c.wg.Wait() @@ -127,15 +127,15 @@ func TestDecompressingCacheReader(t *testing.T) { framePath := makeFrameFilename(c.path, Range{Offset: 0, Length: len(compressed)}) capturing := newCaptureReader(bytesRangeReader(compressed), len(compressed)+100, true, - c.compressedFrameWriteback(framePath, 0, len(compressed)+100)) // wrong expected size - rc, err := NewDecompressingReader(capturing, CompressionLZ4) + c.compressedFrameWriteback(framePath, 0, len(compressed)+100, SourceFS, CompressionLZ4)) // wrong expected size + rc, err := newDecompressReader(capturing, CompressionLZ4, SourceFS, c.objType) require.NoError(t, err) got, err := io.ReadAll(rc) require.NoError(t, err) require.Equal(t, original, got, "decompressed data should be correct regardless") - closeErr := rc.Close(t.Context()) + _, closeErr := rc.Close(t.Context()) require.NoError(t, closeErr, "writeback failure must not surface as a read error") c.wg.Wait() diff --git a/packages/shared/pkg/storage/storage_cache_seekable.go b/packages/shared/pkg/storage/storage_cache_seekable.go index 2891047aaf..e8f42334be 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable.go +++ b/packages/shared/pkg/storage/storage_cache_seekable.go @@ -9,6 +9,7 @@ import ( "path/filepath" "strconv" "sync" + "time" "github.com/google/uuid" "github.com/launchdarkly/go-sdk-common/v3/ldcontext" @@ -70,6 +71,7 @@ type cachedSeekable struct { inner Seekable flags featureFlagsClient tracer trace.Tracer + objType SeekableObjectType wg sync.WaitGroup } @@ -79,7 +81,7 @@ var ( _ RangeOpener = (*cachedSeekable)(nil) ) -func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *FrameTable) (RangeReader, error) { +func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length int64, frameTable *FrameTable) (RangeReader, Source, error) { compressed := frameTable.IsCompressed() ctx, span := c.tracer.Start(ctx, "read", trace.WithAttributes( @@ -89,24 +91,28 @@ func (c *cachedSeekable) OpenRangeReader(ctx context.Context, off int64, length )) var rc RangeReader + var source Source var err error - if compressed { - rc, err = c.openReaderCompressed(ctx, off, frameTable) - } else if err = c.validateReadParams(length, off); err == nil { - rc, err = c.openReaderUncompressed(ctx, off, length) + switch { + case compressed: + rc, source, err = c.openReaderCompressed(ctx, off, frameTable) + default: + if err = c.validateReadParams(length, off); err == nil { + rc, source, err = c.openReaderUncompressed(ctx, off, length) + } } if err != nil { recordError(span, err) span.End() - return nil, err + return nil, source, err } - return newObservableReader(rc, nil, span), nil + return newObservableReader(rc).withSpan(span), source, nil } -func (c *cachedSeekable) openReaderUncompressed(ctx context.Context, off, length int64) (RangeReader, error) { +func (c *cachedSeekable) openReaderUncompressed(ctx context.Context, off, length int64) (RangeReader, Source, error) { timer := cacheSlabReadTimerFactory.Begin( attribute.String(nfsCacheOperationAttr, nfsCacheOperationAttrReadAt), attribute.Bool("compressed", false), @@ -118,8 +124,9 @@ func (c *cachedSeekable) openReaderUncompressed(ctx context.Context, off, length if err == nil { recordCacheRead(ctx, true, length, cacheTypeSeekable, cacheOpOpenRangeReader) timer.Success(ctx, length) + readCache.Add(ctx, 1, CacheHitAttrs(c.objType, SourceNFS, CompressionNone)) - return withNFSGauge(ctx, newSectionReader(fp, 0, length)), nil + return withNFSGauge(ctx, newSectionReader(fp, 0, length)), SourceNFS, nil } if !os.IsNotExist(err) { @@ -127,20 +134,21 @@ func (c *cachedSeekable) openReaderUncompressed(ctx context.Context, off, length } timer.Failure(ctx, 0) + readCache.Add(ctx, 1, CacheMissAttrs(c.objType, SourceNFS, CompressionNone)) - rc, err := c.inner.OpenRangeReader(ctx, off, length, nil) + rc, innerSource, err := c.inner.OpenRangeReader(ctx, off, length, nil) if err != nil { - return nil, fmt.Errorf("failed to open inner range reader: %w", err) + return nil, innerSource, fmt.Errorf("failed to open inner range reader: %w", err) } recordCacheRead(ctx, false, length, cacheTypeSeekable, cacheOpOpenRangeReader) if !skipCacheWriteback(ctx) { rc = newCaptureReader(rc, int(length), false, - c.uncompressedChunkWriteback(chunkPath, off, length)) + c.uncompressedChunkWriteback(chunkPath, off, length, innerSource)) } - return rc, nil + return rc, innerSource, nil } // nfsGaugeReadCloser wraps a reader and decrements the NFS concurrent reads @@ -149,7 +157,7 @@ type nfsGaugeReadCloser struct { RangeReader } -func (r *nfsGaugeReadCloser) Close(ctx context.Context) error { +func (r *nfsGaugeReadCloser) Close(ctx context.Context) (*ReadStats, error) { nfsCacheConcurrentReads.Add(ctx, -1) return r.RangeReader.Close(ctx) @@ -165,7 +173,7 @@ func withNFSGauge(ctx context.Context, rc RangeReader) RangeReader { // the captured chunk to the NFS cache in a detached goroutine. Best-effort: // a short capture (e.g. upstream truncation) is dropped silently — a streaming // reader always ends in EOF, so byte count is the only reliable signal. -func (c *cachedSeekable) uncompressedChunkWriteback(chunkPath string, off, expectedLen int64) func(context.Context, []byte) { +func (c *cachedSeekable) uncompressedChunkWriteback(chunkPath string, off, expectedLen int64, src Source) func(context.Context, []byte) { return func(ctx context.Context, captured []byte) { if !isCompleteRead(len(captured), int(expectedLen), nil) { return @@ -175,7 +183,10 @@ func (c *cachedSeekable) uncompressedChunkWriteback(chunkPath string, off, expec ctx, span := c.tracer.Start(ctx, "write range reader chunk back to cache") defer span.End() + start := time.Now() err := c.writeToCache(ctx, off, chunkPath, captured) + recordWriteback(ctx, time.Since(start), int64(len(captured)), c.objType, src, CompressionNone, err) + if err != nil { recordError(span, err) recordCacheWriteError(ctx, cacheTypeSeekable, cacheOpOpenRangeReader, err) diff --git a/packages/shared/pkg/storage/storage_cache_seekable_compressed.go b/packages/shared/pkg/storage/storage_cache_seekable_compressed.go index 2e8077c4bb..05f8e0dc51 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable_compressed.go +++ b/packages/shared/pkg/storage/storage_cache_seekable_compressed.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "os" + "time" "go.opentelemetry.io/otel/attribute" ) @@ -17,10 +18,10 @@ var compressedCacheReadAttrs = []attribute.KeyValue{ // openReaderCompressed handles the compressed cache path for OpenRangeReader. // NFS stores compressed frames (.frm); on hit we decompress, on miss we fetch // raw compressed bytes and tee them to NFS on Close. -func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64, frameTable *FrameTable) (RangeReader, error) { +func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64, frameTable *FrameTable) (RangeReader, Source, error) { r, err := frameTable.LocateCompressed(offsetU) if err != nil { - return nil, fmt.Errorf("frame lookup for offset %d: %w", offsetU, err) + return nil, UnknownSource, fmt.Errorf("frame lookup for offset %d: %w", offsetU, err) } path := makeFrameFilename(c.path, r) @@ -36,14 +37,15 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64 recordCacheRead(ctx, true, int64(r.Length), cacheTypeSeekable, cacheOpOpenRangeReader) timer.Success(ctx, int64(r.Length)) - dec, err := NewDecompressingReader(NewRangeReader(f), ct) + dec, err := newDecompressReader(NewRangeReader(f), ct, SourceNFS, c.objType) if err != nil { f.Close() - return nil, fmt.Errorf("decompress cached frame: %w", err) + return nil, SourceNFS, fmt.Errorf("decompress cached frame: %w", err) } + readCache.Add(ctx, 1, CacheHitAttrs(c.objType, SourceNFS, ct)) - return withNFSGauge(ctx, dec), nil + return withNFSGauge(ctx, dec), SourceNFS, nil case statErr == nil: // Confirmed size mismatch: drop the file so the miss path rewrites it. f.Close() @@ -60,36 +62,37 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64 } timer.Failure(ctx, 0) + readCache.Add(ctx, 1, CacheMissAttrs(c.objType, SourceNFS, ct)) // Cache miss: fetch raw compressed bytes via OpenRangeReader(nil frameTable). - raw, err := c.inner.OpenRangeReader(ctx, r.Offset, int64(r.Length), nil) + raw, innerSource, err := c.inner.OpenRangeReader(ctx, r.Offset, int64(r.Length), nil) if err != nil { - return nil, fmt.Errorf("raw fetch at C=%d: %w", r.Offset, err) + return nil, innerSource, fmt.Errorf("raw fetch at C=%d: %w", r.Offset, err) } recordCacheRead(ctx, false, int64(r.Length), cacheTypeSeekable, cacheOpOpenRangeReader) - in := raw + src := raw if !skipCacheWriteback(ctx) { - in = newCaptureReader(raw, r.Length, true, - c.compressedFrameWriteback(path, offsetU, r.Length)) + src = newCaptureReader(raw, r.Length, true, + c.compressedFrameWriteback(path, offsetU, r.Length, innerSource, ct)) } - dec, err := NewDecompressingReader(in, ct) + dec, err := newDecompressReader(src, ct, innerSource, c.objType) if err != nil { raw.Close(ctx) - return nil, fmt.Errorf("create decompressor: %w", err) + return nil, innerSource, fmt.Errorf("create decompressor: %w", err) } - return dec, nil + return dec, innerSource, nil } // compressedFrameWriteback returns a captureReader callback that // persists the captured frame to the NFS cache in a detached goroutine. // Best-effort: a short capture is logged and skipped — the caller already // got valid decompressed bytes. -func (c *cachedSeekable) compressedFrameWriteback(framePath string, offset int64, expectedSize int) func(context.Context, []byte) { +func (c *cachedSeekable) compressedFrameWriteback(framePath string, offset int64, expectedSize int, src Source, codec CompressionType) func(context.Context, []byte) { return func(ctx context.Context, frame []byte) { if !isCompleteRead(len(frame), expectedSize, nil) { recordCacheWriteError(ctx, cacheTypeSeekable, cacheOpOpenRangeReader, @@ -102,7 +105,10 @@ func (c *cachedSeekable) compressedFrameWriteback(framePath string, offset int64 ctx, span := c.tracer.Start(ctx, "write compressed frame back to cache") defer span.End() + start := time.Now() err := c.writeToCache(ctx, offset, framePath, frame) + recordWriteback(ctx, time.Since(start), int64(len(frame)), c.objType, src, codec, err) + if err != nil { recordError(span, err) recordCacheWriteError(ctx, cacheTypeSeekable, cacheOpOpenRangeReader, err) diff --git a/packages/shared/pkg/storage/storage_cache_seekable_test.go b/packages/shared/pkg/storage/storage_cache_seekable_test.go index 81f7364c0a..56cbb0942a 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable_test.go +++ b/packages/shared/pkg/storage/storage_cache_seekable_test.go @@ -14,28 +14,49 @@ import ( "github.com/stretchr/testify/require" ) -// mustClose closes a RangeReader and asserts no error. -func mustClose(t *testing.T, rc RangeReader) { +// testCache returns a cachedSeekable with a fresh temp dir, the given chunk +// size, and a MockSeekable for inner. Cast c.inner to *MockSeekable to attach +// EXPECTs; for cache-only tests, leave it untouched (mockery only flags +// unexpected calls, not unused mocks). +func testCache(t *testing.T, chunkSize int64) cachedSeekable { t.Helper() - require.NoError(t, rc.Close(t.Context())) + + return cachedSeekable{ + path: t.TempDir(), + chunkSize: chunkSize, + inner: NewMockSeekable(t), + tracer: noopTracer, + } } -// bytesRangeReader wraps an in-memory byte slice as a RangeReader for tests. -func bytesRangeReader(b []byte) RangeReader { - return NewRangeReader(io.NopCloser(bytes.NewReader(b))) +// testCacheMock is testCache plus the inner mock already cast, for tests that +// attach EXPECTs. The returned cachedSeekable shares c.inner with the mock, so +// callers should not reassign c.inner before setting expectations. +func testCacheMock(t *testing.T, chunkSize int64) (c cachedSeekable, m *MockSeekable) { + t.Helper() + + m = NewMockSeekable(t) + c = cachedSeekable{ + path: t.TempDir(), + chunkSize: chunkSize, + inner: m, + tracer: noopTracer, + } + + return } // testReadAt emulates the removed cachedSeekable.ReadAt via OpenRangeReader. // This preserves the base test structure after ReadAt was removed from the Seekable interface. func testReadAt(ctx context.Context, c *cachedSeekable, buff []byte, off int64) (int, error) { - rc, err := c.OpenRangeReader(ctx, off, int64(len(buff)), nil) + rc, _, err := c.OpenRangeReader(ctx, off, int64(len(buff)), nil) if err != nil { return 0, err } n, err := io.ReadFull(rc, buff) - closeErr := rc.Close(ctx) + _, closeErr := rc.Close(ctx) if errors.Is(err, io.ErrUnexpectedEOF) { err = io.EOF } @@ -47,6 +68,18 @@ func testReadAt(ctx context.Context, c *cachedSeekable, buff []byte, off int64) return n, err } +// mustClose closes a RangeReader and asserts no error, discarding the stats. +func mustClose(t *testing.T, rc RangeReader) { + t.Helper() + _, err := rc.Close(t.Context()) + require.NoError(t, err) +} + +// bytesRangeReader wraps an in-memory byte slice as a RangeReader for tests. +func bytesRangeReader(b []byte) RangeReader { + return NewRangeReader(io.NopCloser(bytes.NewReader(b))) +} + func TestCachedFileObjectProvider_MakeChunkFilename(t *testing.T) { t.Parallel() @@ -63,10 +96,8 @@ func TestCachedFileObjectProvider_Size(t *testing.T) { const expectedSize int64 = 1024 - inner := NewMockSeekable(t) - inner.EXPECT().Size(mock.Anything).Return(expectedSize, nil) - - c := cachedSeekable{path: t.TempDir(), inner: inner, tracer: noopTracer} + c, inner := testCacheMock(t, 10) + inner.EXPECT().Size(mock.Anything).Return(expectedSize, nil).Once() // first call will write to cache size, err := c.Size(t.Context()) @@ -91,38 +122,33 @@ func TestCachedFileObjectProvider_WriteFromFileSystem(t *testing.T) { t.Run("can be cached successfully", func(t *testing.T) { t.Parallel() - tempDir := t.TempDir() - cacheDir := filepath.Join(tempDir, "cache") - tempFilename := filepath.Join(tempDir, "temp.bin") + tempFilename := filepath.Join(t.TempDir(), "temp.bin") data := []byte("hello world") - err := os.MkdirAll(cacheDir, os.ModePerm) - require.NoError(t, err) - - err = os.WriteFile(tempFilename, data, 0o644) - require.NoError(t, err) + require.NoError(t, os.WriteFile(tempFilename, data, 0o644)) - inner := NewMockSeekable(t) + c, inner := testCacheMock(t, 1024) inner.EXPECT(). StoreFile(mock.Anything, mock.Anything). - Return(nil, [32]byte{}, nil) + Return(nil, [32]byte{}, nil). + Once() featureFlags := NewMockFeatureFlagsClient(t) featureFlags.EXPECT().BoolFlag(mock.Anything, mock.Anything).Return(true) featureFlags.EXPECT().IntFlag(mock.Anything, mock.Anything).Return(10) - - c := cachedSeekable{path: cacheDir, inner: inner, chunkSize: 1024, flags: featureFlags, tracer: noopTracer} + c.flags = featureFlags // write temp file - _, _, err = c.StoreFile(t.Context(), tempFilename) + _, _, err := c.StoreFile(t.Context(), tempFilename) require.NoError(t, err) // file is written asynchronously, wait for it to finish c.wg.Wait() + // Remaining reads should come from cache only. c.inner = nil + c.flags = nil - // size should be cached size, err := c.Size(t.Context()) require.NoError(t, err) assert.Equal(t, int64(len(data)), size) @@ -142,18 +168,12 @@ func TestCachedFileObjectProvider_WriteTo(t *testing.T) { t.Run("read from cache when the file exists", func(t *testing.T) { t.Parallel() - tempDir := t.TempDir() - - tempPath := filepath.Join(tempDir, "a", "b", "c") - c := cachedSeekable{path: tempPath, chunkSize: 3, tracer: noopTracer} + c := testCache(t, 3) // create cache file cacheFilename := c.makeChunkFilename(0) - dirName := filepath.Dir(cacheFilename) - err := os.MkdirAll(dirName, 0o755) - require.NoError(t, err) - err = os.WriteFile(cacheFilename, []byte{1, 2, 3}, 0o600) - require.NoError(t, err) + require.NoError(t, os.MkdirAll(filepath.Dir(cacheFilename), 0o755)) + require.NoError(t, os.WriteFile(cacheFilename, []byte{1, 2, 3}, 0o600)) buffer := make([]byte, 3) read, err := testReadAt(t.Context(), &c, buffer, 0) @@ -165,9 +185,7 @@ func TestCachedFileObjectProvider_WriteTo(t *testing.T) { t.Run("short cache file returns EOF via ReadAt", func(t *testing.T) { t.Parallel() - tempDir := t.TempDir() - - c := cachedSeekable{path: tempDir, chunkSize: 10, tracer: noopTracer} + c := testCache(t, 10) // Plant a 3-byte cache file (valid last chunk). chunkPath := c.makeChunkFilename(0) @@ -188,23 +206,15 @@ func TestCachedFileObjectProvider_WriteTo(t *testing.T) { t.Parallel() fakeData := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} - inner := NewMockSeekable(t) + c, inner := testCacheMock(t, 3) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - RunAndReturn(func(_ context.Context, off int64, length int64, _ *FrameTable) (RangeReader, error) { + RunAndReturn(func(_ context.Context, off int64, length int64, _ *FrameTable) (RangeReader, Source, error) { end := min(int(off)+int(length), len(fakeData)) - return NewRangeReader(io.NopCloser(bytes.NewReader(fakeData[off:end]))), nil - }) - - tempDir := t.TempDir() - c := cachedSeekable{ - path: tempDir, - chunkSize: 3, - inner: inner, - tracer: noopTracer, - } + return bytesRangeReader(fakeData[off:end]), SourceFS, nil + }).Once() // first read goes to source buffer := make([]byte, 3) @@ -218,6 +228,7 @@ func TestCachedFileObjectProvider_WriteTo(t *testing.T) { // second read pulls from cache c.inner = nil // prevent remote reads, force cache read + buffer = make([]byte, 3) read, err = testReadAt(t.Context(), &c, buffer, 3) require.NoError(t, err) @@ -323,18 +334,10 @@ func TestCachedSeekableObjectProvider_ReadAt(t *testing.T) { t.Run("zero byte read with EOF is not cached", func(t *testing.T) { t.Parallel() - tempDir := t.TempDir() - inner := NewMockSeekable(t) + c, inner := testCacheMock(t, 10) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - Return(NewRangeReader(io.NopCloser(bytes.NewReader(nil))), nil) - - c := cachedSeekable{ - path: tempDir, - chunkSize: 10, - inner: inner, - tracer: noopTracer, - } + Return(bytesRangeReader(nil), SourceFS, nil) buff := make([]byte, 10) count, err := testReadAt(t.Context(), &c, buff, 0) @@ -351,19 +354,12 @@ func TestCachedSeekableObjectProvider_ReadAt(t *testing.T) { t.Run("full read without EOF is cached", func(t *testing.T) { t.Parallel() - tempDir := t.TempDir() data := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} - inner := NewMockSeekable(t) + + c, inner := testCacheMock(t, 10) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - Return(NewRangeReader(io.NopCloser(bytes.NewReader(data))), nil) - - c := cachedSeekable{ - path: tempDir, - chunkSize: 10, - inner: inner, - tracer: noopTracer, - } + Return(bytesRangeReader(data), SourceFS, nil) buff := make([]byte, 10) count, err := testReadAt(t.Context(), &c, buff, 0) @@ -414,18 +410,10 @@ func TestCachedSeekable_ReadAt_PreservesEOF(t *testing.T) { t.Run("EOF from inner is returned to caller unchanged", func(t *testing.T) { t.Parallel() - tempDir := t.TempDir() - inner := NewMockSeekable(t) + c, inner := testCacheMock(t, 10) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - Return(NewRangeReader(io.NopCloser(bytes.NewReader([]byte{1, 2, 3}))), nil) - - c := cachedSeekable{ - path: tempDir, - chunkSize: 10, - inner: inner, - tracer: noopTracer, - } + Return(bytesRangeReader([]byte{1, 2, 3}), SourceFS, nil) buff := make([]byte, 10) n, err := testReadAt(t.Context(), &c, buff, 0) @@ -438,18 +426,10 @@ func TestCachedSeekable_ReadAt_PreservesEOF(t *testing.T) { t.Run("nil error from inner is returned to caller unchanged", func(t *testing.T) { t.Parallel() - tempDir := t.TempDir() - inner := NewMockSeekable(t) + c, inner := testCacheMock(t, 10) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - Return(NewRangeReader(io.NopCloser(bytes.NewReader([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}))), nil) - - c := cachedSeekable{ - path: tempDir, - chunkSize: 10, - inner: inner, - tracer: noopTracer, - } + Return(bytesRangeReader([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}), SourceFS, nil) buff := make([]byte, 10) n, err := testReadAt(t.Context(), &c, buff, 0) @@ -463,22 +443,15 @@ func TestCachedSeekable_ReadAt_PreservesEOF(t *testing.T) { func TestCachedSeekable_ReadAt_SkipCacheWriteback(t *testing.T) { t.Parallel() - tempDir := t.TempDir() data := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} - inner := NewMockSeekable(t) + + c, inner := testCacheMock(t, 10) inner.EXPECT(). OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, (*FrameTable)(nil)). - RunAndReturn(func(_ context.Context, _ int64, _ int64, _ *FrameTable) (RangeReader, error) { - return NewRangeReader(io.NopCloser(bytes.NewReader(data))), nil + RunAndReturn(func(_ context.Context, _ int64, _ int64, _ *FrameTable) (RangeReader, Source, error) { + return bytesRangeReader(data), SourceFS, nil }) - c := cachedSeekable{ - path: tempDir, - chunkSize: 10, - inner: inner, - tracer: noopTracer, - } - ctx := WithSkipCacheWriteback(t.Context()) buff := make([]byte, 10) n, err := testReadAt(ctx, &c, buff, 0) @@ -498,74 +471,59 @@ func TestCachedSeekable_OpenRangeReader(t *testing.T) { t.Run("cache miss then full read populates cache for next call", func(t *testing.T) { t.Parallel() - tempDir := t.TempDir() data := []byte("hello") - inner := NewMockSeekable(t) + c, inner := testCacheMock(t, 10) inner.EXPECT(). OpenRangeReader(mock.Anything, int64(0), int64(len(data)), (*FrameTable)(nil)). - Return(NewRangeReader(io.NopCloser(bytes.NewReader(data))), nil). + Return(bytesRangeReader(data), SourceFS, nil). Once() - c := cachedSeekable{ - path: tempDir, - chunkSize: 10, - inner: inner, - tracer: noopTracer, - } - // First call: cache miss, reads from inner. - rc, err := c.OpenRangeReader(t.Context(), 0, int64(len(data)), nil) + rc, _, err := c.OpenRangeReader(t.Context(), 0, int64(len(data)), nil) require.NoError(t, err) got, err := io.ReadAll(rc) require.NoError(t, err) assert.Equal(t, data, got) - require.NoError(t, rc.Close(t.Context())) + mustClose(t, rc) c.wg.Wait() // Second call: should serve from NFS cache, inner not called again. c.inner = nil - rc2, err := c.OpenRangeReader(t.Context(), 0, int64(len(data)), nil) + + rc2, _, err := c.OpenRangeReader(t.Context(), 0, int64(len(data)), nil) require.NoError(t, err) got2, err := io.ReadAll(rc2) require.NoError(t, err) assert.Equal(t, data, got2) - require.NoError(t, rc2.Close(t.Context())) + mustClose(t, rc2) }) t.Run("skip cache writeback returns inner directly", func(t *testing.T) { t.Parallel() - tempDir := t.TempDir() data := []byte("hello") - inner := NewMockSeekable(t) + c, inner := testCacheMock(t, 10) inner.EXPECT(). OpenRangeReader(mock.Anything, int64(0), int64(len(data)), (*FrameTable)(nil)). - RunAndReturn(func(_ context.Context, _ int64, _ int64, _ *FrameTable) (RangeReader, error) { - return NewRangeReader(io.NopCloser(bytes.NewReader(data))), nil + RunAndReturn(func(_ context.Context, _ int64, _ int64, _ *FrameTable) (RangeReader, Source, error) { + return bytesRangeReader(data), SourceFS, nil }). Times(2) - c := cachedSeekable{ - path: tempDir, - chunkSize: 10, - inner: inner, - tracer: noopTracer, - } - ctx := WithSkipCacheWriteback(t.Context()) - rc, err := c.OpenRangeReader(ctx, 0, int64(len(data)), nil) + rc, _, err := c.OpenRangeReader(ctx, 0, int64(len(data)), nil) require.NoError(t, err) got, err := io.ReadAll(rc) require.NoError(t, err) assert.Equal(t, data, got) - require.NoError(t, rc.Close(ctx)) + mustClose(t, rc) c.wg.Wait() @@ -574,39 +532,30 @@ func TestCachedSeekable_OpenRangeReader(t *testing.T) { _, err = os.Stat(chunkPath) assert.True(t, os.IsNotExist(err), "skip writeback should not populate cache") - rc2, err := c.OpenRangeReader(ctx, 0, int64(len(data)), nil) + rc2, _, err := c.OpenRangeReader(ctx, 0, int64(len(data)), nil) require.NoError(t, err) got2, err := io.ReadAll(rc2) require.NoError(t, err) assert.Equal(t, data, got2) - require.NoError(t, rc2.Close(ctx)) + mustClose(t, rc2) }) t.Run("truncated inner read does not populate cache", func(t *testing.T) { t.Parallel() - tempDir := t.TempDir() - - inner := NewMockSeekable(t) + c, inner := testCacheMock(t, 10) inner.EXPECT(). OpenRangeReader(mock.Anything, int64(0), int64(5), (*FrameTable)(nil)). - Return(NewRangeReader(io.NopCloser(bytes.NewReader([]byte{0xAA, 0xBB}))), nil) - - c := cachedSeekable{ - path: tempDir, - chunkSize: 10, - inner: inner, - tracer: noopTracer, - } + Return(bytesRangeReader([]byte{0xAA, 0xBB}), SourceFS, nil) - rc, err := c.OpenRangeReader(t.Context(), 0, 5, nil) + rc, _, err := c.OpenRangeReader(t.Context(), 0, 5, nil) require.NoError(t, err) got, err := io.ReadAll(rc) require.NoError(t, err) assert.Equal(t, []byte{0xAA, 0xBB}, got) - require.NoError(t, rc.Close(t.Context())) + mustClose(t, rc) c.wg.Wait() @@ -677,31 +626,21 @@ func TestCachedSeekable_StoreFile_Compressed_WriteThrough(t *testing.T) { func TestCacheWriteThroughReader(t *testing.T) { t.Parallel() - newTestCache := func(t *testing.T) cachedSeekable { - t.Helper() - - return cachedSeekable{ - path: t.TempDir(), - chunkSize: 10, - tracer: noopTracer, - } - } - t.Run("complete read is cached", func(t *testing.T) { t.Parallel() - c := newTestCache(t) + c := testCache(t, 10) data := []byte("hello") - inner := NewRangeReader(io.NopCloser(bytes.NewReader(data))) + inner := bytesRangeReader(data) r := newCaptureReader(inner, len(data), false, - c.uncompressedChunkWriteback(c.makeChunkFilename(0), 0, int64(len(data)))) + c.uncompressedChunkWriteback(c.makeChunkFilename(0), 0, int64(len(data)), SourceFS)) got, err := io.ReadAll(r) require.NoError(t, err) assert.Equal(t, data, got) - require.NoError(t, r.Close(t.Context())) + mustClose(t, r) c.wg.Wait() cached, err := os.ReadFile(c.makeChunkFilename(0)) @@ -712,20 +651,20 @@ func TestCacheWriteThroughReader(t *testing.T) { t.Run("truncated upstream fully consumed is not cached", func(t *testing.T) { t.Parallel() - c := newTestCache(t) + c := testCache(t, 10) // Inner has only 2 bytes but expectedLen is 5. The reader is // fully consumed (EOF is reached), yet the total doesn't match // the expected length so it must not be cached. - inner := NewRangeReader(io.NopCloser(bytes.NewReader([]byte{0xAA, 0xBB}))) + inner := bytesRangeReader([]byte{0xAA, 0xBB}) r := newCaptureReader(inner, 5, false, - c.uncompressedChunkWriteback(c.makeChunkFilename(0), 0, 5)) + c.uncompressedChunkWriteback(c.makeChunkFilename(0), 0, 5, SourceFS)) got, err := io.ReadAll(r) require.NoError(t, err) assert.Equal(t, []byte{0xAA, 0xBB}, got) - require.NoError(t, r.Close(t.Context())) + mustClose(t, r) c.wg.Wait() _, err = os.Stat(c.makeChunkFilename(0)) @@ -735,12 +674,12 @@ func TestCacheWriteThroughReader(t *testing.T) { t.Run("partially consumed reader closed early is not cached", func(t *testing.T) { t.Parallel() - c := newTestCache(t) + c := testCache(t, 10) data := []byte("hello") - inner := NewRangeReader(io.NopCloser(bytes.NewReader(data))) + inner := bytesRangeReader(data) r := newCaptureReader(inner, len(data), false, - c.uncompressedChunkWriteback(c.makeChunkFilename(0), 0, int64(len(data)))) + c.uncompressedChunkWriteback(c.makeChunkFilename(0), 0, int64(len(data)), SourceFS)) // Read only 2 of 5 bytes, then close without reaching EOF. buf := make([]byte, 2) @@ -748,7 +687,7 @@ func TestCacheWriteThroughReader(t *testing.T) { require.NoError(t, err) assert.Equal(t, 2, n) - require.NoError(t, r.Close(t.Context())) + mustClose(t, r) c.wg.Wait() _, err = os.Stat(c.makeChunkFilename(0)) diff --git a/packages/shared/pkg/storage/storage_fs.go b/packages/shared/pkg/storage/storage_fs.go index fcaac07366..fa6a6d817d 100644 --- a/packages/shared/pkg/storage/storage_fs.go +++ b/packages/shared/pkg/storage/storage_fs.go @@ -30,7 +30,8 @@ type fsStorage struct { var _ StorageProvider = (*fsStorage)(nil) type fsObject struct { - path string + path string + objType SeekableObjectType } var ( @@ -71,14 +72,15 @@ func (s *fsStorage) UploadSignedURL(_ context.Context, path string, ttl time.Dur return u, nil } -func (s *fsStorage) OpenSeekable(_ context.Context, path string, _ SeekableObjectType) (Seekable, error) { +func (s *fsStorage) OpenSeekable(_ context.Context, path string, objectType SeekableObjectType) (Seekable, error) { dir := filepath.Dir(s.getPath(path)) if err := os.MkdirAll(dir, 0o755); err != nil { return nil, err } return &fsObject{ - path: s.getPath(path), + path: s.getPath(path), + objType: objectType, }, nil } @@ -303,27 +305,32 @@ func (u *fsPartUploader) Complete(_ context.Context) error { return os.WriteFile(u.fullPath, u.Assemble(), 0o644) } -func (o *fsObject) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, error) { +func (o *fsObject) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, Source, error) { if frameTable.IsCompressed() { r, err := frameTable.LocateCompressed(offsetU) if err != nil { - return nil, fmt.Errorf("get frame for offset %d, FS:%s: %w", offsetU, o.path, err) + return nil, SourceFS, fmt.Errorf("get frame for offset %d, FS:%s: %w", offsetU, o.path, err) } raw, err := o.openRangeReader(ctx, r.Offset, int64(r.Length)) if err != nil { - return nil, err + return nil, SourceFS, err } - dec, err := NewDecompressingReader(raw, frameTable.CompressionType()) + dec, err := newDecompressReader(raw, frameTable.CompressionType(), SourceFS, o.objType) if err != nil { raw.Close(ctx) - return nil, err + return nil, SourceFS, err } - return dec, nil + return dec, SourceFS, nil + } + + raw, err := o.openRangeReader(ctx, offsetU, length) + if err != nil { + return nil, SourceFS, err } - return o.openRangeReader(ctx, offsetU, length) + return newObservableReader(raw), SourceFS, nil } diff --git a/packages/shared/pkg/storage/storage_google.go b/packages/shared/pkg/storage/storage_google.go index e1fe7c6dc9..e0d3fa254a 100644 --- a/packages/shared/pkg/storage/storage_google.go +++ b/packages/shared/pkg/storage/storage_google.go @@ -91,6 +91,7 @@ type gcpObject struct { storage *gcpStorage path string handle *storage.ObjectHandle + objType SeekableObjectType limiter *limit.Limiter } @@ -173,7 +174,7 @@ func (s *gcpStorage) UploadSignedURL(_ context.Context, path string, ttl time.Du return url, nil } -func (s *gcpStorage) OpenSeekable(_ context.Context, path string, _ SeekableObjectType) (Seekable, error) { +func (s *gcpStorage) OpenSeekable(_ context.Context, path string, objectType SeekableObjectType) (Seekable, error) { handle := s.bucket.Object(path).Retryer( storage.WithMaxAttempts(googleMaxAttempts), storage.WithPolicy(storage.RetryAlways), @@ -190,6 +191,7 @@ func (s *gcpStorage) OpenSeekable(_ context.Context, path string, _ SeekableObje storage: s, path: path, handle: handle, + objType: objectType, limiter: s.limiter, }, nil @@ -316,12 +318,12 @@ func (r *idleTimeoutReader) Read(p []byte) (int, error) { return n, err } -func (r *idleTimeoutReader) Close(ctx context.Context) error { +func (r *idleTimeoutReader) Close(ctx context.Context) (*ReadStats, error) { r.timer.Stop() defer r.cancel() gcsConcurrentReads.Add(ctx, -1) - return r.ReadCloser.Close() + return nil, r.ReadCloser.Close() } func (o *gcpObject) Put(ctx context.Context, data []byte, opts ...PutOption) error { @@ -602,7 +604,7 @@ func parseServiceAccountBase64(serviceAccount string) (*gcpServiceToken, error) return &sa, nil } -func (o *gcpObject) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, error) { +func (o *gcpObject) OpenRangeReader(ctx context.Context, offsetU int64, length int64, frameTable *FrameTable) (RangeReader, Source, error) { timer := googleReadTimerFactory.Begin(attribute.String(gcsOperationAttr, gcsOperationAttrReadAt)) if !frameTable.IsCompressed() { @@ -610,35 +612,37 @@ func (o *gcpObject) OpenRangeReader(ctx context.Context, offsetU int64, length i if err != nil { timer.Failure(ctx, 0) - return nil, err + return nil, SourceGCS, err } - return newObservableReader(rc, timer, nil), nil + return newObservableReader(rc). + withTimer(timer), SourceGCS, nil } r, err := frameTable.LocateCompressed(offsetU) if err != nil { timer.Failure(ctx, 0) - return nil, fmt.Errorf("get frame for offset %d, GCS:%s: %w", offsetU, o.path, err) + return nil, SourceGCS, fmt.Errorf("get frame for offset %d, GCS:%s: %w", offsetU, o.path, err) } raw, err := o.openRangeReader(ctx, r.Offset, int64(r.Length)) if err != nil { timer.Failure(ctx, 0) - return nil, err + return nil, SourceGCS, err } - dec, err := NewDecompressingReader(raw, frameTable.CompressionType()) + dec, err := newDecompressReader(raw, frameTable.CompressionType(), SourceGCS, o.objType) if err != nil { raw.Close(ctx) timer.Failure(ctx, 0) - return nil, err + return nil, SourceGCS, err } - return newObservableReader(dec, timer, nil), nil + return newObservableReader(dec). + withTimer(timer), SourceGCS, nil } func isResourceExhausted(err error) bool { diff --git a/packages/shared/pkg/telemetry/meters.go b/packages/shared/pkg/telemetry/meters.go index 477a9a879c..4736b3cbca 100644 --- a/packages/shared/pkg/telemetry/meters.go +++ b/packages/shared/pkg/telemetry/meters.go @@ -531,6 +531,48 @@ func NewTimerFactory( return TimerFactory{duration, bytes, count}, nil } +// FloatTimerFactory records duration as fractional milliseconds so +// sub-millisecond operations aren't truncated to 0. Callers supply the +// duration via Record. +// +// The histogram and bytes counter are registered under distinct OTEL +// instrument names ( for the histogram, .size for the +// bytes counter) so Grafana's metric-metadata-driven unit detection isn't +// ambiguous. The histogram already emits a _count series, so no separate +// event-count instrument is needed. +type FloatTimerFactory struct { + duration metric.Float64Histogram + bytes metric.Int64Counter +} + +func NewFloatTimerFactory( + meter metric.Meter, + metricName, durationDescription, bytesDescription, _ string, +) (FloatTimerFactory, error) { + duration, err := meter.Float64Histogram(metricName, + metric.WithDescription(durationDescription), + metric.WithUnit("ms"), + ) + if err != nil { + return FloatTimerFactory{}, fmt.Errorf("failed to create duration histogram: %w", err) + } + + bytes, err := meter.Int64Counter(metricName+".size", + metric.WithDescription(bytesDescription), + metric.WithUnit("By"), + ) + if err != nil { + return FloatTimerFactory{}, fmt.Errorf("failed to create bytes counter: %w", err) + } + + return FloatTimerFactory{duration, bytes}, nil +} + +func (f *FloatTimerFactory) Record(ctx context.Context, dur time.Duration, total int64, attrs metric.MeasurementOption) { + f.duration.Record(ctx, float64(dur)/float64(time.Millisecond), attrs) + f.bytes.Add(ctx, total, attrs) +} + func (f *TimerFactory) Begin(kv ...attribute.KeyValue) *Stopwatch { return &Stopwatch{ histogram: f.duration, @@ -585,9 +627,9 @@ func PrecomputeAttrs(kv ...attribute.KeyValue) metric.MeasurementOption { // RecordRaw records an operation using a precomputed attribute option, it does // not include any previous attributes passed at Begin(). Zero-allocation // alternative to Success/Failure for hot paths. -func (t Stopwatch) RecordRaw(ctx context.Context, total int64, precomputedAttrs metric.MeasurementOption) { +func (t Stopwatch) RecordRaw(ctx context.Context, total int64, allAttrs metric.MeasurementOption) { amount := time.Since(t.start).Milliseconds() - t.histogram.Record(ctx, amount, precomputedAttrs) - t.sum.Add(ctx, total, precomputedAttrs) - t.count.Add(ctx, 1, precomputedAttrs) + t.histogram.Record(ctx, amount, allAttrs) + t.sum.Add(ctx, total, allAttrs) + t.count.Add(ctx, 1, allAttrs) } 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 a02fb9a01e..d5e11842b0 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 @@ -147,7 +147,7 @@ func verifyChecksum(t *testing.T, ctx context.Context, persistence storage.Stora obj, err := persistence.OpenSeekable(ctx, dataPath, objType) require.NoErrorf(t, err, "%s/%s: open data file %s", node.name, fileName, dataPath) - rc, err := obj.OpenRangeReader(ctx, 0, bd.Size, bd.FrameData) + rc, _, err := obj.OpenRangeReader(ctx, 0, bd.Size, bd.FrameData) require.NoErrorf(t, err, "%s/%s: open range reader", node.name, fileName) defer rc.Close(ctx) From ba1e232d8296f3beaa3e86941f1e192bb147ef59 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Tue, 16 Jun 2026 15:49:38 -0700 Subject: [PATCH 08/11] touchups --- .../pkg/sandbox/block/streaming_chunk.go | 9 ++++ packages/shared/pkg/storage/read_attrs.go | 17 ++----- .../pkg/storage/read_attrs_precomputed.go | 46 ------------------- .../shared/pkg/storage/read_attrs_test.go | 4 -- packages/shared/pkg/storage/read_metrics.go | 36 ++++++++++----- packages/shared/pkg/storage/storage.go | 1 - .../pkg/storage/storage_cache_seekable.go | 17 ++++--- .../storage_cache_seekable_compressed.go | 10 +++- 8 files changed, 57 insertions(+), 83 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go index be5c64da07..be3073ce21 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go @@ -186,6 +186,15 @@ func (c *Chunker) fetch(ctx context.Context, off, length int64, upstream storage startBlock := (off / blockSize) * blockSize endBlock := ((off + length - 1) / blockSize) * blockSize chunkEnd := chunkOff + chunkLen + + // If the session has already streamed past every byte + // we need, the data is in the mmap and we never have to wait. Report + // source=mmap. + endByte := min(endBlock+blockSize, chunkEnd) - chunkOff + if session.bytesReady.Load() >= endByte { + return storage.SourceMmap, nil + } + for b := startBlock; b <= endBlock; b += blockSize { if b >= chunkEnd { break // tail belongs to the caller's next chunk fetch. diff --git a/packages/shared/pkg/storage/read_attrs.go b/packages/shared/pkg/storage/read_attrs.go index 2c1559c12d..db0c1b4808 100644 --- a/packages/shared/pkg/storage/read_attrs.go +++ b/packages/shared/pkg/storage/read_attrs.go @@ -7,20 +7,13 @@ const ( AttrSource = "source" AttrCodec = "codec" AttrOutcome = "outcome" - AttrEvent = "event" AttrFileType = "file_type" ) const ( - OutcomeOK = "ok" - OutcomeErrCanceled = "err_canceled" - OutcomeErrIO = "err_io" - OutcomeErrTimeout = "err_timeout" -) - -const ( - CacheEventHit = "hit" - CacheEventMiss = "miss" - CacheEventWritebackOK = "writeback_ok" - CacheEventWritebackErr = "writeback_err" + OutcomeOK = "ok" + OutcomeErrCanceled = "err_canceled" + OutcomeErrIO = "err_io" + OutcomeErrTimeout = "err_timeout" + OutcomeTransitioned = "transitioned" ) diff --git a/packages/shared/pkg/storage/read_attrs_precomputed.go b/packages/shared/pkg/storage/read_attrs_precomputed.go index 08944bd0c8..efba50ccad 100644 --- a/packages/shared/pkg/storage/read_attrs_precomputed.go +++ b/packages/shared/pkg/storage/read_attrs_precomputed.go @@ -44,11 +44,6 @@ const numCodecs = 3 // CompressionNone, Zstd, LZ4 var ( tableOK [numSeekableObjectTypes][numSources][numCodecs]metric.MeasurementOption - tableCacheHit [numSeekableObjectTypes][numSources][numCodecs]metric.MeasurementOption - tableCacheMiss [numSeekableObjectTypes][numSources][numCodecs]metric.MeasurementOption - tableCacheWritebackOK [numSeekableObjectTypes][numSources][numCodecs]metric.MeasurementOption - tableCacheWritebackErr [numSeekableObjectTypes][numSources][numCodecs]metric.MeasurementOption - // keyed by file_type only: inflight is incremented before the source is // known (the OpenRangeReader call itself dominates GCS latency). tableInflightFetch [numSeekableObjectTypes]metric.MeasurementOption @@ -74,23 +69,6 @@ func init() { tableOK[ot][s][ct] = set( ftAttr, srcAttr, codecAttr, outcomeOK, ) - - tableCacheHit[ot][s][ct] = set( - ftAttr, attribute.String(AttrEvent, CacheEventHit), - srcAttr, codecAttr, - ) - tableCacheMiss[ot][s][ct] = set( - ftAttr, attribute.String(AttrEvent, CacheEventMiss), - srcAttr, codecAttr, - ) - tableCacheWritebackOK[ot][s][ct] = set( - ftAttr, attribute.String(AttrEvent, CacheEventWritebackOK), - srcAttr, codecAttr, - ) - tableCacheWritebackErr[ot][s][ct] = set( - ftAttr, attribute.String(AttrEvent, CacheEventWritebackErr), - srcAttr, codecAttr, - ) } } } @@ -116,30 +94,6 @@ func OKAttrs(o SeekableObjectType, s Source, c CompressionType) metric.Measureme return tableOK[o][s][c] } -func CacheHitAttrs(o SeekableObjectType, s Source, c CompressionType) metric.MeasurementOption { - o, s, c = safeAttrIdx(o, s, c) - - return tableCacheHit[o][s][c] -} - -func CacheMissAttrs(o SeekableObjectType, s Source, c CompressionType) metric.MeasurementOption { - o, s, c = safeAttrIdx(o, s, c) - - return tableCacheMiss[o][s][c] -} - -func CacheWritebackOKAttrs(o SeekableObjectType, s Source, c CompressionType) metric.MeasurementOption { - o, s, c = safeAttrIdx(o, s, c) - - return tableCacheWritebackOK[o][s][c] -} - -func CacheWritebackErrAttrs(o SeekableObjectType, s Source, c CompressionType) metric.MeasurementOption { - o, s, c = safeAttrIdx(o, s, c) - - return tableCacheWritebackErr[o][s][c] -} - func InflightFetchAttrs(o SeekableObjectType) metric.MeasurementOption { return tableInflightFetch[o] } diff --git a/packages/shared/pkg/storage/read_attrs_test.go b/packages/shared/pkg/storage/read_attrs_test.go index 7d1b549f5f..ee6e3008b8 100644 --- a/packages/shared/pkg/storage/read_attrs_test.go +++ b/packages/shared/pkg/storage/read_attrs_test.go @@ -47,10 +47,6 @@ func TestPrecomputedAttrsPopulated(t *testing.T) { for s := range numSources { for c := range CompressionType(numCodecs) { require.NotNil(t, OKAttrs(o, s, c)) - require.NotNil(t, CacheHitAttrs(o, s, c)) - require.NotNil(t, CacheMissAttrs(o, s, c)) - require.NotNil(t, CacheWritebackOKAttrs(o, s, c)) - require.NotNil(t, CacheWritebackErrAttrs(o, s, c)) } } } diff --git a/packages/shared/pkg/storage/read_metrics.go b/packages/shared/pkg/storage/read_metrics.go index ae88776442..b78b2b76e7 100644 --- a/packages/shared/pkg/storage/read_metrics.go +++ b/packages/shared/pkg/storage/read_metrics.go @@ -60,17 +60,17 @@ var ( "fetch / (open + read + decompress) — 1.0 = fetch wall fully explained by work, >1 = overhead", "1", ) - readCache = utils.Must(meter.Int64Counter( - "orchestrator.read.cache", - metric.WithDescription("NFS read-cache events (hit / miss / writeback). The mmap tier is orchestrator.chunk.cache."), - metric.WithUnit("1"), - )) - readInflight = utils.Must(meter.Int64UpDownCounter( "orchestrator.read.inflight", metric.WithDescription("In-flight read-path fetches (cache miss → backend), by file_type"), metric.WithUnit("1"), )) + + readWritebackContended = utils.Must(meter.Int64Counter( + "orchestrator.read.writeback.contended", + metric.WithDescription("Writebacks skipped because the NFS chunk lock was already held — another goroutine is writing the same chunk (normal cache dedup, not an error)"), + metric.WithUnit("1"), + )) ) func RecordReadOpen(ctx context.Context, dur time.Duration, bytes int64, attrs metric.MeasurementOption) { @@ -102,7 +102,11 @@ func StartInflight(ctx context.Context, attrs metric.MeasurementOption) func() { } // Outcome maps a read-path error to the closed read.* outcome enum. +// PeerTransitionedError is a routing signal — the peer told us to refresh +// the header and reopen against storage — so it gets its own bucket +// instead of polluting err_io with sub-ms "errors" at every transition. func Outcome(err error) string { + var transErr *PeerTransitionedError switch { case err == nil: return OutcomeOK @@ -110,6 +114,8 @@ func Outcome(err error) string { return OutcomeErrCanceled case errors.Is(err, context.DeadlineExceeded): return OutcomeErrTimeout + case errors.As(err, &transErr): + return OutcomeTransitioned default: return OutcomeErrIO } @@ -136,16 +142,24 @@ func recordDecompressStep(ctx context.Context, r *decompressReader, stats *ReadS readDecompress.Record(ctx, stats.Decompress, stats.UncompressedBytes, ErrAttrs(r.objType, r.source, r.ct, readErr)) } -// recordWriteback emits the read.writeback timer and its read.cache event. -// src is the originating fetch source (kept for cross-correlation); writebacks -// always target NFS. +// recordWritebackContended emits the read.writeback.contended counter — a +// writeback that didn't run because the NFS chunk lock was already held. +// Cold path; attrs built inline. +func recordWritebackContended(ctx context.Context, ot SeekableObjectType, src Source, ct CompressionType) { + readWritebackContended.Add(ctx, 1, metric.WithAttributes( + attribute.String(AttrFileType, ot.String()), + attribute.String(AttrSource, src.String()), + attribute.String(AttrCodec, ct.String()), + )) +} + +// recordWriteback emits the read.writeback timer. src is the originating +// fetch source (kept for cross-correlation); writebacks always target NFS. func recordWriteback(ctx context.Context, dur time.Duration, bytes int64, ot SeekableObjectType, src Source, ct CompressionType, err error) { if err == nil { readWriteback.Record(ctx, dur, bytes, OKAttrs(ot, src, ct)) - readCache.Add(ctx, 1, CacheWritebackOKAttrs(ot, SourceNFS, ct)) return } readWriteback.Record(ctx, dur, bytes, ErrAttrs(ot, src, ct, err)) - readCache.Add(ctx, 1, CacheWritebackErrAttrs(ot, SourceNFS, ct)) } diff --git a/packages/shared/pkg/storage/storage.go b/packages/shared/pkg/storage/storage.go index cf145515a2..bdbefe65fa 100644 --- a/packages/shared/pkg/storage/storage.go +++ b/packages/shared/pkg/storage/storage.go @@ -166,7 +166,6 @@ type ReadStats struct { Decompress time.Duration } - type RangeReader interface { io.Reader // Close returns stats from the reader's lifetime, or nil if the reader did diff --git a/packages/shared/pkg/storage/storage_cache_seekable.go b/packages/shared/pkg/storage/storage_cache_seekable.go index e8f42334be..eef5458e25 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable.go +++ b/packages/shared/pkg/storage/storage_cache_seekable.go @@ -124,7 +124,6 @@ func (c *cachedSeekable) openReaderUncompressed(ctx context.Context, off, length if err == nil { recordCacheRead(ctx, true, length, cacheTypeSeekable, cacheOpOpenRangeReader) timer.Success(ctx, length) - readCache.Add(ctx, 1, CacheHitAttrs(c.objType, SourceNFS, CompressionNone)) return withNFSGauge(ctx, newSectionReader(fp, 0, length)), SourceNFS, nil } @@ -134,7 +133,6 @@ func (c *cachedSeekable) openReaderUncompressed(ctx context.Context, off, length } timer.Failure(ctx, 0) - readCache.Add(ctx, 1, CacheMissAttrs(c.objType, SourceNFS, CompressionNone)) rc, innerSource, err := c.inner.OpenRangeReader(ctx, off, length, nil) if err != nil { @@ -185,6 +183,11 @@ func (c *cachedSeekable) uncompressedChunkWriteback(chunkPath string, off, expec start := time.Now() err := c.writeToCache(ctx, off, chunkPath, captured) + if errors.Is(err, lock.ErrLockAlreadyHeld) { + recordWritebackContended(ctx, c.objType, src, CompressionNone) + + return + } recordWriteback(ctx, time.Since(start), int64(len(captured)), c.objType, src, CompressionNone, err) if err != nil { @@ -366,15 +369,15 @@ func (c *cachedSeekable) validateReadParams(buffSize, offset int64) error { func (c *cachedSeekable) writeToCache(ctx context.Context, offset int64, finalPath string, bytes []byte) error { writeTimer := cacheSlabWriteTimerFactory.Begin() - // Try to acquire lock for this chunk write to NFS cache + // Try to acquire lock for this chunk write to NFS cache. + // Lock contention propagates as ErrLockAlreadyHeld; callers treat that as + // "not our work" and skip writeback metrics — counting it as success would + // inflate writeback bytes by the chunk size on every contention event. lockFile, err := lock.TryAcquireLock(ctx, finalPath) if err != nil { - // failed to acquire lock, which is a different category of failure than "write failed" - recordCacheWriteError(ctx, cacheTypeSeekable, cacheOpOpenRangeReader, err) - writeTimer.Failure(ctx, 0) - return nil + return err } // Release lock after write completes diff --git a/packages/shared/pkg/storage/storage_cache_seekable_compressed.go b/packages/shared/pkg/storage/storage_cache_seekable_compressed.go index 05f8e0dc51..621da9bd09 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable_compressed.go +++ b/packages/shared/pkg/storage/storage_cache_seekable_compressed.go @@ -2,11 +2,14 @@ package storage import ( "context" + "errors" "fmt" "os" "time" "go.opentelemetry.io/otel/attribute" + + "github.com/e2b-dev/infra/packages/shared/pkg/storage/lock" ) // Precomputed OTEL attributes for compressed cache reads (avoids per-read allocation). @@ -43,7 +46,6 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64 return nil, SourceNFS, fmt.Errorf("decompress cached frame: %w", err) } - readCache.Add(ctx, 1, CacheHitAttrs(c.objType, SourceNFS, ct)) return withNFSGauge(ctx, dec), SourceNFS, nil case statErr == nil: @@ -62,7 +64,6 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64 } timer.Failure(ctx, 0) - readCache.Add(ctx, 1, CacheMissAttrs(c.objType, SourceNFS, ct)) // Cache miss: fetch raw compressed bytes via OpenRangeReader(nil frameTable). raw, innerSource, err := c.inner.OpenRangeReader(ctx, r.Offset, int64(r.Length), nil) @@ -107,6 +108,11 @@ func (c *cachedSeekable) compressedFrameWriteback(framePath string, offset int64 start := time.Now() err := c.writeToCache(ctx, offset, framePath, frame) + if errors.Is(err, lock.ErrLockAlreadyHeld) { + recordWritebackContended(ctx, c.objType, src, codec) + + return + } recordWriteback(ctx, time.Since(start), int64(len(frame)), c.objType, src, codec, err) if err != nil { From 4d6a841c7630e16202858ce04034eb40e98e5425 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Wed, 17 Jun 2026 09:21:50 -0700 Subject: [PATCH 09/11] PR feedback: context.WithoutCancel --- packages/orchestrator/pkg/sandbox/block/streaming_chunk.go | 2 +- packages/orchestrator/pkg/sandbox/template/peerclient/blob.go | 2 +- .../tests/api/sandboxes/sandbox_rapid_pause_resume_test.go | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go index f7d9213b0c..5c05ff7255 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go @@ -241,7 +241,7 @@ func (c *Chunker) progressiveRead(ctx context.Context, s *fetchSession, mmapSlic return 0, fmt.Errorf("failed to open range reader at %d: %w", s.chunkOff, err) } defer func() { - if closeErr := reader.Close(ctx); closeErr != nil && err == nil { + if closeErr := reader.Close(context.WithoutCancel(ctx)); closeErr != nil && err == nil { err = closeErr } }() diff --git a/packages/orchestrator/pkg/sandbox/template/peerclient/blob.go b/packages/orchestrator/pkg/sandbox/template/peerclient/blob.go index c546f1467c..92067e6011 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerclient/blob.go +++ b/packages/orchestrator/pkg/sandbox/template/peerclient/blob.go @@ -65,7 +65,7 @@ func (b *peerBlob) WriteTo(ctx context.Context, dst io.Writer) (int64, error) { } reader := newPeerStreamReader(recv, cancel) - defer reader.Close(ctx) + defer reader.Close(context.WithoutCancel(ctx)) n, err := io.Copy(dst, reader) if err != nil { 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 a02fb9a01e..3b36793ff7 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 @@ -149,7 +149,7 @@ func verifyChecksum(t *testing.T, ctx context.Context, persistence storage.Stora rc, err := obj.OpenRangeReader(ctx, 0, bd.Size, bd.FrameData) require.NoErrorf(t, err, "%s/%s: open range reader", node.name, fileName) - defer rc.Close(ctx) + defer rc.Close(context.WithoutCancel(ctx)) hasher := sha256.New() n, err := io.Copy(hasher, rc) From 57aaaceb6363f158af4795a2eb87c396cb428e64 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Wed, 17 Jun 2026 10:03:20 -0700 Subject: [PATCH 10/11] PR feedback: unused param --- packages/orchestrator/pkg/sandbox/block/metrics/main.go | 1 - packages/orchestrator/pkg/sandbox/build/build.go | 2 +- packages/shared/pkg/storage/read_metrics.go | 5 ----- packages/shared/pkg/telemetry/meters.go | 2 +- 4 files changed, 2 insertions(+), 8 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/block/metrics/main.go b/packages/orchestrator/pkg/sandbox/block/metrics/main.go index 25da8b0b71..2ae23f3610 100644 --- a/packages/orchestrator/pkg/sandbox/block/metrics/main.go +++ b/packages/orchestrator/pkg/sandbox/block/metrics/main.go @@ -65,7 +65,6 @@ func NewMetrics(meterProvider metric.MeterProvider) (Metrics, error) { blocksMeter, orchestratorChunkSlice, "Time taken by Chunker to serve a Slice() (source=mmap when served from cache)", "Bytes returned", - "Slice call count", ); err != nil { return m, fmt.Errorf("error creating chunk slice timer factory: %w", err) } diff --git a/packages/orchestrator/pkg/sandbox/build/build.go b/packages/orchestrator/pkg/sandbox/build/build.go index 06da10b07f..5925f08510 100644 --- a/packages/orchestrator/pkg/sandbox/build/build.go +++ b/packages/orchestrator/pkg/sandbox/build/build.go @@ -49,7 +49,7 @@ var ( var fileReadAtTimer = utils.Must(telemetry.NewFloatTimerFactory(meter, "orchestrator.file.read_at", "Time to serve a build.File ReadAt across all source builds", - "Bytes read", "ReadAt call count")) + "Bytes read")) type File struct { header atomic.Pointer[header.Header] diff --git a/packages/shared/pkg/storage/read_metrics.go b/packages/shared/pkg/storage/read_metrics.go index b78b2b76e7..9de1f24ec9 100644 --- a/packages/shared/pkg/storage/read_metrics.go +++ b/packages/shared/pkg/storage/read_metrics.go @@ -27,31 +27,26 @@ var ( "orchestrator.read.open", "OpenRangeReader (open / TTFB) wall", "Bytes (always 0 — open transfers no payload)", - "Number of opens", )) readRead = utils.Must(telemetry.NewFloatTimerFactory(meter, "orchestrator.read.read", "Raw source-read wall (decompression excluded)", "Compressed/stored bytes read from the source", - "Number of source reads", )) readDecompress = utils.Must(telemetry.NewFloatTimerFactory(meter, "orchestrator.read.decompress", "Decompression CPU wall (decoder read time minus source transfer)", "Uncompressed bytes produced", - "Number of decompress records", )) readFetch = utils.Must(telemetry.NewFloatTimerFactory(meter, "orchestrator.read.fetch", "Total fetch wall — should ≈ open + read + decompress; any excess is overhead (see read.pipeline.efficiency)", "Bytes delivered to the app", - "Number of fetches", )) readWriteback = utils.Must(telemetry.NewFloatTimerFactory(meter, "orchestrator.read.writeback", "Async NFS cache writeback wall", "Bytes written to NFS", - "Number of writebacks", )) // fetch / (open + read + decompress); 1.0 = fully explained, >1 = overhead. diff --git a/packages/shared/pkg/telemetry/meters.go b/packages/shared/pkg/telemetry/meters.go index f502d27d19..227c629b1a 100644 --- a/packages/shared/pkg/telemetry/meters.go +++ b/packages/shared/pkg/telemetry/meters.go @@ -564,7 +564,7 @@ type FloatTimerFactory struct { func NewFloatTimerFactory( meter metric.Meter, - metricName, durationDescription, bytesDescription, _ string, + metricName, durationDescription, bytesDescription string, ) (FloatTimerFactory, error) { duration, err := meter.Float64Histogram(metricName, metric.WithDescription(durationDescription), From 300e86222286722610ee73f6221da057e0c1eb37 Mon Sep 17 00:00:00 2001 From: Lev Brouk Date: Wed, 17 Jun 2026 10:23:57 -0700 Subject: [PATCH 11/11] Removed a merge-generated unused interface --- packages/shared/pkg/storage/storage.go | 6 ------ 1 file changed, 6 deletions(-) diff --git a/packages/shared/pkg/storage/storage.go b/packages/shared/pkg/storage/storage.go index bdbefe65fa..5c7d4d8a27 100644 --- a/packages/shared/pkg/storage/storage.go +++ b/packages/shared/pkg/storage/storage.go @@ -152,12 +152,6 @@ 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) -} - // ReadStats is what a RangeReader did over its lifetime; returned from Close. type ReadStats struct { CompressedBytes int64