From 7d743bc911be3ec4cf11302132af327d4a03c5be Mon Sep 17 00:00:00 2001 From: francoposa Date: Tue, 25 Aug 2026 19:56:03 -0700 Subject: [PATCH 1/9] scratch: WIP index-header bucket reader impl async read ahead --- .../encoding/bucket_async_reader.go | 268 ++++++++++++++++++ .../encoding/bucket_async_reader_test.go | 46 +++ .../indexheader/encoding/bucket_reader.go | 22 +- 3 files changed, 326 insertions(+), 10 deletions(-) create mode 100644 pkg/storage/indexheader/encoding/bucket_async_reader.go create mode 100644 pkg/storage/indexheader/encoding/bucket_async_reader_test.go diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader.go b/pkg/storage/indexheader/encoding/bucket_async_reader.go new file mode 100644 index 00000000000..e61081cffc3 --- /dev/null +++ b/pkg/storage/indexheader/encoding/bucket_async_reader.go @@ -0,0 +1,268 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +package encoding + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "sync" + + "github.com/thanos-io/objstore" + "golang.org/x/sync/errgroup" +) + +const ( + ReadAheadFactor = 4 +) + +type BufPromise struct { + //r *BucketReader + eg *errgroup.Group + //bkt objstore.BucketReader + + //ctx context.Context + //name string + //base int + //length int + //off int + + //buf []byte + bufioReader *bufio.Reader + //readN int +} + +func NewBufPromise( + ctx context.Context, + bkt objstore.BucketReader, + name string, + base int, + length int, + //buf []byte, + bufioReader *bufio.Reader, +) *BufPromise { + bucketReader := NewBucketReader(ctx, bkt, name, base, length) + bufioReader.Reset(bucketReader) + + bp := &BufPromise{ + eg: &errgroup.Group{}, + bufioReader: bufioReader, + } + + bp.eg.Go(func() error { + // Send the fill to the background. + // Peek reads length bytes into the buffer, and does not consume them. + _, err := bp.bufioReader.Peek(length) + return err + }) + return bp +} + +func (bp *BufPromise) Read(p []byte) (n int, err error) { + if err := bp.eg.Wait(); err != nil { + // Return any error in the underlying bucket read from the initial fill-via-Peek. + // A Read after a bad Peek overwrites the bucket read error with a generic short read error. + return 0, err + } + return bp.bufioReader.Read(p) +} + +func (bp *BufPromise) Buffered() (n int, err error) { + if err := bp.eg.Wait(); err != nil { + // Return any error in the underlying bucket read from the initial fill-via-Peek. + // A Read after a bad Peek overwrites the bucket read error with a generic short read error. + return 0, err + } + return bp.bufioReader.Buffered(), nil +} + +// release returns the buffer of the promise to the pool. +// release waits for the fill, because a fill writes to the buffer. +// A buffer in the pool belongs to the next reader that gets it. +// The promise is not usable after release. +func (bp *BufPromise) release(bufioPool *sync.Pool) { + // The error is not relevant here, because release discards the contents of the buffer. + _ = bp.eg.Wait() + bufioPool.Put(bp.bufioReader) + bp.bufioReader = nil +} + +type BucketAsyncBufReader struct { + ctx context.Context + bkt objstore.BucketReader + name string + base int + length int + readOffset int + + resetReader func(off int) error + + bufSize int + + bufIdx int + bufPromises []*BufPromise + bufferedOffset int + + peekBuf []byte + + // pool reference to return to on Close + bufioPool *sync.Pool +} + +func NewBucketAsyncBufReader( + ctx context.Context, bkt objstore.BucketReader, name string, base int, length int, +) *BucketAsyncBufReader { + return newBucketAsyncBufReader( + ctx, bkt, name, base, length, + &bucketBufioPool, ReadBufferSize, ReadAheadFactor, + ) +} + +func newBucketAsyncBufReader( + ctx context.Context, + bkt objstore.BucketReader, + name string, + base int, + length int, + bufioPool *sync.Pool, + bufSize int, + maxBufCount int, +) *BucketAsyncBufReader { + bufsForLength := (length + bufSize - 1) / bufSize + numBufs := min(maxBufCount, bufsForLength) + bufPromises := make([]*BufPromise, numBufs) + + iBase := base + bufferedOffset := 0 + for i := range numBufs { + bufioReader := bufioPool.Get().(*bufio.Reader) + bufLen := min(length-bufferedOffset, bufSize) + bufPromises[i] = NewBufPromise(ctx, bkt, name, iBase, bufLen, bufioReader) + iBase += bufLen + bufferedOffset += bufLen + } + + return &BucketAsyncBufReader{ + ctx: ctx, + bkt: bkt, + name: name, + base: base, + length: length, + bufSize: bufSize, + peekBuf: make([]byte, 0, bufSize), + bufPromises: bufPromises, + bufferedOffset: bufferedOffset, + bufioPool: bufioPool, + } +} + +func (bbar *BucketAsyncBufReader) Reset() error { + //TODO implement me + panic("implement me") +} + +func (bbar *BucketAsyncBufReader) ResetAt(off int) error { + //TODO implement me + panic("implement me") +} + +func (bbar *BucketAsyncBufReader) Skip(l int) error { + //TODO implement me + panic("implement me") +} + +func (bbar *BucketAsyncBufReader) Peek(n int) ([]byte, error) { + //TODO implement me + panic("implement me") +} + +func (bbar *BucketAsyncBufReader) Read(n int) ([]byte, error) { + b := make([]byte, n) + + err := bbar.ReadInto(b) + if err != nil { + return nil, err + } + + return b, nil +} + +func (bbar *BucketAsyncBufReader) ReadInto(b []byte) error { + resultBufWritten := 0 + for resultBufWritten < len(b) { + headPromise := bbar.bufPromises[bbar.bufIdx] + headPromiseBuffered, err := headPromise.Buffered() + if err != nil { + return err + } + + toRead := min(len(b)-resultBufWritten, headPromiseBuffered) + n, err := io.ReadFull(headPromise, b[resultBufWritten:resultBufWritten+toRead]) + bbar.readOffset += n + resultBufWritten += n + + headPromiseBuffered, err = headPromise.Buffered() + if err != nil { + return err + } + + if headPromiseBuffered <= 0 { + // Rotate & replace + headPromise.release(bbar.bufioPool) + bufioReader := bbar.bufioPool.Get().(*bufio.Reader) + bufLen := min(bbar.length-bbar.bufferedOffset, bbar.bufSize) + bbar.bufPromises[bbar.bufIdx] = NewBufPromise( + bbar.ctx, bbar.bkt, bbar.name, bbar.bufferedOffset, bufLen, bufioReader, + ) + bbar.bufIdx = (bbar.bufIdx + 1) % len(bbar.bufPromises) + bbar.bufferedOffset += bufLen + } + + // Now we can surface any error + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(b), err) + } else if err != nil { + return err + } + } + return nil +} + +func (bbar *BucketAsyncBufReader) Size() int { + return bbar.bufSize +} + +func (bbar *BucketAsyncBufReader) Len() int { + return bbar.length - bbar.readOffset +} + +func (bbar *BucketAsyncBufReader) Offset() int { + return bbar.readOffset +} + +func (bbar *BucketAsyncBufReader) Buffered() int { + //TODO implement me + panic("implement me") +} + +// Close releases each promise that the reader still holds, +// and discards the data in the buffers of those promises. +// Close is safe to call more than one time. +func (bbar *BucketAsyncBufReader) Close() error { + for i, bufPromise := range bbar.bufPromises { + if bufPromise == nil { + // A rotate or an earlier Close released the promise in this slot. + continue + } + // Note that we don't do anything to clean up the buffer before returning it to the pool here: + // we reset the buffer when we retrieve it from the pool instead. + bufPromise.release(bbar.bufioPool) + bbar.bufPromises[i] = nil + } + + // The BucketReader of a promise does not need a Close call. + // It closes the reader from bkt.GetRange in each Read call. + return nil +} diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader_test.go b/pkg/storage/indexheader/encoding/bucket_async_reader_test.go new file mode 100644 index 00000000000..3d210bd9261 --- /dev/null +++ b/pkg/storage/indexheader/encoding/bucket_async_reader_test.go @@ -0,0 +1,46 @@ +package encoding + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func newTestAsyncBufReader(t *testing.T, base, length int) (*BucketAsyncBufReader, *trackingBucket) { + t.Helper() + ctx := context.Background() + + objectData := make([]byte, 0, length) + objectData = append(objectData, testBucketContents...) + bkt := newTrackingBucket(t, objectData) + + return newBucketAsyncBufReader( + ctx, bkt, testBucketObjectName, base, length, + &testBucketBufPool, testBufPoolSize, 4, + ), bkt +} + +func TestBucketAsyncBufReader_Read_Sequential(t *testing.T) { + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + + b, err := r.Read(5) + require.NoError(t, err) + require.Equal(t, testBucketContents[:5], b) + require.Equal(t, 5, r.Offset()) + + b, err = r.Read(5) + require.NoError(t, err) + require.Equal(t, testBucketContents[5:10], b) + require.Equal(t, 10, r.Offset()) +} + +func TestBucketAsyncBufReader_Read_ExactLength(t *testing.T) { + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + + b, err := r.Read(len(testBucketContents)) + require.NoError(t, err) + require.Equal(t, testBucketContents, b) + require.Equal(t, len(testBucketContents), r.Offset()) + require.Equal(t, 0, r.Len()) +} diff --git a/pkg/storage/indexheader/encoding/bucket_reader.go b/pkg/storage/indexheader/encoding/bucket_reader.go index 30e973131db..f25a60e4f9d 100644 --- a/pkg/storage/indexheader/encoding/bucket_reader.go +++ b/pkg/storage/indexheader/encoding/bucket_reader.go @@ -13,6 +13,16 @@ import ( "github.com/thanos-io/objstore" ) +const ( + ReadBufferSize = 1 << 20 // 1 MiB +) + +var bucketBufioPool = sync.Pool{ + New: func() any { + return bufio.NewReaderSize(nil, ReadBufferSize) + }, +} + type BucketReader struct { ctx context.Context bkt objstore.BucketReader @@ -70,14 +80,6 @@ func (r *BucketReader) Seek(offset int64, whence int) (int64, error) { return offset, nil } -var bucketBufPool = sync.Pool{ - New: func() any { - // 1MiB buffer chosen as starting point; - // we could make this configurable and benchmark. - return bufio.NewReaderSize(nil, 1<<20) - }, -} - type BucketBufReader struct { ctx context.Context bkt objstore.BucketReader @@ -88,7 +90,7 @@ type BucketBufReader struct { r *BucketReader resetReader func(off int) error buf *bufio.Reader - // Hold a reference to the pool for returning on Close - allows tests to use different pool. + // bufPool reference to return to on Close bufPool *sync.Pool } @@ -107,7 +109,7 @@ func resetReaderFunc(bufReader *BucketBufReader) func(off int) error { func NewBucketBufReader( ctx context.Context, bkt objstore.BucketReader, name string, base int, length int, ) *BucketBufReader { - return newBucketBufReader(ctx, &bucketBufPool, bkt, name, base, length) + return newBucketBufReader(ctx, &bucketBufioPool, bkt, name, base, length) } func newBucketBufReader( From fe25cae8ac0dd2a80d9d85ba2b01681e632db727 Mon Sep 17 00:00:00 2001 From: francoposa Date: Thu, 27 Aug 2026 17:01:50 -0700 Subject: [PATCH 2/9] WIP implement Peek and Skip for index-header bucket reader async read ahead --- .../encoding/bucket_async_reader.go | 185 +++++++++++++--- .../encoding/bucket_async_reader_test.go | 204 ++++++++++++++++++ .../indexheader/encoding/bucket_reader.go | 18 +- .../encoding/bucket_reader_test.go | 22 +- .../indexheader/encoding/file_reader.go | 6 +- pkg/storage/indexheader/encoding/reader.go | 4 +- 6 files changed, 396 insertions(+), 43 deletions(-) diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader.go b/pkg/storage/indexheader/encoding/bucket_async_reader.go index e61081cffc3..fd1c5541479 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader.go @@ -8,6 +8,7 @@ import ( "errors" "fmt" "io" + "slices" "sync" "github.com/thanos-io/objstore" @@ -60,13 +61,13 @@ func NewBufPromise( return bp } -func (bp *BufPromise) Read(p []byte) (n int, err error) { +func (bp *BufPromise) Read(dst []byte) (n int, err error) { if err := bp.eg.Wait(); err != nil { // Return any error in the underlying bucket read from the initial fill-via-Peek. // A Read after a bad Peek overwrites the bucket read error with a generic short read error. return 0, err } - return bp.bufioReader.Read(p) + return bp.bufioReader.Read(dst) } func (bp *BufPromise) Buffered() (n int, err error) { @@ -78,6 +79,26 @@ func (bp *BufPromise) Buffered() (n int, err error) { return bp.bufioReader.Buffered(), nil } +// Peek returns at most n bytes from the promise, without consuming them. +// The byte slice points into the buffer of the promise. +// It becomes invalid when the promise is released. +func (bp *BufPromise) Peek(n int) ([]byte, error) { + if err := bp.eg.Wait(); err != nil { + // Return any error in the underlying bucket read from the initial fill-via-Peek. + return nil, err + } + return bp.bufioReader.Peek(n) +} + +// Discard consumes and drops at most n bytes from the promise. +func (bp *BufPromise) Discard(n int) (discarded int, err error) { + if err := bp.eg.Wait(); err != nil { + // Return any error in the underlying bucket read from the initial fill-via-Peek. + return 0, err + } + return bp.bufioReader.Discard(n) +} + // release returns the buffer of the promise to the pool. // release waits for the fill, because a fill writes to the buffer. // A buffer in the pool belongs to the next reader that gets it. @@ -168,14 +189,122 @@ func (bbar *BucketAsyncBufReader) ResetAt(off int) error { panic("implement me") } +// rotateHead releases the drained head promise, starts a promise for the next chunk +// of the data segment in the same slot, and advances the head index. +// Peek copies its result into peekBuf, so no slice that the reader returned +// points into the buffer of a promise. The release is safe at once. +func (bbar *BucketAsyncBufReader) rotateHead() { + // Note that we don't do anything to clean up the buffer before returning it to the pool here: + // we reset the buffer when we retrieve it from the pool instead. + bbar.bufPromises[bbar.bufIdx].release(bbar.bufioPool) + bufioReader := bbar.bufioPool.Get().(*bufio.Reader) + + // Create a new buffer promise in the same spot in the buffer promise queue. + // The promise must not reach past the end of the reader. + // bufferedOffset is relative to the data segment. + // Add base to get the offset in the object. + bufLen := min(bbar.length-bbar.bufferedOffset, bbar.bufSize) + bbar.bufPromises[bbar.bufIdx] = NewBufPromise( + bbar.ctx, bbar.bkt, bbar.name, bbar.base+bbar.bufferedOffset, bufLen, bufioReader, + ) + bbar.bufferedOffset += bufLen + + // Advance current buffer queue index - modulo wraps around to the front of the slice if at end. + bbar.bufIdx = (bbar.bufIdx + 1) % len(bbar.bufPromises) +} + +// Skip advances the cursor by l bytes in the data segment and discards those bytes. +// Skip returns ErrInvalidSize if l is greater than the number of bytes that remain. func (bbar *BucketAsyncBufReader) Skip(l int) error { - //TODO implement me - panic("implement me") + if l > bbar.Len() { + return ErrInvalidSize + } + + bytesSkipped := 0 + // First try to complete the skip from previously-peeked bytes. + // If peekBuf is non-empty, those bytes were not skipped or read yet. + n := min(len(bbar.peekBuf), l) + bbar.readOffset += n + bytesSkipped += n + // Truncate the peekBuf even if we did not skip all the previously-peeked bytes. + // Peek interface contract says "byte slice returned becomes invalid at the next read" (which includes Skip). + bbar.peekBuf = bbar.peekBuf[:0] + + // Move on to skip the data from promises if we have not complete the skip yet. + for bytesSkipped < l { + headPromise := bbar.bufPromises[bbar.bufIdx] + headPromiseBuffered, err := headPromise.Buffered() + if err != nil { + return err + } + + toSkip := min(l-bytesSkipped, headPromiseBuffered) + n, err := headPromise.Discard(toSkip) + if err != nil { + return err + } + bbar.readOffset += n + bytesSkipped += n + + headPromiseBuffered, err = headPromise.Buffered() + if err != nil { + return err + } + + if headPromiseBuffered <= 0 { + bbar.rotateHead() + } + } + + return nil } +// Peek returns at most n bytes from the data segment, without consuming them. +// Peek always copies the bytes into peekBuf for now. +// This keeps the logic simple for rotating out exhausted buffer promises +// in the case where a peek crosses the promise boundaries. +// Since peekBuf is pre-allocated, this still avoids the extra slice allocation +// which occurs when callers Read instead of Peek. func (bbar *BucketAsyncBufReader) Peek(n int) ([]byte, error) { - //TODO implement me - panic("implement me") + // Ensure peekBuf has capacity of n by calling Grow against the truncated slice. + // This should never trigger a new alloc as peekBuf is pre-allocated to larger than we need - + // at most it needs to hold the length of one Prometheus label or value. + // Length must be truncated before return in the case of a short Peek. + bbar.peekBuf = slices.Grow(bbar.peekBuf[:0], n)[:n] + + peekableBytes := bbar.Size() + peekBytesWritten := 0 + for peekBytesWritten < n && peekableBytes > 0 { + headPromise := bbar.bufPromises[bbar.bufIdx] + headPromiseBuffered, err := headPromise.Buffered() + if err != nil { + return nil, err + } + + toRead := min(n-peekBytesWritten, headPromiseBuffered) + readN, err := io.ReadFull(headPromise, bbar.peekBuf[peekBytesWritten:peekBytesWritten+toRead]) + peekBytesWritten += readN + peekableBytes -= readN + if err != nil { + return nil, err + } + + headPromiseBuffered, err = headPromise.Buffered() + if err != nil { + return nil, err + } + + if headPromiseBuffered <= 0 { + bbar.rotateHead() + } + } + // A short Peek is valid; truncate to what was actually read. + bbar.peekBuf = bbar.peekBuf[:peekBytesWritten] + + if peekBytesWritten == 0 { + return nil, nil + } + return bbar.peekBuf[:peekBytesWritten], nil } func (bbar *BucketAsyncBufReader) Read(n int) ([]byte, error) { @@ -189,19 +318,33 @@ func (bbar *BucketAsyncBufReader) Read(n int) ([]byte, error) { return b, nil } -func (bbar *BucketAsyncBufReader) ReadInto(b []byte) error { - resultBufWritten := 0 - for resultBufWritten < len(b) { +func (bbar *BucketAsyncBufReader) ReadInto(dst []byte) error { + // First try to satisfy the read from previously-peeked bytes. + dstBytesWritten := 0 + n := copy(dst, bbar.peekBuf) + bbar.readOffset += n + dstBytesWritten += n + + // Truncate the peekBuf even if we did not read all the previously-peeked bytes. + // Peek interface contract says "byte slice returned becomes invalid at the next read". + // We do not need to try to serve two subsequent reads from the peekBuf even if they fit. + bbar.peekBuf = bbar.peekBuf[:0] + + // Move on to read from the promises if we have not satisfied the read yet. + for dstBytesWritten < len(dst) { headPromise := bbar.bufPromises[bbar.bufIdx] headPromiseBuffered, err := headPromise.Buffered() if err != nil { return err } - toRead := min(len(b)-resultBufWritten, headPromiseBuffered) - n, err := io.ReadFull(headPromise, b[resultBufWritten:resultBufWritten+toRead]) + toRead := min(len(dst)-dstBytesWritten, headPromiseBuffered) + n, err := io.ReadFull(headPromise, dst[dstBytesWritten:dstBytesWritten+toRead]) bbar.readOffset += n - resultBufWritten += n + dstBytesWritten += n + if err != nil { + return err + } headPromiseBuffered, err = headPromise.Buffered() if err != nil { @@ -209,20 +352,12 @@ func (bbar *BucketAsyncBufReader) ReadInto(b []byte) error { } if headPromiseBuffered <= 0 { - // Rotate & replace - headPromise.release(bbar.bufioPool) - bufioReader := bbar.bufioPool.Get().(*bufio.Reader) - bufLen := min(bbar.length-bbar.bufferedOffset, bbar.bufSize) - bbar.bufPromises[bbar.bufIdx] = NewBufPromise( - bbar.ctx, bbar.bkt, bbar.name, bbar.bufferedOffset, bufLen, bufioReader, - ) - bbar.bufIdx = (bbar.bufIdx + 1) % len(bbar.bufPromises) - bbar.bufferedOffset += bufLen + bbar.rotateHead() } // Now we can surface any error if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { - return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(b), err) + return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(dst), err) } else if err != nil { return err } @@ -230,8 +365,10 @@ func (bbar *BucketAsyncBufReader) ReadInto(b []byte) error { return nil } +// Size returns the largest number of bytes that a single Peek can return. +// Peek assembles its result in peekBuf, so the capacity of peekBuf is the limit. func (bbar *BucketAsyncBufReader) Size() int { - return bbar.bufSize + return cap(bbar.peekBuf) } func (bbar *BucketAsyncBufReader) Len() int { @@ -253,7 +390,7 @@ func (bbar *BucketAsyncBufReader) Buffered() int { func (bbar *BucketAsyncBufReader) Close() error { for i, bufPromise := range bbar.bufPromises { if bufPromise == nil { - // A rotate or an earlier Close released the promise in this slot. + // An earlier Close released the promise in this slot. continue } // Note that we don't do anything to clean up the buffer before returning it to the pool here: diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader_test.go b/pkg/storage/indexheader/encoding/bucket_async_reader_test.go index 3d210bd9261..bc5244d47e0 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader_test.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader_test.go @@ -7,6 +7,11 @@ import ( "github.com/stretchr/testify/require" ) +// testBucketContentsLong is a 64-byte payload for tests that need more data than +// the read-ahead window holds. Every byte is distinct, so a read at a wrong offset +// gives a different result. +var testBucketContentsLong = []byte("abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ+/") + func newTestAsyncBufReader(t *testing.T, base, length int) (*BucketAsyncBufReader, *trackingBucket) { t.Helper() ctx := context.Background() @@ -21,6 +26,23 @@ func newTestAsyncBufReader(t *testing.T, base, length int) (*BucketAsyncBufReade ), bkt } +// newTestAsyncBufReaderWithData builds a reader over the given object data, +// with an explicit read-ahead buffer count. +// A test that needs a rotation of the read-ahead window must use this helper, +// because a rotation needs a length greater than maxBufCount*testBufPoolSize bytes. +func newTestAsyncBufReaderWithData( + t *testing.T, objectData []byte, base, length, maxBufCount int, +) (*BucketAsyncBufReader, *trackingBucket) { + t.Helper() + + bkt := newTrackingBucket(t, objectData) + + return newBucketAsyncBufReader( + t.Context(), bkt, testBucketObjectName, base, length, + &testBucketBufPool, testBufPoolSize, maxBufCount, + ), bkt +} + func TestBucketAsyncBufReader_Read_Sequential(t *testing.T) { r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) @@ -44,3 +66,185 @@ func TestBucketAsyncBufReader_Read_ExactLength(t *testing.T) { require.Equal(t, len(testBucketContents), r.Offset()) require.Equal(t, 0, r.Len()) } + +func TestBucketAsyncBufReader_Peek_Basic(t *testing.T) { + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + + b, err := r.Peek(5) + require.NoError(t, err) + require.Equal(t, testBucketContents[:5], b) + require.Equal(t, 0, r.Offset(), "Peek does not consume") + + // Read returns the same bytes. + got, err := r.Read(5) + require.NoError(t, err) + require.Equal(t, testBucketContents[:5], got) +} + +func TestBucketAsyncBufReader_Peek_AcrossPromiseBoundary(t *testing.T) { + // Each buffer promise holds testBufPoolSize bytes. + // A peek of 8 bytes at offset 12 takes 4 bytes from the first promise + // and 4 bytes from the second promise. + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + require.NoError(t, r.Skip(testBufPoolSize-4)) + + b, err := r.Peek(8) + require.NoError(t, err) + require.Equal(t, testBucketContents[testBufPoolSize-4:testBufPoolSize+4], b) + require.Equal(t, testBufPoolSize-4, r.Offset(), "Peek does not consume") +} + +func TestBucketAsyncBufReader_Peek_PastSegmentEnd(t *testing.T) { + // Peek must return the bytes read and suppress the EOF error + // when peeking past the configured length or peeking beyond the true end of the object. + const sectionLen = 5 + r, _ := newTestAsyncBufReader(t, 0, sectionLen) + + b, err := r.Peek(sectionLen + 5) + require.NoError(t, err) + require.Equal(t, testBucketContents[:sectionLen], b) + require.Equal(t, 0, r.Offset()) +} + +func TestBucketAsyncBufReader_Peek_AtEnd(t *testing.T) { + const sectionLen = 5 + r, _ := newTestAsyncBufReader(t, 0, sectionLen) + require.NoError(t, r.Skip(sectionLen)) + + b, err := r.Peek(1) + require.NoError(t, err) + require.Nil(t, b) +} + +//func TestBucketAsyncBufReader_Peek_BeyondPeekBuffer(t *testing.T) { +// r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) +// +// b, err := r.Peek(r.Size() + 1) +// require.ErrorIs(t, err, ErrInvalidSize) +// require.Nil(t, b) +//} + +// TestBucketAsyncBufReader_Peek_ThenSkip covers the access pattern of +// Decbuf.UnsafeUvarintBytes, which peeks bytes and then skips exactly those bytes. +// The skip drains the head promise and returns the buffer of that promise to the pool. +// The reader then takes that same buffer back and fills it with the third chunk. +// Peek copies into peekBuf, so the peeked bytes stay valid through that refill. +func TestBucketAsyncBufReader_Peek_ThenSkip(t *testing.T) { + // The read-ahead window holds maxBufCount*testBufPoolSize = 32 bytes of the 48-byte segment, + // so the reader must refill a buffer to reach the third chunk. + const ( + length = 48 + maxBufCount = 2 + ) + r, _ := newTestAsyncBufReaderWithData(t, testBucketContentsLong, 0, length, maxBufCount) + + // Peek and skip the whole head promise, which drains and releases it. + b, err := r.Peek(testBufPoolSize) + require.NoError(t, err) + require.NoError(t, r.Skip(len(b))) + require.Equal(t, testBufPoolSize, r.Offset()) + + // Read the second chunk, then one byte of the third chunk. + // The read of the third chunk waits for the refill of the released buffer. + got, err := r.Read(testBufPoolSize) + require.NoError(t, err) + require.Equal(t, testBucketContentsLong[testBufPoolSize:2*testBufPoolSize], got) + + got, err = r.Read(1) + require.NoError(t, err) + require.Equal(t, testBucketContentsLong[2*testBufPoolSize:2*testBufPoolSize+1], got) + + require.Equal(t, testBucketContentsLong[:testBufPoolSize], b, "peeked bytes survive the refill") +} + +// TestBucketAsyncBufReader_Peek_ThenSkip_AcrossPromiseBoundary makes sure that a skip +// after a straddling peek discards the bytes that the peek returned. +// Peek does not consume, so the cursor still points at the first peeked byte. +func TestBucketAsyncBufReader_Peek_ThenSkip_AcrossPromiseBoundary(t *testing.T) { + // A peek of 8 bytes at offset 12 takes 4 bytes from the first promise + // and 4 bytes from the second promise, so it goes through peekBuf. + const peekAt = testBufPoolSize - 4 + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + require.NoError(t, r.Skip(peekAt)) + + b, err := r.Peek(8) + require.NoError(t, err) + require.Equal(t, testBucketContents[peekAt:peekAt+8], b) + + require.NoError(t, r.Skip(len(b))) + require.Equal(t, peekAt+8, r.Offset(), "the skip consumes the peeked bytes, not the bytes after them") + require.Equal(t, testBucketContents[peekAt:peekAt+8], b, "the skip does not disturb peekBuf") + + // The next read starts one byte past the peeked bytes. + got, err := r.Read(4) + require.NoError(t, err) + require.Equal(t, testBucketContents[peekAt+8:peekAt+12], got) +} + +func TestBucketAsyncBufReader_Skip_Basic(t *testing.T) { + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + + require.NoError(t, r.Skip(10)) + require.Equal(t, 10, r.Offset()) + require.Equal(t, len(testBucketContents)-10, r.Len()) + + b, err := r.Read(3) + require.NoError(t, err) + require.Equal(t, testBucketContents[10:13], b) +} + +func TestBucketAsyncBufReader_Skip_AcrossPromises(t *testing.T) { + // A skip of 20 bytes drains the first promise of testBufPoolSize bytes, + // then takes the remainder from the second promise. + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + + require.NoError(t, r.Skip(testBufPoolSize+4)) + require.Equal(t, testBufPoolSize+4, r.Offset()) + + b, err := r.Read(3) + require.NoError(t, err) + require.Equal(t, testBucketContents[testBufPoolSize+4:testBufPoolSize+7], b) +} + +func TestBucketAsyncBufReader_Skip_ToEnd(t *testing.T) { + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + + require.NoError(t, r.Skip(len(testBucketContents))) + require.Equal(t, len(testBucketContents), r.Offset()) + require.Equal(t, 0, r.Len()) +} + +func TestBucketAsyncBufReader_Skip_BeyondEnd(t *testing.T) { + const sectionLen = 10 + r, _ := newTestAsyncBufReader(t, 0, sectionLen) + + require.ErrorIs(t, r.Skip(sectionLen+1), ErrInvalidSize) +} + +//func TestBucketAsyncBufReader_Skip_Negative(t *testing.T) { +// // Decbuf.SkipUvarintBytes converts a uint64 from the object to an int, +// // so a corrupt length can arrive as a negative number. +// r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) +// +// require.ErrorIs(t, r.Skip(-1), ErrInvalidSize) +// require.Equal(t, 0, r.Offset()) +//} + +func TestBucketAsyncBufReader_Read_ExactLength_NonZeroBaseWithRotation(t *testing.T) { + // The read-ahead window holds maxBufCount*testBufPoolSize = 32 bytes. + // A length of 48 forces the reader to release one buffer promise and refill it. + // The refilled promise must add base to the buffered offset. + // Without base, the promise reads the object 4 bytes too early. + const ( + base = 4 + length = 48 + maxBufCount = 2 + ) + r, _ := newTestAsyncBufReaderWithData(t, testBucketContentsLong, base, length, maxBufCount) + + b, err := r.Read(length) + require.NoError(t, err) + require.Equal(t, testBucketContentsLong[base:base+length], b) + require.Equal(t, length, r.Offset()) + require.Equal(t, 0, r.Len()) +} diff --git a/pkg/storage/indexheader/encoding/bucket_reader.go b/pkg/storage/indexheader/encoding/bucket_reader.go index f25a60e4f9d..08dddf21183 100644 --- a/pkg/storage/indexheader/encoding/bucket_reader.go +++ b/pkg/storage/indexheader/encoding/bucket_reader.go @@ -44,14 +44,14 @@ func NewBucketReader( } } -func (r *BucketReader) Read(p []byte) (n int, err error) { - if len(p) == 0 { +func (r *BucketReader) Read(dst []byte) (n int, err error) { + if len(dst) == 0 { return 0, nil } if r.off >= r.length { return 0, io.EOF } - toRead := len(p) + toRead := len(dst) remaining := r.length - r.off if toRead > remaining { toRead = remaining @@ -61,7 +61,7 @@ func (r *BucketReader) Read(p []byte) (n int, err error) { return 0, err } defer rc.Close() - n, err = io.ReadFull(rc, p[:toRead]) + n, err = io.ReadFull(rc, dst[:toRead]) r.off += n if errors.Is(err, io.ErrUnexpectedEOF) { err = io.EOF @@ -164,9 +164,7 @@ func (bbr *BucketBufReader) Skip(l int) error { } n, err := bbr.buf.Discard(l) - if n > 0 { - bbr.off += n - } + bbr.off += n return err } @@ -197,14 +195,14 @@ func (bbr *BucketBufReader) Read(n int) ([]byte, error) { return b, nil } -func (bbr *BucketBufReader) ReadInto(b []byte) error { - n, err := io.ReadFull(bbr.buf, b) +func (bbr *BucketBufReader) ReadInto(dst []byte) error { + n, err := io.ReadFull(bbr.buf, dst) if n > 0 { bbr.off += n } if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { - return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(b), err) + return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(dst), err) } else if err != nil { return err } diff --git a/pkg/storage/indexheader/encoding/bucket_reader_test.go b/pkg/storage/indexheader/encoding/bucket_reader_test.go index 5310428eb4d..e6a688dcd66 100644 --- a/pkg/storage/indexheader/encoding/bucket_reader_test.go +++ b/pkg/storage/indexheader/encoding/bucket_reader_test.go @@ -8,6 +8,7 @@ import ( "context" "errors" "io" + "slices" "sync" "testing" @@ -278,12 +279,12 @@ func TestBucketBufReader_GetRangeCalls_Buffering(t *testing.T) { _, err := r.Read(1) require.NoError(t, err) } - require.Len(t, bkt.calls, 1) + require.Len(t, bkt.rangeCalls(), 1) // Buffer depleted; next read operation triggers another bufio fill and GetRange call. _, err := r.Read(1) require.NoError(t, err) - require.Len(t, bkt.calls, 2) + require.Len(t, bkt.rangeCalls(), 2) } func TestBucketBufReader_GetRangeCalls_ResetRefetches(t *testing.T) { @@ -291,13 +292,13 @@ func TestBucketBufReader_GetRangeCalls_ResetRefetches(t *testing.T) { _, err := r.Read(1) require.NoError(t, err) - require.Len(t, bkt.calls, 1) + require.Len(t, bkt.rangeCalls(), 1) // After Reset the buffer is discarded; the next read must refetch from the bucket. require.NoError(t, r.Reset()) _, err = r.Read(1) require.NoError(t, err) - require.Len(t, bkt.calls, 2) + require.Len(t, bkt.rangeCalls(), 2) } func TestBucketBufReader_Read_GetRangeError(t *testing.T) { @@ -317,8 +318,12 @@ func TestBucketBufReader_ReadInto_GetRangeError(t *testing.T) { } // trackingBucket wraps an InstrumentedBucketReader and records every GetRange call. +// The read-ahead reader fills its buffer promises from several goroutines at the same time, +// so the mutex protects the record of the calls. type trackingBucket struct { objstore.InstrumentedBucketReader + + mtx sync.Mutex calls []rangeCall } @@ -328,10 +333,19 @@ type rangeCall struct { } func (b *trackingBucket) GetRange(ctx context.Context, name string, off, length int64) (io.ReadCloser, error) { + b.mtx.Lock() b.calls = append(b.calls, rangeCall{off, length}) + b.mtx.Unlock() return b.InstrumentedBucketReader.GetRange(ctx, name, off, length) } +// rangeCalls returns a copy of the recorded GetRange calls. +func (b *trackingBucket) rangeCalls() []rangeCall { + b.mtx.Lock() + defer b.mtx.Unlock() + return slices.Clone(b.calls) +} + func newTrackingBucket(t *testing.T, objectData []byte) *trackingBucket { t.Helper() inmem := objstore.NewInMemBucket() diff --git a/pkg/storage/indexheader/encoding/file_reader.go b/pkg/storage/indexheader/encoding/file_reader.go index 5fe6bdb1d83..bfba1d81b79 100644 --- a/pkg/storage/indexheader/encoding/file_reader.go +++ b/pkg/storage/indexheader/encoding/file_reader.go @@ -110,14 +110,14 @@ func (f *FileReader) Read(n int) ([]byte, error) { return b, nil } -func (f *FileReader) ReadInto(b []byte) error { - r, err := io.ReadFull(f.buf, b) +func (f *FileReader) ReadInto(dst []byte) error { + r, err := io.ReadFull(f.buf, dst) if r > 0 { f.off += r } if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { - return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(b), err) + return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(dst), err) } else if err != nil { return err } diff --git a/pkg/storage/indexheader/encoding/reader.go b/pkg/storage/indexheader/encoding/reader.go index 9c3222b1ab4..31c1a7ec700 100644 --- a/pkg/storage/indexheader/encoding/reader.go +++ b/pkg/storage/indexheader/encoding/reader.go @@ -35,11 +35,11 @@ type BufReader interface { // and the remaining bytes MUST be consumed. Read(n int) ([]byte, error) - // ReadInto reads len(b) bytes from the data segment into b, consuming them. + // ReadInto reads len(dst) bytes from the data segment into dst, consuming them. // It is NOT valid to read beyond the end of the data segment; // in this case implementations MUST return a nil byte slice and an ErrInvalidSize error, // and the remaining bytes MUST be consumed. - ReadInto(b []byte) error + ReadInto(dst []byte) error // Size returns the length of the underlying buffer in bytes. Size() int From 5387161eea9b41d080a5b88e17f2db00c8200521 Mon Sep 17 00:00:00 2001 From: francoposa Date: Sat, 29 Aug 2026 17:51:48 -0700 Subject: [PATCH 3/9] Fix up Peek and ReadInto impls --- .../encoding/bucket_async_reader.go | 95 +++++++++++-------- .../encoding/bucket_async_reader_test.go | 83 ++++++++++++---- pkg/storage/indexheader/encoding/reader.go | 6 +- 3 files changed, 121 insertions(+), 63 deletions(-) diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader.go b/pkg/storage/indexheader/encoding/bucket_async_reader.go index fd1c5541479..e1ef97e2898 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader.go @@ -5,10 +5,7 @@ package encoding import ( "bufio" "context" - "errors" - "fmt" "io" - "slices" "sync" "github.com/thanos-io/objstore" @@ -259,22 +256,23 @@ func (bbar *BucketAsyncBufReader) Skip(l int) error { return nil } -// Peek returns at most n bytes from the data segment, without consuming them. -// Peek always copies the bytes into peekBuf for now. -// This keeps the logic simple for rotating out exhausted buffer promises -// in the case where a peek crosses the promise boundaries. -// Since peekBuf is pre-allocated, this still avoids the extra slice allocation -// which occurs when callers Read instead of Peek. func (bbar *BucketAsyncBufReader) Peek(n int) ([]byte, error) { - // Ensure peekBuf has capacity of n by calling Grow against the truncated slice. - // This should never trigger a new alloc as peekBuf is pre-allocated to larger than we need - - // at most it needs to hold the length of one Prometheus label or value. - // Length must be truncated before return in the case of a short Peek. - bbar.peekBuf = slices.Grow(bbar.peekBuf[:0], n)[:n] - - peekableBytes := bbar.Size() - peekBytesWritten := 0 - for peekBytesWritten < n && peekableBytes > 0 { + // Clamp n to the lesser of the capacity of peekBuf or the length of the section. + n = min(n, cap(bbar.peekBuf), bbar.Len()) + + // Start with any previously-peeked bytes - Peek-after-Peek is a valid access pattern. + // Any data remaining in peekBuf is assumed to still be the valid start of a Peek. + // The read operations (Read/ReadInto and Skip) are required to update peekBuf + // to discard any previously-peeked bytes which were consumed by the read. + peekBytesAvailable := len(bbar.peekBuf) + if n > peekBytesAvailable { + bbar.peekBuf = bbar.peekBuf[:n] // Grow length + } + peekBytesWritten := min(n, peekBytesAvailable) + + // Move on to the promises if we have not satisfied the peek yet. + // Promises are consumed to copy into the peekBuf and rotated if exhausted. + for peekBytesWritten < n { headPromise := bbar.bufPromises[bbar.bufIdx] headPromiseBuffered, err := headPromise.Buffered() if err != nil { @@ -284,7 +282,6 @@ func (bbar *BucketAsyncBufReader) Peek(n int) ([]byte, error) { toRead := min(n-peekBytesWritten, headPromiseBuffered) readN, err := io.ReadFull(headPromise, bbar.peekBuf[peekBytesWritten:peekBytesWritten+toRead]) peekBytesWritten += readN - peekableBytes -= readN if err != nil { return nil, err } @@ -298,8 +295,6 @@ func (bbar *BucketAsyncBufReader) Peek(n int) ([]byte, error) { bbar.rotateHead() } } - // A short Peek is valid; truncate to what was actually read. - bbar.peekBuf = bbar.peekBuf[:peekBytesWritten] if peekBytesWritten == 0 { return nil, nil @@ -319,16 +314,32 @@ func (bbar *BucketAsyncBufReader) Read(n int) ([]byte, error) { } func (bbar *BucketAsyncBufReader) ReadInto(dst []byte) error { - // First try to satisfy the read from previously-peeked bytes. - dstBytesWritten := 0 - n := copy(dst, bbar.peekBuf) - bbar.readOffset += n - dstBytesWritten += n - - // Truncate the peekBuf even if we did not read all the previously-peeked bytes. - // Peek interface contract says "byte slice returned becomes invalid at the next read". - // We do not need to try to serve two subsequent reads from the peekBuf even if they fit. - bbar.peekBuf = bbar.peekBuf[:0] + // TODO consistency in error handling + //// A read past the end of the data segment is not valid. + //// The guard in Skip returns ErrInvalidSize for the same condition. + //// The reader contract also requires this read to consume the bytes that remain, + //// so move the cursor to the end before the return. + //// Without this guard the loop below never ends, + //// because a drained head promise rotates to an empty promise forever. + //if len(dst) > bbar.Len() { + // remaining := bbar.Len() + // if err := bbar.Skip(remaining); err != nil { + // return err + // } + // // io.ReadFull reports io.EOF for no bytes and io.ErrUnexpectedEOF for a partial read. + // // BucketBufReader passes that error through, so report the same error here. + // shortErr := io.ErrUnexpectedEOF + // if remaining == 0 { + // shortErr = io.EOF + // } + // return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(dst), shortErr) + //} + + // Start with any previously-peeked bytes + dstBytesWritten := copy(dst, bbar.peekBuf) + bbar.readOffset += dstBytesWritten + // Slide any unconsumed bytes from peekBuf to the beginning of the slice and truncate. + bbar.peekBuf = bbar.peekBuf[dstBytesWritten:] // Move on to read from the promises if we have not satisfied the read yet. for dstBytesWritten < len(dst) { @@ -339,9 +350,9 @@ func (bbar *BucketAsyncBufReader) ReadInto(dst []byte) error { } toRead := min(len(dst)-dstBytesWritten, headPromiseBuffered) - n, err := io.ReadFull(headPromise, dst[dstBytesWritten:dstBytesWritten+toRead]) - bbar.readOffset += n - dstBytesWritten += n + readN, err := io.ReadFull(headPromise, dst[dstBytesWritten:dstBytesWritten+toRead]) + bbar.readOffset += readN + dstBytesWritten += readN if err != nil { return err } @@ -355,12 +366,13 @@ func (bbar *BucketAsyncBufReader) ReadInto(dst []byte) error { bbar.rotateHead() } - // Now we can surface any error - if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { - return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(dst), err) - } else if err != nil { - return err - } + // TODO consistency in error handling + //// Now we can surface any error + //if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + // return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(dst), err) + //} else if err != nil { + // return err + //} } return nil } @@ -380,8 +392,7 @@ func (bbar *BucketAsyncBufReader) Offset() int { } func (bbar *BucketAsyncBufReader) Buffered() int { - //TODO implement me - panic("implement me") + return bbar.bufferedOffset - bbar.readOffset } // Close releases each promise that the reader still holds, diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader_test.go b/pkg/storage/indexheader/encoding/bucket_async_reader_test.go index bc5244d47e0..63194941ec7 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader_test.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader_test.go @@ -43,42 +43,61 @@ func newTestAsyncBufReaderWithData( ), bkt } -func TestBucketAsyncBufReader_Read_Sequential(t *testing.T) { +func TestBucketAsyncBufReader_Peek_Peek(t *testing.T) { r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) - b, err := r.Read(5) + peek1, err := r.Peek(5) require.NoError(t, err) - require.Equal(t, testBucketContents[:5], b) - require.Equal(t, 5, r.Offset()) + require.Equal(t, testBucketContents[:5], peek1) + require.Equal(t, 0, r.Offset(), "Peek does not consume") - b, err = r.Read(5) + // Same-size Peek returns the same bytes. + peek2, err := r.Peek(5) require.NoError(t, err) - require.Equal(t, testBucketContents[5:10], b) - require.Equal(t, 10, r.Offset()) -} + require.Equal(t, testBucketContents[:5], peek2) -func TestBucketAsyncBufReader_Read_ExactLength(t *testing.T) { - r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + // Smaller Peek returns the same starting bytes. + peek3, err := r.Peek(3) + require.NoError(t, err) + require.Equal(t, testBucketContents[:3], peek3) - b, err := r.Read(len(testBucketContents)) + // Larger Peek returns the original bytes plus more + peek4, err := r.Peek(10) require.NoError(t, err) - require.Equal(t, testBucketContents, b) - require.Equal(t, len(testBucketContents), r.Offset()) - require.Equal(t, 0, r.Len()) + require.Equal(t, testBucketContents[:10], peek4) } -func TestBucketAsyncBufReader_Peek_Basic(t *testing.T) { +func TestBucketAsyncBufReader_Peak_Read(t *testing.T) { r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) - b, err := r.Peek(5) + peek1, err := r.Peek(5) require.NoError(t, err) - require.Equal(t, testBucketContents[:5], b) - require.Equal(t, 0, r.Offset(), "Peek does not consume") + require.Equal(t, testBucketContents[:5], peek1) + require.Equal(t, 0, r.Offset()) // Read returns the same bytes. - got, err := r.Read(5) + read1, err := r.Read(5) require.NoError(t, err) - require.Equal(t, testBucketContents[:5], got) + require.Equal(t, testBucketContents[:5], read1) + require.Equal(t, 5, r.Offset()) + + // Another Peek returns the next bytes after Read consumed previous bytes. + peek2, err := r.Peek(5) + require.NoError(t, err) + require.Equal(t, testBucketContents[5:10], peek2) + require.Equal(t, 5, r.Offset()) + + // A Read short of the peeked bytes consumes some of them. + read2, err := r.Read(3) + require.NoError(t, err) + require.Equal(t, testBucketContents[5:8], read2) + require.Equal(t, 8, r.Offset()) + + // A Read beyond the remaining peeked bytes consumes them and more. + read3, err := r.Read(8) + require.NoError(t, err) + require.Equal(t, testBucketContents[8:16], read3) + require.Equal(t, 16, r.Offset()) } func TestBucketAsyncBufReader_Peek_AcrossPromiseBoundary(t *testing.T) { @@ -116,6 +135,30 @@ func TestBucketAsyncBufReader_Peek_AtEnd(t *testing.T) { require.Nil(t, b) } +func TestBucketAsyncBufReader_Read_Sequential(t *testing.T) { + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + + b, err := r.Read(5) + require.NoError(t, err) + require.Equal(t, testBucketContents[:5], b) + require.Equal(t, 5, r.Offset()) + + b, err = r.Read(5) + require.NoError(t, err) + require.Equal(t, testBucketContents[5:10], b) + require.Equal(t, 10, r.Offset()) +} + +func TestBucketAsyncBufReader_Read_ExactLength(t *testing.T) { + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + + b, err := r.Read(len(testBucketContents)) + require.NoError(t, err) + require.Equal(t, testBucketContents, b) + require.Equal(t, len(testBucketContents), r.Offset()) + require.Equal(t, 0, r.Len()) +} + //func TestBucketAsyncBufReader_Peek_BeyondPeekBuffer(t *testing.T) { // r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) // diff --git a/pkg/storage/indexheader/encoding/reader.go b/pkg/storage/indexheader/encoding/reader.go index 31c1a7ec700..165c4a76a70 100644 --- a/pkg/storage/indexheader/encoding/reader.go +++ b/pkg/storage/indexheader/encoding/reader.go @@ -24,9 +24,13 @@ type BufReader interface { Skip(l int) error // Peek returns at most the given number of bytes from the data segment, without consuming them. - // The byte slice returned becomes invalid at the next read. // It is valid to Peek beyond the end of the data segment; // in this case implementations MUST return the available bytes up to the end and a nil error. + // The byte slice returned becomes invalid at the next read. + // + // Peek is limited to and only must support reads up to the underlying buffer length. + // Caller checks Size first to see if the next read operation length fits in the underlying buffer, + // then a Peek-Skip pattern is used to avoid the slice allocation which must occur in Read. Peek(n int) ([]byte, error) // Read returns the given number of bytes from the data segment, consuming them. From 98bd2abe2860b3799f2559ba5a6971ac9f28be42 Mon Sep 17 00:00:00 2001 From: francoposa Date: Sat, 29 Aug 2026 18:02:45 -0700 Subject: [PATCH 4/9] Update Skip to match Peek and ReadInto and test --- .../encoding/bucket_async_reader.go | 33 +++++++++---------- .../encoding/bucket_async_reader_test.go | 31 +++++++++++++++-- 2 files changed, 44 insertions(+), 20 deletions(-) diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader.go b/pkg/storage/indexheader/encoding/bucket_async_reader.go index e1ef97e2898..9d2639d040b 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader.go @@ -210,24 +210,20 @@ func (bbar *BucketAsyncBufReader) rotateHead() { bbar.bufIdx = (bbar.bufIdx + 1) % len(bbar.bufPromises) } -// Skip advances the cursor by l bytes in the data segment and discards those bytes. -// Skip returns ErrInvalidSize if l is greater than the number of bytes that remain. func (bbar *BucketAsyncBufReader) Skip(l int) error { if l > bbar.Len() { return ErrInvalidSize } - bytesSkipped := 0 - // First try to complete the skip from previously-peeked bytes. - // If peekBuf is non-empty, those bytes were not skipped or read yet. - n := min(len(bbar.peekBuf), l) - bbar.readOffset += n - bytesSkipped += n - // Truncate the peekBuf even if we did not skip all the previously-peeked bytes. - // Peek interface contract says "byte slice returned becomes invalid at the next read" (which includes Skip). - bbar.peekBuf = bbar.peekBuf[:0] - - // Move on to skip the data from promises if we have not complete the skip yet. + // Start with any previously-peeked bytes. + bytesSkipped := min(len(bbar.peekBuf), l) + bbar.readOffset += bytesSkipped + + // Slide any unconsumed bytes from peekBuf to the beginning of the slice and truncate. + bbar.peekBuf = bbar.peekBuf[bytesSkipped:] + + // Move on to the promises if we have not satisfied the skip yet. + // Promises are consumed to discard the data and rotated if exhausted. for bytesSkipped < l { headPromise := bbar.bufPromises[bbar.bufIdx] headPromiseBuffered, err := headPromise.Buffered() @@ -236,12 +232,12 @@ func (bbar *BucketAsyncBufReader) Skip(l int) error { } toSkip := min(l-bytesSkipped, headPromiseBuffered) - n, err := headPromise.Discard(toSkip) + skipN, err := headPromise.Discard(toSkip) if err != nil { return err } - bbar.readOffset += n - bytesSkipped += n + bbar.readOffset += skipN + bytesSkipped += skipN headPromiseBuffered, err = headPromise.Buffered() if err != nil { @@ -335,13 +331,14 @@ func (bbar *BucketAsyncBufReader) ReadInto(dst []byte) error { // return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(dst), shortErr) //} - // Start with any previously-peeked bytes + // Start with any previously-peeked bytes. dstBytesWritten := copy(dst, bbar.peekBuf) bbar.readOffset += dstBytesWritten // Slide any unconsumed bytes from peekBuf to the beginning of the slice and truncate. bbar.peekBuf = bbar.peekBuf[dstBytesWritten:] - // Move on to read from the promises if we have not satisfied the read yet. + // Move on to the promises if we have not satisfied the read yet. + // Promises are consumed to copy into dst and rotated if exhausted. for dstBytesWritten < len(dst) { headPromise := bbar.bufPromises[bbar.bufIdx] headPromiseBuffered, err := headPromise.Buffered() diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader_test.go b/pkg/storage/indexheader/encoding/bucket_async_reader_test.go index 63194941ec7..935d42e9e4b 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader_test.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader_test.go @@ -43,7 +43,7 @@ func newTestAsyncBufReaderWithData( ), bkt } -func TestBucketAsyncBufReader_Peek_Peek(t *testing.T) { +func TestBucketAsyncBufReader_Peek(t *testing.T) { r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) peek1, err := r.Peek(5) @@ -67,7 +67,34 @@ func TestBucketAsyncBufReader_Peek_Peek(t *testing.T) { require.Equal(t, testBucketContents[:10], peek4) } -func TestBucketAsyncBufReader_Peak_Read(t *testing.T) { +func TestBucketAsyncBufReader_Peek_Skip(t *testing.T) { + r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) + + peek1, err := r.Peek(5) + require.NoError(t, err) + require.Equal(t, testBucketContents[:5], peek1) + require.Equal(t, 0, r.Offset()) + + // Skip the same bytes. + require.NoError(t, r.Skip(5)) + require.Equal(t, 5, r.Offset()) + + // Another Peek returns the next bytes after Skip consumed previous bytes. + peek2, err := r.Peek(5) + require.NoError(t, err) + require.Equal(t, testBucketContents[5:10], peek2) + require.Equal(t, 5, r.Offset()) + + // A Skip short of the peeked bytes consumes some of them. + require.NoError(t, r.Skip(3)) + require.Equal(t, 8, r.Offset()) + + // A Skip beyond the remaining peeked bytes consumes them and more. + require.NoError(t, r.Skip(8)) + require.Equal(t, 16, r.Offset()) +} + +func TestBucketAsyncBufReader_Peek_Read(t *testing.T) { r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) peek1, err := r.Peek(5) From ca47b411b2b8faf5ebceed7fc06dadad260eed8e Mon Sep 17 00:00:00 2001 From: francoposa Date: Sat, 29 Aug 2026 20:15:22 -0700 Subject: [PATCH 5/9] Finsh bucket async reader impl (more tests TODO); fix capacity bug in sliding unconsumed peekBuf bytes;implement bucket read cancellation for buffer promise release; convert resetReader func from bucket reader to something normal --- .../encoding/bucket_async_reader.go | 160 +++++++----------- .../encoding/bucket_async_reader_test.go | 14 +- .../indexheader/encoding/bucket_reader.go | 51 +++--- .../indexheader/encoding/file_reader.go | 3 +- pkg/storage/indexheader/encoding/reader.go | 6 +- 5 files changed, 98 insertions(+), 136 deletions(-) diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader.go b/pkg/storage/indexheader/encoding/bucket_async_reader.go index 9d2639d040b..cfe384489bd 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader.go @@ -5,6 +5,7 @@ package encoding import ( "bufio" "context" + "errors" "io" "sync" @@ -17,19 +18,10 @@ const ( ) type BufPromise struct { - //r *BucketReader - eg *errgroup.Group - //bkt objstore.BucketReader - - //ctx context.Context - //name string - //base int - //length int - //off int - - //buf []byte bufioReader *bufio.Reader - //readN int + + cancel context.CancelCauseFunc + eg *errgroup.Group } func NewBufPromise( @@ -38,14 +30,18 @@ func NewBufPromise( name string, base int, length int, - //buf []byte, bufioReader *bufio.Reader, ) *BufPromise { + // Create handle to cancel an inflight bucket read + ctx, cancel := context.WithCancelCause(ctx) bucketReader := NewBucketReader(ctx, bkt, name, base, length) bufioReader.Reset(bucketReader) + // Propagate cancellable context to errgroup. + eg, _ := errgroup.WithContext(ctx) bp := &BufPromise{ - eg: &errgroup.Group{}, + cancel: cancel, + eg: eg, bufioReader: bufioReader, } @@ -76,17 +72,6 @@ func (bp *BufPromise) Buffered() (n int, err error) { return bp.bufioReader.Buffered(), nil } -// Peek returns at most n bytes from the promise, without consuming them. -// The byte slice points into the buffer of the promise. -// It becomes invalid when the promise is released. -func (bp *BufPromise) Peek(n int) ([]byte, error) { - if err := bp.eg.Wait(); err != nil { - // Return any error in the underlying bucket read from the initial fill-via-Peek. - return nil, err - } - return bp.bufioReader.Peek(n) -} - // Discard consumes and drops at most n bytes from the promise. func (bp *BufPromise) Discard(n int) (discarded int, err error) { if err := bp.eg.Wait(); err != nil { @@ -96,13 +81,9 @@ func (bp *BufPromise) Discard(n int) (discarded int, err error) { return bp.bufioReader.Discard(n) } -// release returns the buffer of the promise to the pool. -// release waits for the fill, because a fill writes to the buffer. -// A buffer in the pool belongs to the next reader that gets it. -// The promise is not usable after release. -func (bp *BufPromise) release(bufioPool *sync.Pool) { - // The error is not relevant here, because release discards the contents of the buffer. - _ = bp.eg.Wait() +func (bp *BufPromise) Release(bufioPool *sync.Pool, cancelCause error) { + bp.cancel(cancelCause) + _ = bp.eg.Wait() // Ensure any write to the buffer completes. bufioPool.Put(bp.bufioReader) bp.bufioReader = nil } @@ -115,25 +96,22 @@ type BucketAsyncBufReader struct { length int readOffset int - resetReader func(off int) error - - bufSize int - + bufSize int + peekBuf []byte bufIdx int bufPromises []*BufPromise bufferedOffset int - - peekBuf []byte - - // pool reference to return to on Close - bufioPool *sync.Pool + bufioPool *sync.Pool } func NewBucketAsyncBufReader( - ctx context.Context, bkt objstore.BucketReader, name string, base int, length int, + ctx context.Context, + bkt objstore.BucketReader, + name string, base int, length int, + ) *BucketAsyncBufReader { return newBucketAsyncBufReader( - ctx, bkt, name, base, length, + ctx, bkt, name, base, length, 0, &bucketBufioPool, ReadBufferSize, ReadAheadFactor, ) } @@ -144,6 +122,7 @@ func newBucketAsyncBufReader( name string, base int, length int, + startOffset int, bufioPool *sync.Pool, bufSize int, maxBufCount int, @@ -152,8 +131,8 @@ func newBucketAsyncBufReader( numBufs := min(maxBufCount, bufsForLength) bufPromises := make([]*BufPromise, numBufs) - iBase := base - bufferedOffset := 0 + iBase := base + startOffset + bufferedOffset := startOffset for i := range numBufs { bufioReader := bufioPool.Get().(*bufio.Reader) bufLen := min(length-bufferedOffset, bufSize) @@ -168,6 +147,7 @@ func newBucketAsyncBufReader( name: name, base: base, length: length, + readOffset: startOffset, bufSize: bufSize, peekBuf: make([]byte, 0, bufSize), bufPromises: bufPromises, @@ -177,13 +157,26 @@ func newBucketAsyncBufReader( } func (bbar *BucketAsyncBufReader) Reset() error { - //TODO implement me - panic("implement me") + return bbar.ResetAt(0) } func (bbar *BucketAsyncBufReader) ResetAt(off int) error { - //TODO implement me - panic("implement me") + if off > bbar.length { + return ErrInvalidSize + } + + if dist := off - bbar.readOffset; dist > 0 && dist < bbar.Buffered() { + // Reset via Skip to avoid discarding all buffered bytes. + return bbar.Skip(dist) + } + + bbar.Close() + newBbar := newBucketAsyncBufReader( + bbar.ctx, bbar.bkt, bbar.name, bbar.base, bbar.length, off, + bbar.bufioPool, bbar.bufSize, ReadAheadFactor, + ) + *bbar = *newBbar + return nil } // rotateHead releases the drained head promise, starts a promise for the next chunk @@ -191,15 +184,10 @@ func (bbar *BucketAsyncBufReader) ResetAt(off int) error { // Peek copies its result into peekBuf, so no slice that the reader returned // points into the buffer of a promise. The release is safe at once. func (bbar *BucketAsyncBufReader) rotateHead() { - // Note that we don't do anything to clean up the buffer before returning it to the pool here: - // we reset the buffer when we retrieve it from the pool instead. - bbar.bufPromises[bbar.bufIdx].release(bbar.bufioPool) + // No need to clean up buffer, we reset when we retrieve it from the pool + bbar.bufPromises[bbar.bufIdx].Release(bbar.bufioPool, nil) bufioReader := bbar.bufioPool.Get().(*bufio.Reader) - // Create a new buffer promise in the same spot in the buffer promise queue. - // The promise must not reach past the end of the reader. - // bufferedOffset is relative to the data segment. - // Add base to get the offset in the object. bufLen := min(bbar.length-bbar.bufferedOffset, bbar.bufSize) bbar.bufPromises[bbar.bufIdx] = NewBufPromise( bbar.ctx, bbar.bkt, bbar.name, bbar.base+bbar.bufferedOffset, bufLen, bufioReader, @@ -220,7 +208,8 @@ func (bbar *BucketAsyncBufReader) Skip(l int) error { bbar.readOffset += bytesSkipped // Slide any unconsumed bytes from peekBuf to the beginning of the slice and truncate. - bbar.peekBuf = bbar.peekBuf[bytesSkipped:] + n := copy(bbar.peekBuf, bbar.peekBuf[bytesSkipped:]) + bbar.peekBuf = bbar.peekBuf[:n] // Move on to the promises if we have not satisfied the skip yet. // Promises are consumed to discard the data and rotated if exhausted. @@ -256,7 +245,7 @@ func (bbar *BucketAsyncBufReader) Peek(n int) ([]byte, error) { // Clamp n to the lesser of the capacity of peekBuf or the length of the section. n = min(n, cap(bbar.peekBuf), bbar.Len()) - // Start with any previously-peeked bytes - Peek-after-Peek is a valid access pattern. + // Start with any previously-peeked bytes. // Any data remaining in peekBuf is assumed to still be the valid start of a Peek. // The read operations (Read/ReadInto and Skip) are required to update peekBuf // to discard any previously-peeked bytes which were consumed by the read. @@ -310,32 +299,19 @@ func (bbar *BucketAsyncBufReader) Read(n int) ([]byte, error) { } func (bbar *BucketAsyncBufReader) ReadInto(dst []byte) error { - // TODO consistency in error handling - //// A read past the end of the data segment is not valid. - //// The guard in Skip returns ErrInvalidSize for the same condition. - //// The reader contract also requires this read to consume the bytes that remain, - //// so move the cursor to the end before the return. - //// Without this guard the loop below never ends, - //// because a drained head promise rotates to an empty promise forever. - //if len(dst) > bbar.Len() { - // remaining := bbar.Len() - // if err := bbar.Skip(remaining); err != nil { - // return err - // } - // // io.ReadFull reports io.EOF for no bytes and io.ErrUnexpectedEOF for a partial read. - // // BucketBufReader passes that error through, so report the same error here. - // shortErr := io.ErrUnexpectedEOF - // if remaining == 0 { - // shortErr = io.EOF - // } - // return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(dst), shortErr) - //} + if len(dst) > bbar.Len() { + if err := bbar.Skip(bbar.Len()); err != nil { + return err + } + return ErrInvalidSize + } // Start with any previously-peeked bytes. dstBytesWritten := copy(dst, bbar.peekBuf) bbar.readOffset += dstBytesWritten // Slide any unconsumed bytes from peekBuf to the beginning of the slice and truncate. - bbar.peekBuf = bbar.peekBuf[dstBytesWritten:] + n := copy(bbar.peekBuf, bbar.peekBuf[dstBytesWritten:]) + bbar.peekBuf = bbar.peekBuf[:n] // Move on to the promises if we have not satisfied the read yet. // Promises are consumed to copy into dst and rotated if exhausted. @@ -362,14 +338,6 @@ func (bbar *BucketAsyncBufReader) ReadInto(dst []byte) error { if headPromiseBuffered <= 0 { bbar.rotateHead() } - - // TODO consistency in error handling - //// Now we can surface any error - //if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { - // return fmt.Errorf("%w reading %d bytes: %s", ErrInvalidSize, len(dst), err) - //} else if err != nil { - // return err - //} } return nil } @@ -392,22 +360,18 @@ func (bbar *BucketAsyncBufReader) Buffered() int { return bbar.bufferedOffset - bbar.readOffset } -// Close releases each promise that the reader still holds, -// and discards the data in the buffers of those promises. -// Close is safe to call more than one time. +// errBufPromiseReleased is the cancellation cause when release stops a fill that is still in flight. +var errBufReaderClosed = errors.New("BufReader closed") + +// Close cancels all promises and releases buffers back to the pool. func (bbar *BucketAsyncBufReader) Close() error { for i, bufPromise := range bbar.bufPromises { - if bufPromise == nil { - // An earlier Close released the promise in this slot. - continue - } - // Note that we don't do anything to clean up the buffer before returning it to the pool here: - // we reset the buffer when we retrieve it from the pool instead. - bufPromise.release(bbar.bufioPool) + // No need to clean up buffer, we reset when we retrieve it from the pool. + bufPromise.Release(bbar.bufioPool, errBufReaderClosed) bbar.bufPromises[i] = nil } // The BucketReader of a promise does not need a Close call. - // It closes the reader from bkt.GetRange in each Read call. + // It closes the reader created by bkt.GetRange on each Read call. return nil } diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader_test.go b/pkg/storage/indexheader/encoding/bucket_async_reader_test.go index 935d42e9e4b..89ebef5e6bb 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader_test.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader_test.go @@ -21,7 +21,7 @@ func newTestAsyncBufReader(t *testing.T, base, length int) (*BucketAsyncBufReade bkt := newTrackingBucket(t, objectData) return newBucketAsyncBufReader( - ctx, bkt, testBucketObjectName, base, length, + ctx, bkt, testBucketObjectName, base, length, 0, &testBucketBufPool, testBufPoolSize, 4, ), bkt } @@ -38,7 +38,7 @@ func newTestAsyncBufReaderWithData( bkt := newTrackingBucket(t, objectData) return newBucketAsyncBufReader( - t.Context(), bkt, testBucketObjectName, base, length, + t.Context(), bkt, testBucketObjectName, base, length, 0, &testBucketBufPool, testBufPoolSize, maxBufCount, ), bkt } @@ -186,6 +186,16 @@ func TestBucketAsyncBufReader_Read_ExactLength(t *testing.T) { require.Equal(t, 0, r.Len()) } +func TestBucketAsyncBufReader_Read_BeyondEnd(t *testing.T) { + const sectionLen = 10 + r, _ := newTestAsyncBufReader(t, 0, sectionLen) + + b, err := r.Read(sectionLen + 1) + require.ErrorIs(t, err, ErrInvalidSize) + require.Nil(t, b) + require.Equal(t, sectionLen, r.Offset()) +} + //func TestBucketAsyncBufReader_Peek_BeyondPeekBuffer(t *testing.T) { // r, _ := newTestAsyncBufReader(t, 0, len(testBucketContents)) // diff --git a/pkg/storage/indexheader/encoding/bucket_reader.go b/pkg/storage/indexheader/encoding/bucket_reader.go index 08dddf21183..21c90be2128 100644 --- a/pkg/storage/indexheader/encoding/bucket_reader.go +++ b/pkg/storage/indexheader/encoding/bucket_reader.go @@ -81,29 +81,16 @@ func (r *BucketReader) Seek(offset int64, whence int) (int64, error) { } type BucketBufReader struct { - ctx context.Context - bkt objstore.BucketReader - name string - base int - length int - off int - r *BucketReader - resetReader func(off int) error - buf *bufio.Reader - // bufPool reference to return to on Close - bufPool *sync.Pool -} - -func resetReaderFunc(bufReader *BucketBufReader) func(off int) error { - return func(off int) error { - r := NewBucketReader(bufReader.ctx, bufReader.bkt, bufReader.name, bufReader.base, bufReader.length) - _, err := r.Seek(int64(off), io.SeekStart) - if err != nil { - return err - } - bufReader.r = r - return nil - } + ctx context.Context + bkt objstore.BucketReader + name string + base int + length int + off int + + r *BucketReader + buf *bufio.Reader + bufPool *sync.Pool // Reference to return to on Close } func NewBucketBufReader( @@ -130,7 +117,6 @@ func newBucketBufReader( bufPool: bufioPool, } - bufReader.resetReader = resetReaderFunc(bufReader) return bufReader } @@ -144,15 +130,17 @@ func (bbr *BucketBufReader) ResetAt(off int) error { } if dist := off - bbr.off; dist > 0 && dist < bbr.Buffered() { - // skip ahead by discarding the distance bytes + // Reset via Skip to avoid discarding all buffered bytes. return bbr.Skip(dist) } - if err := bbr.resetReader(off); err != nil { + r := NewBucketReader(bbr.ctx, bbr.bkt, bbr.name, bbr.base, bbr.length) + _, err := r.Seek(int64(off), io.SeekStart) + if err != nil { return err } - - bbr.buf.Reset(bbr.r) + bbr.r = r + bbr.buf.Reset(r) bbr.off = off return nil @@ -227,10 +215,9 @@ func (bbr *BucketBufReader) Buffered() int { } func (bbr *BucketBufReader) Close() error { - // Note that we don't do anything to clean up the buffer before returning it to the pool here: - // we reset the buffer when we retrieve it from the pool instead. + // No need to clean up buffer, we reset when we retrieve it from the pool bbr.bufPool.Put(bbr.buf) - // The BucketReader does not need closed - - // it closes the reader generated from bkt.GetRange on each Read call. + // The BucketReader of a promise does not need a Close call. + // It closes the reader created by bkt.GetRange on each Read call. return nil } diff --git a/pkg/storage/indexheader/encoding/file_reader.go b/pkg/storage/indexheader/encoding/file_reader.go index bfba1d81b79..0ae0abba2cc 100644 --- a/pkg/storage/indexheader/encoding/file_reader.go +++ b/pkg/storage/indexheader/encoding/file_reader.go @@ -143,8 +143,7 @@ func (f *FileReader) Buffered() int { // Close cleans up the underlying resources used by this FileReader. func (f *FileReader) Close() error { - // Note that we don't do anything to clean up the buffer before returning it to the pool here: - // we reset the buffer when we retrieve it from the pool instead. + // No need to clean up buffer, we reset when we retrieve it from the pool bufferPool.Put(f.buf) // File handles are pooled, so we don't actually close the handle here, just return it. return f.closer.Put(f.file) diff --git a/pkg/storage/indexheader/encoding/reader.go b/pkg/storage/indexheader/encoding/reader.go index 165c4a76a70..c3e93b2ed99 100644 --- a/pkg/storage/indexheader/encoding/reader.go +++ b/pkg/storage/indexheader/encoding/reader.go @@ -19,8 +19,10 @@ type BufReader interface { ResetAt(off int) error // Skip advances the cursor by the given number of bytes in the data segment. - // Attempting to skip to the end of the data segment is valid. - // Attempting to skip _beyond_ the end of the data segment will return an error. + // It is valid to skip to exactly the end of the data segment. + // It is NOT valid to skip beyond the end of the data segment; + // in this case implementations MUST return an ErrInvalidSize error, + // but MUST NOT advance the cursor or consume any remaining bytes. Skip(l int) error // Peek returns at most the given number of bytes from the data segment, without consuming them. From efa30fb6c052c04e0bb40be68092278f1b8b3494 Mon Sep 17 00:00:00 2001 From: francoposa Date: Mon, 31 Aug 2026 02:55:59 +0200 Subject: [PATCH 6/9] Fix async bucket reader bug due to unclear interface contract; improve interface documentation --- .../encoding/bucket_async_reader.go | 83 +++++++++++-------- pkg/storage/indexheader/encoding/reader.go | 6 +- 2 files changed, 53 insertions(+), 36 deletions(-) diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader.go b/pkg/storage/indexheader/encoding/bucket_async_reader.go index cfe384489bd..87488e81c80 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader.go @@ -96,8 +96,17 @@ type BucketAsyncBufReader struct { length int readOffset int - bufSize int - peekBuf []byte + bufSize int + + // peekBuf holds peeked bytes until they are read or discarded. + // Peek is intended to return a slice of bytes without an extra allocation, + // but we cannot return a slice which spans two underlying promise buffers. + // We allocate peekBuf with a large capacity and Peek returns subslices of it. + // peekBuf's slice bounds slides forward through the backing array until it runs out of capacity, + // then it is compacted by copying remaining elements back to the start of the array. + peekBuf []byte + peekBufBase []byte // Holds the reference to the start of the slice for compaction + bufIdx int bufPromises []*BufPromise bufferedOffset int @@ -141,6 +150,7 @@ func newBucketAsyncBufReader( bufferedOffset += bufLen } + peekBufBase := make([]byte, 0, bufSize) return &BucketAsyncBufReader{ ctx: ctx, bkt: bkt, @@ -149,13 +159,31 @@ func newBucketAsyncBufReader( length: length, readOffset: startOffset, bufSize: bufSize, - peekBuf: make([]byte, 0, bufSize), + peekBufBase: peekBufBase, + peekBuf: peekBufBase[:0], bufPromises: bufPromises, bufferedOffset: bufferedOffset, bufioPool: bufioPool, } } +// rotateHead releases the exhausted head promise back to the pool +// and queues a promise to buffer the next read range in its place. +func (bbar *BucketAsyncBufReader) rotateHead() { + // No need to clean up buffer, we reset when we retrieve it from the pool + bbar.bufPromises[bbar.bufIdx].Release(bbar.bufioPool, nil) + bufioReader := bbar.bufioPool.Get().(*bufio.Reader) + + bufLen := min(bbar.length-bbar.bufferedOffset, bbar.bufSize) + bbar.bufPromises[bbar.bufIdx] = NewBufPromise( + bbar.ctx, bbar.bkt, bbar.name, bbar.base+bbar.bufferedOffset, bufLen, bufioReader, + ) + bbar.bufferedOffset += bufLen + + // Advance current buffer queue index - modulo wraps to the front of the slice. + bbar.bufIdx = (bbar.bufIdx + 1) % len(bbar.bufPromises) +} + func (bbar *BucketAsyncBufReader) Reset() error { return bbar.ResetAt(0) } @@ -179,25 +207,6 @@ func (bbar *BucketAsyncBufReader) ResetAt(off int) error { return nil } -// rotateHead releases the drained head promise, starts a promise for the next chunk -// of the data segment in the same slot, and advances the head index. -// Peek copies its result into peekBuf, so no slice that the reader returned -// points into the buffer of a promise. The release is safe at once. -func (bbar *BucketAsyncBufReader) rotateHead() { - // No need to clean up buffer, we reset when we retrieve it from the pool - bbar.bufPromises[bbar.bufIdx].Release(bbar.bufioPool, nil) - bufioReader := bbar.bufioPool.Get().(*bufio.Reader) - - bufLen := min(bbar.length-bbar.bufferedOffset, bbar.bufSize) - bbar.bufPromises[bbar.bufIdx] = NewBufPromise( - bbar.ctx, bbar.bkt, bbar.name, bbar.base+bbar.bufferedOffset, bufLen, bufioReader, - ) - bbar.bufferedOffset += bufLen - - // Advance current buffer queue index - modulo wraps around to the front of the slice if at end. - bbar.bufIdx = (bbar.bufIdx + 1) % len(bbar.bufPromises) -} - func (bbar *BucketAsyncBufReader) Skip(l int) error { if l > bbar.Len() { return ErrInvalidSize @@ -207,9 +216,9 @@ func (bbar *BucketAsyncBufReader) Skip(l int) error { bytesSkipped := min(len(bbar.peekBuf), l) bbar.readOffset += bytesSkipped - // Slide any unconsumed bytes from peekBuf to the beginning of the slice and truncate. - n := copy(bbar.peekBuf, bbar.peekBuf[bytesSkipped:]) - bbar.peekBuf = bbar.peekBuf[:n] + // Advance past the consumed bytes without moving them, + // so a slice returned by an earlier Peek stays valid. + bbar.peekBuf = bbar.peekBuf[bytesSkipped:] // Move on to the promises if we have not satisfied the skip yet. // Promises are consumed to discard the data and rotated if exhausted. @@ -242,16 +251,20 @@ func (bbar *BucketAsyncBufReader) Skip(l int) error { } func (bbar *BucketAsyncBufReader) Peek(n int) ([]byte, error) { - // Clamp n to the lesser of the capacity of peekBuf or the length of the section. - n = min(n, cap(bbar.peekBuf), bbar.Len()) + // Clamp n to the lesser of the buffer size or the length of the section. + n = min(n, bbar.bufSize, bbar.Len()) + if n > cap(bbar.peekBuf) { + // Slide remaining peeked bytes + bbar.peekBuf = append(bbar.peekBufBase[:0], bbar.peekBuf...) + } // Start with any previously-peeked bytes. // Any data remaining in peekBuf is assumed to still be the valid start of a Peek. - // The read operations (Read/ReadInto and Skip) are required to update peekBuf + // Read/ReadInto, Reset/ResetAt, and Skip are required to update peekBuf // to discard any previously-peeked bytes which were consumed by the read. peekBytesAvailable := len(bbar.peekBuf) if n > peekBytesAvailable { - bbar.peekBuf = bbar.peekBuf[:n] // Grow length + bbar.peekBuf = bbar.peekBuf[:n] // Grow length; will not allocate. } peekBytesWritten := min(n, peekBytesAvailable) @@ -309,9 +322,9 @@ func (bbar *BucketAsyncBufReader) ReadInto(dst []byte) error { // Start with any previously-peeked bytes. dstBytesWritten := copy(dst, bbar.peekBuf) bbar.readOffset += dstBytesWritten - // Slide any unconsumed bytes from peekBuf to the beginning of the slice and truncate. - n := copy(bbar.peekBuf, bbar.peekBuf[dstBytesWritten:]) - bbar.peekBuf = bbar.peekBuf[:n] + // Advance past the consumed bytes without moving the + // so a slice returned by an earlier Peek stays valid. + bbar.peekBuf = bbar.peekBuf[dstBytesWritten:] // Move on to the promises if we have not satisfied the read yet. // Promises are consumed to copy into dst and rotated if exhausted. @@ -342,10 +355,10 @@ func (bbar *BucketAsyncBufReader) ReadInto(dst []byte) error { return nil } -// Size returns the largest number of bytes that a single Peek can return. -// Peek assembles its result in peekBuf, so the capacity of peekBuf is the limit. func (bbar *BucketAsyncBufReader) Size() int { - return cap(bbar.peekBuf) + // Reported capacity of peekBuf changes as its referenced window slides, + // but we will compact to make use of its full underlying allocated size if needed. + return bbar.bufSize } func (bbar *BucketAsyncBufReader) Len() int { diff --git a/pkg/storage/indexheader/encoding/reader.go b/pkg/storage/indexheader/encoding/reader.go index c3e93b2ed99..1646e40bb0b 100644 --- a/pkg/storage/indexheader/encoding/reader.go +++ b/pkg/storage/indexheader/encoding/reader.go @@ -28,7 +28,11 @@ type BufReader interface { // Peek returns at most the given number of bytes from the data segment, without consuming them. // It is valid to Peek beyond the end of the data segment; // in this case implementations MUST return the available bytes up to the end and a nil error. - // The byte slice returned becomes invalid at the next read. + // + // The byte slice returned MUST remain valid for one subsequent Skip of the returned byte length; + // callers use a Peek-Skip pattern in place of Read to avoid a slice allocation. + // It is NOT valid to read the returned byte slice after any subsequent read operation: + // Peek, Read, ReadInto, Reset, ResetAt, and Skip. // // Peek is limited to and only must support reads up to the underlying buffer length. // Caller checks Size first to see if the next read operation length fits in the underlying buffer, From 06192fc67a1c9741d1eb8291f7a77d12f89c110b Mon Sep 17 00:00:00 2001 From: francoposa Date: Mon, 31 Aug 2026 03:06:35 +0200 Subject: [PATCH 7/9] Fix error propagation to match existing impls --- pkg/storage/indexheader/encoding/bucket_async_reader.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader.go b/pkg/storage/indexheader/encoding/bucket_async_reader.go index 87488e81c80..5ad70252e80 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader.go @@ -6,6 +6,7 @@ import ( "bufio" "context" "errors" + "fmt" "io" "sync" @@ -48,7 +49,13 @@ func NewBufPromise( bp.eg.Go(func() error { // Send the fill to the background. // Peek reads length bytes into the buffer, and does not consume them. - _, err := bp.bufioReader.Peek(length) + b, err := bp.bufioReader.Peek(length) + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return fmt.Errorf( + "%w reading %d bytes at offset %d of %s (got %d bytes): %s", + ErrInvalidSize, length, base, name, len(b), err, + ) + } return err }) return bp From 25bfbcd4612ea6f0928b8517b7bbe073eca5435c Mon Sep 17 00:00:00 2001 From: francoposa Date: Mon, 31 Aug 2026 03:13:41 +0200 Subject: [PATCH 8/9] hardcode in async bucket reader --- pkg/storage/indexheader/encoding/bucket_factory.go | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pkg/storage/indexheader/encoding/bucket_factory.go b/pkg/storage/indexheader/encoding/bucket_factory.go index 7088b227f99..063598b57bb 100644 --- a/pkg/storage/indexheader/encoding/bucket_factory.go +++ b/pkg/storage/indexheader/encoding/bucket_factory.go @@ -88,7 +88,9 @@ func (bf *BucketDecbufFactory) NewDecbufAtChecked(ctx context.Context, offset in bufferLength := numLenBytes + contentLength + crc32.Size - bufReader := NewBucketBufReader(ctx, bf.bkt, bf.objectPath, offset, bufferLength) + //bufReader := NewBucketBufReader(ctx, bf.bkt, bf.objectPath, offset, bufferLength) + bufReader := NewBucketAsyncBufReader(ctx, bf.bkt, bf.objectPath, offset, bufferLength) + // bufReader is expected start at base offset + 4 after consuming length bytes err = bufReader.Skip(numLenBytes) if err != nil { @@ -128,7 +130,7 @@ func (bf *BucketDecbufFactory) NewDecbufInSection(ctx context.Context, tableOffs if sectionLength <= 0 { return Decbuf{E: fmt.Errorf("section length must be greater than 0")} } - bufReader := NewBucketBufReader( + bufReader := NewBucketAsyncBufReader( ctx, bf.bkt, bf.objectPath, @@ -151,7 +153,7 @@ func (bf *BucketDecbufFactory) NewRawDecbuf(ctx context.Context) Decbuf { return Decbuf{E: fmt.Errorf("get size from %s: %w", bf.objectPath, err)} } // Create reader from full file range - r := NewBucketBufReader( + r := NewBucketAsyncBufReader( ctx, bf.bkt, bf.objectPath, offset, int(attrs.Size), ) d := Decbuf{r: r} From 0887f8f5743cd6776e1ad0c4be9a63c234e77e2d Mon Sep 17 00:00:00 2001 From: francoposa Date: Thu, 3 Sep 2026 06:42:45 +0200 Subject: [PATCH 9/9] readahead factor 2 --- pkg/storage/indexheader/encoding/bucket_async_reader.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/storage/indexheader/encoding/bucket_async_reader.go b/pkg/storage/indexheader/encoding/bucket_async_reader.go index 5ad70252e80..b82346d53e8 100644 --- a/pkg/storage/indexheader/encoding/bucket_async_reader.go +++ b/pkg/storage/indexheader/encoding/bucket_async_reader.go @@ -15,7 +15,7 @@ import ( ) const ( - ReadAheadFactor = 4 + ReadAheadFactor = 2 ) type BufPromise struct {