feat: record metrics from rules and export to remote

Author: alsoba13

pkg/experiment/block/compaction.go

@@ -10,14 +10,18 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/cespare/xxhash/v2"
+	"github.com/google/uuid"
 	"github.com/grafana/dskit/multierror"
+	"github.com/oklog/ulid"
 	"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/experiment/block/metadata"
+	"github.com/grafana/pyroscope/pkg/experiment/metrics"
 	phlaremodel "github.com/grafana/pyroscope/pkg/model"
 	"github.com/grafana/pyroscope/pkg/objstore"
 	"github.com/grafana/pyroscope/pkg/phlaredb/block"
@@ -62,6 +66,7 @@ func Compact(
 	ctx context.Context,
 	blocks []*metastorev1.BlockMeta,
 	storage objstore.Bucket,
+	workerId uuid.UUID,
 	options ...CompactionOption,
 ) (m []*metastorev1.BlockMeta, err error) {
 	c := &compactionConfig{
@@ -74,7 +79,7 @@ func Compact(
 	}
 
 	objects := ObjectsFromMetas(storage, blocks, c.objectOptions...)
-	plan, err := PlanCompaction(objects)
+	plan, err := PlanCompaction(objects, workerId)
 	if err != nil {
 		return nil, err
 	}
@@ -93,12 +98,20 @@ func Compact(
 			return nil, compactionErr
 		}
 		compacted = append(compacted, md)
+
+		if p.metricsExporter != nil {
+			go func() {
+				if sendErr := p.SendRecordedMetrics(); sendErr != nil {
+					println("ERROR", sendErr) // TODO
+				}
+			}()
+		}
 	}
 
 	return compacted, nil
 }
 
