github.com/grafana/pyroscope@v1.18.0/pkg/objstore/reader.go (about)

     1  package objstore
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  
     8  	"github.com/thanos-io/objstore"
     9  )
    10  
    11  type ReaderAtCloser interface {
    12  	io.ReaderAt
    13  	io.Closer
    14  }
    15  
    16  func NewBucket(bkt objstore.Bucket) Bucket {
    17  	if bucket, ok := bkt.(Bucket); ok {
    18  		return bucket
    19  	}
    20  	return &ReaderAtBucket{
    21  		Bucket: bkt,
    22  	}
    23  }
    24  
    25  type ReaderAtBucket struct {
    26  	objstore.Bucket
    27  }
    28  
    29  func (b *ReaderAtBucket) ReaderAt(ctx context.Context, name string) (ReaderAtCloser, error) {
    30  	return &ReaderAt{
    31  		GetRangeReader: b.Bucket,
    32  		name:           name,
    33  		ctx:            ctx,
    34  	}, nil
    35  }
    36  
    37  // ReaderWithExpectedErrs implements objstore.Bucket.
    38  func (b *ReaderAtBucket) ReaderWithExpectedErrs(fn IsOpFailureExpectedFunc) BucketReader {
    39  	return b.WithExpectedErrs(fn)
    40  }
    41  
    42  // WithExpectedErrs implements objstore.Bucket.
    43  func (b *ReaderAtBucket) WithExpectedErrs(fn IsOpFailureExpectedFunc) Bucket {
    44  	if ib, ok := b.Bucket.(InstrumentedBucket); ok {
    45  		return &ReaderAtBucket{
    46  			Bucket: ib.WithExpectedErrs(fn),
    47  		}
    48  	}
    49  
    50  	if ib, ok := b.Bucket.(objstore.InstrumentedBucket); ok {
    51  		return &ReaderAtBucket{
    52  			Bucket: ib.WithExpectedErrs(func(err error) bool { return fn(err) }),
    53  		}
    54  	}
    55  
    56  	return b
    57  }
    58  
    59  type GetRangeReader interface {
    60  	GetRange(ctx context.Context, name string, off, length int64) (io.ReadCloser, error)
    61  }
    62  
    63  type ReaderAt struct {
    64  	GetRangeReader
    65  	name string
    66  	ctx  context.Context
    67  }
    68  
    69  func (b *ReaderAt) ReadAt(p []byte, off int64) (int, error) {
    70  	rc, err := b.GetRange(b.ctx, b.name, off, int64(len(p)))
    71  	if err != nil {
    72  		return 0, err
    73  	}
    74  	defer rc.Close()
    75  
    76  	totalBytes := 0
    77  	for {
    78  		byteCount, err := rc.Read(p[totalBytes:])
    79  		totalBytes += byteCount
    80  		if err == io.EOF {
    81  			return totalBytes, nil
    82  		}
    83  		if err != nil {
    84  			return totalBytes, err
    85  		}
    86  		if byteCount == 0 {
    87  			return totalBytes, nil
    88  		}
    89  	}
    90  }
    91  
    92  func (b *ReaderAt) Close() error {
    93  	return nil
    94  }
    95  
    96  func ReadRange(ctx context.Context, reader io.ReaderFrom, name string, storage objstore.BucketReader, off, size int64) error {
    97  	if size == 0 {
    98  		attrs, err := storage.Attributes(ctx, name)
    99  		if err != nil {
   100  			return err
   101  		}
   102  		size = attrs.Size
   103  	}
   104  	if size == 0 {
   105  		return nil
   106  	}
   107  	rc, err := storage.GetRange(ctx, name, off, size)
   108  	if err != nil {
   109  		return err
   110  	}
   111  	defer func() {
   112  		_ = rc.Close()
   113  	}()
   114  	n, err := reader.ReadFrom(io.LimitReader(rc, size))
   115  	if err != nil {
   116  		return err
   117  	}
   118  	if n != size {
   119  		return fmt.Errorf("read %d bytes, expected %d", n, size)
   120  	}
   121  	return nil
   122  }
   123  
   124  type BucketReaderWithOffset struct {
   125  	BucketReader
   126  	offset int64
   127  }
   128  
   129  func NewBucketReaderWithOffset(r BucketReader, offset int64) *BucketReaderWithOffset {
   130  	return &BucketReaderWithOffset{
   131  		BucketReader: r,
   132  		offset:       offset,
   133  	}
   134  }
   135  
   136  func (r *BucketReaderWithOffset) GetRange(ctx context.Context, name string, off, length int64) (io.ReadCloser, error) {
   137  	return r.BucketReader.GetRange(ctx, name, r.offset+off, length)
   138  }