repo_description: Continuous Profiling Platform. Debug performance issues down to a single line of code
repo_topics: continuous-profiling developer-tools devops find-bottlenecks golang hacktoberfest linux monitoring observability performance profiling pyroscope python ruby
code: package block
  
  import (
  	"bytes"
  	"context"
  	"fmt"
  	"io"
  	"os"
  	"slices"
  	"sort"
  	"strings"
  	"sync"
  
  	"github.com/grafana/dskit/multierror"
  	"github.com/parquet-go/parquet-go"
  	"github.com/prometheus/common/model"
  	"github.com/prometheus/prometheus/storage"
  	"golang.org/x/sync/errgroup"
  
  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
  	"github.com/grafana/pyroscope/pkg/block/metadata"
  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
  	"github.com/grafana/pyroscope/pkg/objstore"
  	"github.com/grafana/pyroscope/pkg/phlaredb/symdb"
  	"github.com/grafana/pyroscope/pkg/phlaredb/tsdb/index"
  	memindex "github.com/grafana/pyroscope/pkg/segmentwriter/memdb/index"
  	"github.com/grafana/pyroscope/pkg/util"
  )
  
  var (
  	ErrNoBlocksToMerge    = fmt.Errorf("no blocks to merge")
  	ErrShardMergeMismatch = fmt.Errorf("only blocks from the same shard can be merged")
  )
  
  type CompactionOption func(*compactionConfig)
  
  func WithCompactionObjectOptions(options ...ObjectOption) CompactionOption {
  	return func(p *compactionConfig) {
  		p.objectOptions = append(p.objectOptions, options...)
  	}
  }
  
  func WithCompactionTempDir(tempdir string) CompactionOption {
  	return func(p *compactionConfig) {
  		p.tempdir = tempdir
  	}
  }
  
  func WithCompactionDestination(storage objstore.Bucket) CompactionOption {
  	return func(p *compactionConfig) {
  		p.destination = storage
  	}
  }
  
  func WithSampleObserver(observer SampleObserver) CompactionOption {
  	return func(p *compactionConfig) {
  		p.sampleObserver = observer
  	}
  }
  
  type compactionConfig struct {
  	objectOptions  []ObjectOption
  	source         objstore.BucketReader
  	destination    objstore.Bucket
  	tempdir        string
  	sampleObserver SampleObserver
  }
  
  type SampleObserver interface {
  	symdb.SymbolsObserver
  
  	// Evaluate is called before the compactor rewrites any symbols.
  	// An "observe" callback function is returned to be called after writing the resulting blocks.
  	// This method must not modify the entry.
  	Evaluate(ProfileEntry) (observe func())
  }
  
  func Compact(
  	ctx context.Context,
  	blocks []*metastorev1.BlockMeta,
  	storage objstore.Bucket,
  	options ...CompactionOption,
  ) (m []*metastorev1.BlockMeta, err error) {
  	c := &compactionConfig{
  		source:      storage,
  		destination: storage,
  		tempdir:     os.TempDir(),
  	}
  	for _, option := range options {
  		option(c)
  	}
  
  	objects := ObjectsFromMetas(storage, blocks, c.objectOptions...)
  	plan, err := PlanCompaction(objects)
  	if err != nil {
  		return nil, err
  	}
  
  	if err = objects.Open(ctx); err != nil {
  		return nil, fmt.Errorf("objects.Open: %w", err)
  	}
  	defer func() {
  		_ = objects.Close()
  	}()
  
  	compacted := make([]*metastorev1.BlockMeta, 0, len(plan))
  	for _, p := range plan {
  		md, compactionErr := p.Compact(ctx, c.destination, c.tempdir, c.sampleObserver)
  		if compactionErr != nil {
  			return nil, compactionErr
  		}
  		compacted = append(compacted, md)
  	}
  
  	return compacted, nil
  }
  
  func PlanCompaction(objects Objects) ([]*CompactionPlan, error) {
  	if len(objects) == 0 {
  		// Even if there's just a single object, we still need to rewrite it.
  		return nil, ErrNoBlocksToMerge
  	}
  
  	r := objects[0]
  	var level uint32
  	for _, obj := range objects {
  		if r.meta.Shard != obj.meta.Shard {
  			return nil, ErrShardMergeMismatch
  		}
  		level = max(level, obj.meta.CompactionLevel)
  	}
  	level++
  
  	g := NewULIDGenerator(objects)
  	m := make(map[string]*CompactionPlan)
  	for _, obj := range objects {
  		for _, ds := range obj.meta.Datasets {
  			if ds.Name == 0 {
  				// Anonymous dataset is never compacted:
  				// it is rebuilt based on the actual block contents.
  				continue
  			}
  			tm, ok := m[obj.meta.StringTable[ds.Tenant]]
  			if !ok {
  				tm = newBlockCompaction(
  					g.ULID().String(),
  					obj.meta.StringTable[ds.Tenant],
  					r.meta.Shard,
  					level,
  				)
  				m[obj.meta.StringTable[ds.Tenant]] = tm
  			}
  			// Bind objects to datasets.
  			sm := tm.addDataset(obj.meta, ds)
  			sm.append(NewDataset(ds, obj))
  		}
  	}
  
  	ordered := make([]*CompactionPlan, 0, len(m))
  	for _, tm := range m {
  		ordered = append(ordered, tm)
  		slices.SortFunc(tm.datasets, func(a, b *datasetCompaction) int {
  			return strings.Compare(a.name, b.name)
  		})
  	}
  	slices.SortFunc(ordered, func(a, b *CompactionPlan) int {
  		return strings.Compare(a.tenant, b.tenant)
  	})
  
  	return ordered, nil
  }
  
  type CompactionPlan struct {
  	tenant       string
  	path         string
  	datasetMap   map[int32]*datasetCompaction
  	datasets     []*datasetCompaction
  	meta         *metastorev1.BlockMeta
  	strings      *metadata.StringTable
  	datasetIndex *datasetIndexWriter
  }
  
  func newBlockCompaction(
  	id string,
  	tenant string,
  	shard uint32,
  	compactionLevel uint32,
  ) *CompactionPlan {
  	p := &CompactionPlan{
  		tenant:       tenant,
  		datasetMap:   make(map[int32]*datasetCompaction),
  		strings:      metadata.NewStringTable(),
  		datasetIndex: newDatasetIndexWriter(),
  	}
  	p.path = BuildObjectPath(tenant, shard, compactionLevel, id)
  	p.meta = &metastorev1.BlockMeta{
  		FormatVersion:   1,
  		Id:              id,
  		Tenant:          p.strings.Put(tenant),
  		Shard:           shard,
  		CompactionLevel: compactionLevel,
  	}
  	return p
  }
  
  func (b *CompactionPlan) Compact(
  	ctx context.Context,
  	dst objstore.Bucket,
  	tempdir string,
  	observer SampleObserver,
  ) (m *metastorev1.BlockMeta, err error) {
  	w, err := NewBlockWriter(tempdir)
  	if err != nil {
  		return nil, fmt.Errorf("creating block writer: %w", err)
  	}
  	defer func() {
  		_ = w.Close()
  	}()
  
  	// Datasets are compacted in a strict order.
  	for i, s := range b.datasets {
  		b.datasetIndex.setIndex(uint32(i))
  		s.registerSampleObserver(observer)
  		if err = s.compact(ctx, w); err != nil {
  			return nil, fmt.Errorf("compacting block: %w", err)
  		}
  		b.meta.Datasets = append(b.meta.Datasets, s.meta)
  	}
  	if err = b.writeDatasetIndex(w); err != nil {
  		return nil, fmt.Errorf("writing tenant index: %w", err)
  	}
  	b.meta.StringTable = b.strings.Strings
  	b.meta.MetadataOffset = w.Offset()
  	if err = metadata.Encode(w, b.meta); err != nil {
  		return nil, fmt.Errorf("writing metadata: %w", err)
  	}
  	b.meta.Size = w.Offset()
  	if err = w.Upload(ctx, dst, b.path); err != nil {
  		return nil, fmt.Errorf("uploading block: %w", err)
  	}
  	return b.meta, nil
  }
  
  func (b *CompactionPlan) writeDatasetIndex(w *Writer) error {
  	if err := b.datasetIndex.Flush(); err != nil {
  		return err
  	}
  	off := w.Offset()
  	n, err := io.Copy(w, bytes.NewReader(b.datasetIndex.buf))
  	if err != nil {
  		return err
  	}
  	// We annotate the dataset with the
  	// __tenant_dataset__ = "dataset_tsdb_index" label,
  	// so the dataset index metadata can be queried.
  	labels := metadata.NewLabelBuilder(b.strings).
  		WithLabelSet(metadata.LabelNameTenantDataset, metadata.LabelValueDatasetTSDBIndex).
  		Build()
  	b.meta.Datasets = append(b.meta.Datasets, &metastorev1.Dataset{
  		Format:          1,
  		Tenant:          b.meta.Tenant,
  		Name:            0, // Anonymous.
  		MinTime:         b.meta.MinTime,
  		MaxTime:         b.meta.MaxTime,
  		TableOfContents: []uint64{off},
  		Size:            uint64(n),
  		Labels:          labels,
  	})
  	return nil
  }
  
  func (b *CompactionPlan) addDataset(md *metastorev1.BlockMeta, s *metastorev1.Dataset) *datasetCompaction {
  	name := b.strings.Put(md.StringTable[s.Name])
  	tenant := b.strings.Put(md.StringTable[s.Tenant])
  	sm, ok := b.datasetMap[name]
  	if !ok {
  		sm = b.newDatasetCompaction(tenant, name)
  		b.datasetMap[name] = sm
  		b.datasets = append(b.datasets, sm)
  	}
  	if b.meta.MinTime == 0 || s.MinTime < b.meta.MinTime {
  		b.meta.MinTime = s.MinTime
  	}
  	if s.MaxTime > b.meta.MaxTime {
  		b.meta.MaxTime = s.MaxTime
  	}
  	return sm
  }
  
  type datasetCompaction struct {
  	// Dataset name.
  	name   string
  	parent *CompactionPlan
  	meta   *metastorev1.Dataset
  	labels *metadata.LabelBuilder
  
  	datasets []*Dataset
  
  	indexRewriter   *indexRewriter
  	symbolsRewriter *symbolsRewriter
  	profilesWriter  *profilesWriter
  
  	samples  uint64
  	series   uint64
  	profiles uint64
  
  	flushOnce sync.Once
  
  	observer SampleObserver
  }
  
  func (b *CompactionPlan) newDatasetCompaction(tenant, name int32) *datasetCompaction {
  	return &datasetCompaction{
  		parent: b,
  		name:   b.strings.Strings[name],
  		labels: metadata.NewLabelBuilder(b.strings),
  		meta: &metastorev1.Dataset{
  			Tenant: tenant,
  			Name:   name,
  			// Updated at append.
  			MinTime: 0,
  			MaxTime: 0,
  			// Updated at writeTo.
  			TableOfContents: nil,
  			Size:            0,
  			Labels:          nil,
  		},
  	}
  }
  
  func (m *datasetCompaction) append(s *Dataset) {
  	m.datasets = append(m.datasets, s)
  	if m.meta.MinTime == 0 || s.meta.MinTime < m.meta.MinTime {
  		m.meta.MinTime = s.meta.MinTime
  	}
  	if s.meta.MaxTime > m.meta.MaxTime {
  		m.meta.MaxTime = s.meta.MaxTime
  	}
  	m.labels.Put(s.meta.Labels, s.obj.meta.StringTable)
  }
  
  func (m *datasetCompaction) compact(ctx context.Context, w *Writer) (err error) {
  	off := w.Offset()
  	m.meta.TableOfContents = make([]uint64, 0, 3)
  	m.meta.TableOfContents = append(m.meta.TableOfContents, w.Offset())
  
  	if err = m.open(ctx, w); err != nil {
  		return fmt.Errorf("failed to open sections for compaction: %w", err)
  	}
  	defer func() {
  		_ = m.close()
  	}()
  
  	if err = m.merge(ctx); err != nil {
  		return fmt.Errorf("failed to merge datasets: %w", err)
  	}
  	if err = m.flush(); err != nil {
  		return fmt.Errorf("failed to flush compacted dataset: %w", err)
  	}
  
  	m.meta.TableOfContents = append(m.meta.TableOfContents, w.Offset())
  	if _, err = io.Copy(w, bytes.NewReader(m.indexRewriter.buf)); err != nil {
  		return fmt.Errorf("failed to read index: %w", err)
  	}
  	m.meta.TableOfContents = append(m.meta.TableOfContents, w.Offset())
  	if _, err = io.Copy(w, bytes.NewReader(m.symbolsRewriter.buf.Bytes())); err != nil {
  		return fmt.Errorf("failed to read symbols: %w", err)
  	}
  
  	m.meta.Size = w.Offset() - off
  	m.meta.Labels = m.labels.Build()
  	return nil
  }
  
  func (m *datasetCompaction) registerSampleObserver(observer SampleObserver) {
  	m.observer = observer
  }
  
  func (m *datasetCompaction) open(ctx context.Context, w io.Writer) (err error) {
  	var estimatedProfileTableSize int64
  	for _, ds := range m.datasets {
  		estimatedProfileTableSize += ds.sectionSize(SectionProfiles)
  	}
  	pageBufferSize := estimatePageBufferSize(estimatedProfileTableSize)
  	m.profilesWriter = newProfileWriter(pageBufferSize, w)
  
  	m.indexRewriter = newIndexRewriter()
  	m.symbolsRewriter = newSymbolsRewriter(m.observer)
  
  	g, ctx := errgroup.WithContext(ctx)
  	for _, s := range m.datasets {
  		s := s
  		g.Go(util.RecoverPanic(func() error {
  			if openErr := s.Open(ctx,
  				SectionProfiles,
  				SectionTSDB,
  				SectionSymbols,
  			); openErr != nil {
  				return fmt.Errorf("opening tenant dataset (block %s): %w", s.obj.path, openErr)
  			}
  			return nil
  		}))
  	}
  	if err = g.Wait(); err != nil {
  		merr := multierror.New(err)
  		for _, s := range m.datasets {
  			merr.Add(s.Close())
  		}
  		return merr.Err()
  	}
  	return nil
  }
  
  func (m *datasetCompaction) merge(ctx context.Context) (err error) {
  	rows, err := NewMergeRowProfileIterator(m.datasets)
  	if err != nil {
  		return err
  	}
  	defer func() {
  		err = multierror.New(err, rows.Close()).Err()
  	}()
  	var i int
  	for rows.Next() {
  		if i++; i%1000 == 0 {
  			if err = ctx.Err(); err != nil {
  				return err
  			}
  		}
  		if err = m.writeRow(rows.At()); err != nil {
  			return err
  		}
  	}
  	return rows.Err()
  }
  
  func (m *datasetCompaction) writeRow(r ProfileEntry) (err error) {
  	if m.observer != nil {
  		observe := m.observer.Evaluate(r)
  		defer observe()
  	}
  	m.parent.datasetIndex.writeRow(r)
  	m.indexRewriter.rewriteRow(r)
  	if err = m.symbolsRewriter.rewriteRow(r); err != nil {
  		return err
  	}
  	return m.profilesWriter.writeRow(r)
  }
  
  func (m *datasetCompaction) flush() (err error) {
  	m.flushOnce.Do(func() {
  		merr := multierror.New()
  		merr.Add(m.symbolsRewriter.Flush())
  		merr.Add(m.indexRewriter.Flush())
  		merr.Add(m.profilesWriter.Close())
  		m.samples = m.symbolsRewriter.samples
  		m.series = m.indexRewriter.NumSeries()
  		m.profiles = m.profilesWriter.profiles
  		err = merr.Err()
  	})
  	return err
  }
  
  func (m *datasetCompaction) close() error {
  	err := m.flush()
  	m.symbolsRewriter = nil
  	m.indexRewriter = nil
  	m.profilesWriter = nil
  	m.datasets = nil
  	return err
  }
  
  func newIndexRewriter() *indexRewriter {
  	return &indexRewriter{
  		symbols: make(map[string]struct{}),
  	}
  }
  
  type indexRewriter struct {
  	series     []seriesLabels
  	symbols    map[string]struct{}
  	chunks     []index.ChunkMeta // one chunk per series
  	previousFp model.Fingerprint
  	buf        []byte
  }
  
  type seriesLabels struct {
  	labels      phlaremodel.Labels
  	fingerprint model.Fingerprint
  }
  
  func (rw *indexRewriter) rewriteRow(e ProfileEntry) {
  	if rw.previousFp != e.Fingerprint || len(rw.series) == 0 {
  		series := e.Labels.Clone()
  		for _, l := range series {
  			rw.symbols[l.Name] = struct{}{}
  			rw.symbols[l.Value] = struct{}{}
  		}
  		rw.series = append(rw.series, seriesLabels{
  			labels:      series,
  			fingerprint: e.Fingerprint,
  		})
  		rw.chunks = append(rw.chunks, index.ChunkMeta{
  			MinTime:     e.Timestamp,
  			MaxTime:     e.Timestamp,
  			SeriesIndex: uint32(len(rw.series) - 1),
  		})
  		rw.previousFp = e.Fingerprint
  	}
  	rw.chunks[len(rw.chunks)-1].MaxTime = e.Timestamp
  	e.Row.SetSeriesIndex(rw.chunks[len(rw.chunks)-1].SeriesIndex)
  }
  
  func (rw *indexRewriter) NumSeries() uint64 { return uint64(len(rw.series)) }
  
  func (rw *indexRewriter) Flush() error {
  	// TODO(kolesnikovae):
  	//  * Estimate size.
  	//  * Use buffer pool.
  	w, err := memindex.NewWriter(context.Background(), 256<<10)
  	if err != nil {
  		return err
  	}
  
  	// Sort symbols
  	symbols := make([]string, 0, len(rw.symbols))
  	for s := range rw.symbols {
  		symbols = append(symbols, s)
  	}
  	sort.Strings(symbols)
  
  	// Add symbols
  	for _, symbol := range symbols {
  		if err = w.AddSymbol(symbol); err != nil {
  			return err
  		}
  	}
  
  	// Add Series
  	for i, series := range rw.series {
  		if err = w.AddSeries(storage.SeriesRef(i), series.labels, series.fingerprint, rw.chunks[i]); err != nil {
  			return err
  		}
  	}
  
  	err = w.Close()
  	rw.buf = w.ReleaseIndex()
  	return err
  }
  
  type symbolsRewriter struct {
  	buf      *bytes.Buffer
  	w        *symdb.SymDB
  	rw       map[*Dataset]*symdb.Rewriter
  	samples  uint64
  	observer SampleObserver
  
  	stacktraces []uint32
  }
  
  func newSymbolsRewriter(observer SampleObserver) *symbolsRewriter {
  	// TODO(kolesnikovae):
  	//  * Estimate size.
  	//  * Use buffer pool.
  	buf := bytes.NewBuffer(make([]byte, 0, 1<<20))
  	return &symbolsRewriter{
  		buf: buf,
  		rw:  make(map[*Dataset]*symdb.Rewriter),
  		w: symdb.NewSymDB(&symdb.Config{
  			Version: symdb.FormatV3,
  			Writer:  &nopWriteCloser{buf},
  		}),
  		observer: observer,
  	}
  }
  
  type nopWriteCloser struct{ io.Writer }
  
  func (*nopWriteCloser) Close() error { return nil }
  
  func (s *symbolsRewriter) rewriteRow(e ProfileEntry) (err error) {
  	rw := s.rewriterFor(e.Dataset)
  	e.Row.ForStacktraceIDsValues(func(values []parquet.Value) {
  		s.loadStacktraceIDs(values)
  		if err = rw.Rewrite(e.Row.StacktracePartitionID(), s.stacktraces); err != nil {
  			return
  		}
  		s.samples += uint64(len(values))
  		for i, v := range values {
  			values[i] = parquet.Int64Value(int64(s.stacktraces[i])).Level(v.RepetitionLevel(), v.DefinitionLevel(), v.Column())
  		}
  	})
  	return err
  }
  
  func (s *symbolsRewriter) rewriterFor(x *Dataset) *symdb.Rewriter {
  	rw, ok := s.rw[x]
  	if !ok {
  		rw = symdb.NewRewriter(s.w, x.Symbols(), s.observer)
  		s.rw[x] = rw
  	}
  	return rw
  }
  
  func (s *symbolsRewriter) loadStacktraceIDs(values []parquet.Value) {
  	s.stacktraces = slices.Grow(s.stacktraces[0:], len(values))[:len(values)]
  	for i := range values {
  		s.stacktraces[i] = values[i].Uint32()
  	}
  }
  
  func (s *symbolsRewriter) Flush() error { return s.w.Flush() }
  
  // datasetIndexWriter is identical with indexRewriter,
  // except it writes dataset ID instead of series ID.
  type datasetIndexWriter struct {
  	series   []seriesLabels
  	chunks   []index.ChunkMeta
  	previous model.Fingerprint
  	symbols  map[string]struct{}
  	idx      uint32
  	buf      []byte
  }
  
  func newDatasetIndexWriter() *datasetIndexWriter {
  	return &datasetIndexWriter{
  		symbols: make(map[string]struct{}),
  	}
  }
  
  func (rw *datasetIndexWriter) setIndex(i uint32) { rw.idx = i }
  
  func (rw *datasetIndexWriter) writeRow(e ProfileEntry) {
  	if rw.previous != e.Fingerprint || len(rw.series) == 0 {
  		series := e.Labels.Clone()
  		for _, l := range series {
  			rw.symbols[l.Name] = struct{}{}
  			rw.symbols[l.Value] = struct{}{}
  		}
  		rw.series = append(rw.series, seriesLabels{
  			labels:      series,
  			fingerprint: e.Fingerprint,
  		})
  		rw.chunks = append(rw.chunks, index.ChunkMeta{
  			SeriesIndex: rw.idx,
  		})
  		rw.previous = e.Fingerprint
  	}
  }
  
  func (rw *datasetIndexWriter) Flush() error {
  	// TODO(kolesnikovae):
  	//  * Estimate size.
  	//  * Use buffer pool.
  	w, err := memindex.NewWriter(context.Background(), 1<<20)
  	if err != nil {
  		return err
  	}
  
  	// Sort symbols
  	symbols := make([]string, 0, len(rw.symbols))
  	for s := range rw.symbols {
  		symbols = append(symbols, s)
  	}
  	sort.Strings(symbols)
  
  	// Add symbols
  	for _, symbol := range symbols {
  		if err = w.AddSymbol(symbol); err != nil {
  			return err
  		}
  	}
  
  	// Add Series
  	for i, series := range rw.series {
  		if err = w.AddSeries(storage.SeriesRef(i), series.labels, series.fingerprint, rw.chunks[i]); err != nil {
  			return err
  		}
  	}
  
  	err = w.Close()
  	rw.buf = w.ReleaseIndex()
  	return err
  }
  