github.com/grafana/pyroscope@v1.18.0/pkg/block/section_tsdb.go (about)

     1  package block
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"github.com/grafana/pyroscope/pkg/objstore"
     8  	"github.com/grafana/pyroscope/pkg/phlaredb/tsdb/index"
     9  	"github.com/grafana/pyroscope/pkg/util/bufferpool"
    10  )
    11  
    12  func openDatasetIndex(ctx context.Context, s *Dataset) error {
    13  	return openTSDB(ctx, s)
    14  }
    15  
    16  func openTSDB(ctx context.Context, s *Dataset) (err error) {
    17  	offset := s.sectionOffset(SectionTSDB)
    18  	size := s.sectionSize(SectionTSDB)
    19  	s.tsdb = new(tsdbBuffer)
    20  	defer func() {
    21  		if err != nil {
    22  			_ = s.tsdb.Close()
    23  		}
    24  	}()
    25  	if buf := s.inMemoryBuffer(); buf != nil {
    26  		offset -= int64(s.offset())
    27  		s.tsdb.index, err = index.NewReader(index.RealByteSlice(buf[offset : offset+size]))
    28  	} else {
    29  		s.tsdb.buf = bufferpool.GetBuffer(int(size))
    30  		if err = objstore.ReadRange(ctx, s.tsdb.buf, s.obj.path, s.obj.storage, offset, size); err == nil {
    31  			s.tsdb.index, err = index.NewReader(index.RealByteSlice(s.tsdb.buf.B))
    32  		}
    33  	}
    34  	if err != nil {
    35  		return fmt.Errorf("opening tsdb: %w", err)
    36  	}
    37  	return nil
    38  }
    39  
    40  type tsdbBuffer struct {
    41  	index *index.Reader
    42  	buf   *bufferpool.Buffer
    43  }
    44  
    45  func (b *tsdbBuffer) Close() (err error) {
    46  	if b.index != nil {
    47  		err = b.index.Close()
    48  		b.index = nil
    49  	}
    50  	if b.buf != nil {
    51  		bufferpool.Put(b.buf)
    52  		b.buf = nil
    53  	}
    54  	return err
    55  }