-func PlanCompaction(objects Objects) ([]*CompactionPlan, error) {
+func PlanCompaction(objects Objects, workerId uuid.UUID) ([]*CompactionPlan, error) {
 	if len(objects) == 0 {
 		// Even if there's just a single object, we still need to rewrite it.
 		return nil, ErrNoBlocksToMerge
@@ -125,6 +138,7 @@ func PlanCompaction(objects Objects) ([]*CompactionPlan, error) {
 					obj.meta.StringTable[s.Tenant],
 					r.meta.Shard,
 					level,
+					workerId,
 				)
 				m[obj.meta.StringTable[s.Tenant]] = tm
 			}
@@ -149,24 +163,28 @@ func PlanCompaction(objects Objects) ([]*CompactionPlan, error) {
 }
 
 type CompactionPlan struct {
-	tenant     string
-	path       string
-	datasetMap map[int32]*datasetCompaction
-	datasets   []*datasetCompaction
-	meta       *metastorev1.BlockMeta
-	strings    *metadata.StringTable
+	tenant          string
+	path            string
+	datasetMap      map[int32]*datasetCompaction
+	datasets        []*datasetCompaction
+	meta            *metastorev1.BlockMeta
+	strings         *metadata.StringTable
+	metricsExporter *metrics.Exporter
+	workerId        uuid.UUID
 }
 
 func newBlockCompaction(
 	id string,
 	tenant string,
 	shard uint32,
 	compactionLevel uint32,
+	workerId uuid.UUID,
 ) *CompactionPlan {
 	p := &CompactionPlan{
 		tenant:     tenant,
 		datasetMap: make(map[int32]*datasetCompaction),
 		strings:    metadata.NewStringTable(),
+		workerId:   workerId,
 	}
 	p.path = BuildObjectPath(tenant, shard, compactionLevel, id)
 	p.meta = &metastorev1.BlockMeta{
@@ -176,6 +194,9 @@ func newBlockCompaction(
 		Shard:           shard,
 		CompactionLevel: compactionLevel,
 	}
+	if compactionLevel == 1 {
+		p.metricsExporter = metrics.NewExporter(tenant)
+	}
 	return p
 }
 
@@ -189,6 +210,9 @@ func (b *CompactionPlan) Compact(ctx context.Context, dst objstore.Bucket, tmpdi
 		if err = s.compact(ctx, w); err != nil {
 			return nil, fmt.Errorf("compacting block: %w", err)
 		}
+		if b.metricsExporter != nil {
+			b.metricsExporter.AppendMetrics(s.metricsRecorder.Recordings)
+		}
 		b.meta.Datasets = append(b.meta.Datasets, s.meta)
 	}
 	if err = w.Flush(ctx); err != nil {
@@ -217,6 +241,10 @@ func (b *CompactionPlan) addDataset(md *metastorev1.BlockMeta, s *metastorev1.Da
 	return sm
 }
 
+func (c *CompactionPlan) SendRecordedMetrics() error {
+	return c.metricsExporter.Send()
+}
+
 type datasetCompaction struct {
 	// Dataset name.
 	name   string
@@ -236,6 +264,9 @@ type datasetCompaction struct {
 	profiles uint64
 
 	flushOnce sync.Once
+
+	workerId        uuid.UUID
+	metricsRecorder *metrics.Recorder
 }
 
 func (b *CompactionPlan) newDatasetCompaction(tenant, name int32) *datasetCompaction {
@@ -254,6 +285,7 @@ func (b *CompactionPlan) newDatasetCompaction(tenant, name int32) *datasetCompac
 			Size:            0,
 			Labels:          nil,
 		},
+		workerId: b.workerId,
 	}
 }
 
@@ -309,6 +341,13 @@ func (m *datasetCompaction) open(ctx context.Context, path string) (err error) {
 	m.indexRewriter = newIndexRewriter(m.path)
 	m.symbolsRewriter = newSymbolsRewriter(m.path)
 
+	if m.parent.meta.CompactionLevel == 1 {
+		recordingTime := int64(ulid.MustParse(m.parent.meta.Id).Time())
+		rules := metrics.RecordingRulesFromTenant(m.parent.tenant)
+		pyroscopeInstance := pyroscopeInstanceHash(m.parent.meta.Shard, m.workerId)
+		m.metricsRecorder = metrics.NewRecorder(rules, recordingTime, pyroscopeInstance)
+	}
+
 	g, ctx := errgroup.WithContext(ctx)
 	for _, s := range m.datasets {
 		s := s
@@ -330,6 +369,13 @@ func (m *datasetCompaction) open(ctx context.Context, path string) (err error) {
 	return nil
 }
 
+func pyroscopeInstanceHash(shard uint32, id uuid.UUID) string {
+	buf := make([]byte, 0, 40)
+	buf = append(buf, byte(shard>>24), byte(shard>>16), byte(shard>>8), byte(shard))
+	buf = append(buf, id.String()...)
+	return fmt.Sprintf("%x", xxhash.Sum64(buf))
+}
+
 func (m *datasetCompaction) mergeAndClose(ctx context.Context) (err error) {
 	defer func() {
 		err = multierror.New(err, m.close()).Err()
@@ -366,6 +412,9 @@ func (m *datasetCompaction) writeRow(r ProfileEntry) (err error) {
 	if err = m.symbolsRewriter.rewriteRow(r); err != nil {
 		return err
 	}
+	if m.metricsRecorder != nil {
+		m.metricsRecorder.RecordRow(r.Fingerprint, r.Labels, r.Row.TotalValue())
+	}
 	return m.profilesWriter.writeRow(r)
 }
 

pkg/experiment/block/compaction_test.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/encoding/protojson"
@@ -26,7 +27,7 @@ func Test_CompactBlocks(t *testing.T) {
 	require.NoError(t, err)
 
 	dst, tempdir := testutil.NewFilesystemBucket(t, ctx, t.TempDir())
-	compactedBlocks, err := Compact(ctx, resp.Blocks, bucket,
+	compactedBlocks, err := Compact(ctx, resp.Blocks, bucket, uuid.New(),
 		WithCompactionDestination(dst),
 		WithCompactionTempDir(tempdir),
 		WithCompactionObjectOptions(

pkg/experiment/compactor/compaction_worker.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	"github.com/google/uuid"
 	"github.com/grafana/dskit/services"
 	"github.com/opentracing/opentracing-go"
 	"github.com/pkg/errors"
@@ -46,6 +47,8 @@ type Worker struct {
 	stopped   atomic.Bool
 	closeOnce sync.Once
 	wg        sync.WaitGroup
+
+	id uuid.UUID
 }
 
 type Config struct {
@@ -96,6 +99,7 @@ func New(
 		client:  client,
 		storage: storage,
 		metrics: newMetrics(reg),
+		id:      uuid.New(),
 	}
 	w.threads = config.JobConcurrency
 	if w.threads < 1 {
@@ -363,7 +367,7 @@ func (w *Worker) runCompaction(job *compactionJob) {
 
 	tempdir := filepath.Join(w.config.TempDir, job.Name)
 	sourcedir := filepath.Join(tempdir, "source")
-	compacted, err := block.Compact(ctx, job.blocks, w.storage,
+	compacted, err := block.Compact(ctx, job.blocks, w.storage, w.id,
 		block.WithCompactionTempDir(tempdir),
 		block.WithCompactionObjectOptions(
 			block.WithObjectMaxSizeLoadInMemory(w.config.SmallObjectSize),

pkg/experiment/metrics/exporter.go

@@ -0,0 +1,108 @@
+package metrics
+
+import (
+	"context"
+	"net/url"
+	"time"
+
+	"github.com/gogo/protobuf/proto"
+	"github.com/klauspost/compress/snappy"
+	"github.com/prometheus/common/config"
+	"github.com/prometheus/common/model"
+	"github.com/prometheus/prometheus/prompb"
+	"github.com/prometheus/prometheus/storage/remote"
+)
+
+type Exporter struct {
+	config Config
+	client remote.WriteClient
+	data   map[AggregatedFingerprint]*TimeSeries
+}
+
+type Config struct {
+	url      string
+	username string
+	password config.Secret
+}
+
+func NewExporter(tenant string) *Exporter {
+	cfg := configFromTenant(tenant)
+	return &Exporter{
+		config: cfg,
+		data:   map[AggregatedFingerprint]*TimeSeries{},
+	}
+}
+
+func (e *Exporter) AppendMetrics(recordings []*Recording) {
+	for _, r := range recordings {
+		for fp, ts := range r.data {
+			e.data[fp] = ts
+		}
+	}
+}
+
+func (e *Exporter) Send() error {
+	if e.client == nil {
+		e.client = newClient(e.config)
+	}
+
+	p := &prompb.WriteRequest{Timeseries: make([]prompb.TimeSeries, 0, len(e.data))}
+	for _, ts := range e.data {
+		pts := prompb.TimeSeries{
+			Labels: make([]prompb.Label, 0, len(ts.Labels)),
+		}
+		for _, l := range ts.Labels {
+			pts.Labels = append(pts.Labels, prompb.Label{
+				Name:  l.Name,
+				Value: l.Value,
+			})
+		}
+		for _, s := range ts.Samples {
+			pts.Samples = append(pts.Samples, prompb.Sample{
+				Value:     s.Value,
+				Timestamp: s.Timestamp,
+			})
+		}
+		p.Timeseries = append(p.Timeseries, pts)
+	}
+	buf := proto.NewBuffer(nil)
+	if err := buf.Marshal(p); err != nil {
+		return err
+	}
+	return e.client.Store(context.Background(), snappy.Encode(nil, buf.Bytes()), 0)
+}
+
+func newClient(cfg Config) remote.WriteClient {
+	wURL, err := url.Parse(cfg.url)
+	if err != nil {
+		panic(err)
+	}
+
+	c, err := remote.NewWriteClient("exporter", &remote.ClientConfig{
+		URL:     &config.URL{URL: wURL},
+		Timeout: model.Duration(time.Second * 10),
+		HTTPClientConfig: config.HTTPClientConfig{
+			BasicAuth: &config.BasicAuth{
+				Username: cfg.username,
+				Password: cfg.password,
+			},
+		},
+		SigV4Config:      nil,
+		AzureADConfig:    nil,
+		Headers:          nil,
+		RetryOnRateLimit: false,
+	})
+	if err != nil {
+		panic(err)
+	}
+	return c
+}
+
+func configFromTenant(tenant string) Config {
+	// TODO
+	return Config{
+		url:      "omitted",
+		username: "omitted",
+		password: "omitted",
+	}
+}