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

     1  package objstore
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"path/filepath"
     9  	"sync"
    10  	"time"
    11  
    12  	"github.com/grafana/dskit/multierror"
    13  	"github.com/thanos-io/objstore"
    14  
    15  	"github.com/grafana/pyroscope/pkg/util/bufferpool"
    16  )
    17  
    18  var _ objstore.BucketReader = &ReadOnlyFile{}
    19  
    20  type ReadOnlyFile struct {
    21  	size    int64
    22  	name    string
    23  	path    string
    24  	mu      sync.Mutex
    25  	readers []*fileReader
    26  }
    27  
    28  func Download(ctx context.Context, name string, src BucketReader, dir string) (*ReadOnlyFile, error) {
    29  	f, err := download(ctx, name, src, dir)
    30  	if err != nil {
    31  		return nil, fmt.Errorf("downloading %s: %w", name, err)
    32  	}
    33  	return f, nil
    34  }
    35  
    36  func download(ctx context.Context, name string, src BucketReader, dir string) (*ReadOnlyFile, error) {
    37  	r, err := src.Get(ctx, name)
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  	defer r.Close()
    42  
    43  	path := filepath.Join(dir, filepath.Base(name))
    44  	f := &ReadOnlyFile{
    45  		name: name,
    46  		path: path,
    47  	}
    48  	if err = os.MkdirAll(dir, 0755); err != nil {
    49  		return nil, err
    50  	}
    51  	dst, err := os.Create(path)
    52  	if err != nil {
    53  		return nil, err
    54  	}
    55  	defer dst.Close()
    56  
    57  	buf := bufferpool.GetBuffer(32 << 10)
    58  	defer bufferpool.Put(buf)
    59  	buf.B = buf.B[:cap(buf.B)]
    60  	n, err := io.CopyBuffer(dst, r, buf.B)
    61  	if err != nil {
    62  		_ = os.RemoveAll(path)
    63  		return nil, err
    64  	}
    65  
    66  	f.size = n
    67  	return f, nil
    68  }
    69  
    70  func (f *ReadOnlyFile) Close() error {
    71  	var m multierror.MultiError
    72  	for _, r := range f.readers {
    73  		m.Add(r.Close())
    74  	}
    75  	m.Add(os.RemoveAll(f.path))
    76  	f.readers = f.readers[:0]
    77  	return m.Err()
    78  }
    79  
    80  func (f *ReadOnlyFile) Iter(context.Context, string, func(string) error, ...objstore.IterOption) error {
    81  	return nil
    82  }
    83  
    84  func (f *ReadOnlyFile) IterWithAttributes(context.Context, string, func(attrs objstore.IterObjectAttributes) error, ...objstore.IterOption) error {
    85  	return nil
    86  }
    87  
    88  func (f *ReadOnlyFile) SupportedIterOptions() []objstore.IterOptionType {
    89  	return nil
    90  }
    91  
    92  func (f *ReadOnlyFile) Exists(_ context.Context, name string) (bool, error) {
    93  	return name == f.name, nil
    94  }
    95  
    96  func (f *ReadOnlyFile) IsObjNotFoundErr(err error) bool { return os.IsNotExist(err) }
    97  
    98  func (f *ReadOnlyFile) IsAccessDeniedErr(err error) bool { return os.IsPermission(err) }
    99  
   100  func (f *ReadOnlyFile) Attributes(_ context.Context, name string) (attrs objstore.ObjectAttributes, err error) {
   101  	if name != f.name {
   102  		return attrs, os.ErrNotExist
   103  	}
   104  	return objstore.ObjectAttributes{
   105  		Size:         f.size,
   106  		LastModified: time.Unix(0, 0), // We don't care.
   107  	}, nil
   108  }
   109  
   110  func (f *ReadOnlyFile) ReaderAt(_ context.Context, name string) (ReaderAtCloser, error) {
   111  	return f.borrowOrCreateReader(name)
   112  }
   113  
   114  func (f *ReadOnlyFile) Get(_ context.Context, name string) (io.ReadCloser, error) {
   115  	r, err := f.borrowOrCreateReader(name)
   116  	if err != nil {
   117  		return nil, err
   118  	}
   119  	if _, err = r.Seek(0, io.SeekStart); err != nil {
   120  		_ = r.Close()
   121  		return nil, err
   122  	}
   123  	return r, nil
   124  }
   125  
   126  func (f *ReadOnlyFile) GetRange(_ context.Context, name string, off, length int64) (io.ReadCloser, error) {
   127  	if off < 0 || length < 0 {
   128  		return nil, fmt.Errorf("%w: invalid offset", os.ErrInvalid)
   129  	}
   130  	r, err := f.borrowOrCreateReader(name)
   131  	if err != nil {
   132  		return nil, err
   133  	}
   134  	if _, err = r.Seek(off, io.SeekStart); err != nil {
   135  		_ = r.Close()
   136  		return nil, err
   137  	}
   138  	r.reader = io.LimitReader(r.reader, length)
   139  	return r, nil
   140  }
   141  
   142  func (f *ReadOnlyFile) borrowOrCreateReader(name string) (*fileReader, error) {
   143  	if name != f.name {
   144  		return nil, os.ErrNotExist
   145  	}
   146  	f.mu.Lock()
   147  	defer f.mu.Unlock()
   148  	if len(f.readers) > 0 {
   149  		ff := f.readers[len(f.readers)-1]
   150  		f.readers = f.readers[:len(f.readers)-1]
   151  		ff.reader = ff.File
   152  		return ff, nil
   153  	}
   154  	return f.openReader()
   155  }
   156  
   157  func (f *ReadOnlyFile) returnReader(r *fileReader) {
   158  	f.mu.Lock()
   159  	defer f.mu.Unlock()
   160  	f.readers = append(f.readers, r)
   161  }
   162  
   163  func (f *ReadOnlyFile) openReader() (*fileReader, error) {
   164  	ff, err := os.Open(f.path)
   165  	if err != nil {
   166  		return nil, err
   167  	}
   168  	return &fileReader{
   169  		parent: f,
   170  		File:   ff,
   171  		reader: ff,
   172  	}, nil
   173  }
   174  
   175  type fileReader struct {
   176  	parent *ReadOnlyFile
   177  	reader io.Reader
   178  	*os.File
   179  }
   180  
   181  func (r *fileReader) Close() error {
   182  	r.reader = nil
   183  	r.parent.returnReader(r)
   184  	return nil
   185  }
   186  
   187  func (r *fileReader) Read(p []byte) (int, error) {
   188  	return r.reader.Read(p)
   189  }
   190  
   191  func (r *fileReader) Provider(p []byte) objstore.ObjProvider {
   192  	return objstore.ObjProvider("READONLYFILE")
   193  }