Add metrics for block querying

This refactors the registrt and logger passing to use a context.

Author: simonswine

cmd/firetool/blocks.go

@@ -32,7 +32,7 @@ func blocksList(ctx context.Context) error {
 		return err
 	}
 
-	metas, err := firedb.NewBlockQuerier(logger, bucket).BlockMetas(ctx)
+	metas, err := firedb.NewBlockQuerier(ctx, bucket).BlockMetas(ctx)
 	if err != nil {
 		return err
 	}

cmd/firetool/main.go

@@ -10,6 +10,8 @@ import (
 	"github.com/go-kit/log/level"
 	"github.com/prometheus/common/version"
 	"gopkg.in/alecthomas/kingpin.v2"
+
+	firecontext "github.com/grafana/fire/pkg/fire/context"
 )
 
 var cfg struct {
@@ -26,7 +28,8 @@ var (
 )
 
 func main() {
-	ctx := context.Background()
+	ctx := firecontext.WithLogger(context.Background(), logger)
+
 	app := kingpin.New(filepath.Base(os.Args[0]), "Tooling for Grafana Fire, the continuous profiling aggregation system.").UsageWriter(os.Stdout)
 	app.Version(version.Print("firetool"))
 	app.HelpFlag.Short('h')

pkg/fire/context/context.go

@@ -0,0 +1,54 @@
+package context
+
+import (
+	"context"
+
+	"github.com/go-kit/log"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+type contextKey int
+
+const (
+	loggerKey contextKey = iota
+	registryKey
+)
+
+func WithLogger(ctx context.Context, logger log.Logger) context.Context {
+	return context.WithValue(ctx, loggerKey, logger)
+}
+
+func Logger(ctx context.Context) log.Logger {
+	if logger, ok := ctx.Value(loggerKey).(log.Logger); ok {
+		return logger
+	}
+	return log.NewNopLogger()
+}
+
+func WithRegistry(ctx context.Context, registry prometheus.Registerer) context.Context {
+	return context.WithValue(ctx, registryKey, registry)
+}
+
+func Registry(ctx context.Context) prometheus.Registerer {
+	if registry, ok := ctx.Value(registryKey).(prometheus.Registerer); ok {
+		return registry
+	}
+	return prometheus.NewRegistry()
+}
+
+func WrapTenant(ctx context.Context, tenantID string) context.Context {
+	// wrap registry
+	reg := Registry(ctx)
+	ctx = WithRegistry(ctx, prometheus.WrapRegistererWith(
+		prometheus.Labels{"tenant": tenantID},
+		reg,
+	))
+
+	// add field to logger
+	logger := Logger(ctx)
+	ctx = WithLogger(ctx,
+		log.With(logger, "tenant", tenantID),
+	)
+
+	return ctx
+}

pkg/fire/modules.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/grafana/fire/pkg/agent"
 	"github.com/grafana/fire/pkg/distributor"
+	firecontext "github.com/grafana/fire/pkg/fire/context"
 	agentv1 "github.com/grafana/fire/pkg/gen/agent/v1"
 	"github.com/grafana/fire/pkg/gen/agent/v1/agentv1connect"
 	"github.com/grafana/fire/pkg/gen/ingester/v1/ingesterv1connect"
@@ -171,7 +172,11 @@ func (f *Fire) initStorage() (_ services.Service, err error) {
 func (f *Fire) initIngester() (_ services.Service, err error) {
 	f.Cfg.Ingester.LifecyclerConfig.ListenPort = f.Cfg.Server.HTTPListenPort
 
-	ingester, err := ingester.New(f.Cfg.Ingester, f.Cfg.FireDB, f.logger, f.reg, f.storageBucket)
+	// TODO: This should be passed to all other services and could also be used to signal shutdown
+	firectx := firecontext.WithLogger(context.Background(), f.logger)
+	firectx = firecontext.WithRegistry(firectx, f.reg)
+
+	ingester, err := ingester.New(firectx, f.Cfg.Ingester, f.Cfg.FireDB, f.storageBucket)
 	if err != nil {
 		return nil, err
 	}

pkg/firedb/block_querier.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"sort"
 	"sync"
+	"time"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
@@ -26,6 +27,7 @@ import (
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc/codes"
 
+	firecontext "github.com/grafana/fire/pkg/fire/context"
 	"github.com/grafana/fire/pkg/firedb/block"
 	query "github.com/grafana/fire/pkg/firedb/query"
 	schemav1 "github.com/grafana/fire/pkg/firedb/schemas/v1"
@@ -44,20 +46,22 @@ type tableReader interface {
 }
 
 type BlockQuerier struct {
-	logger log.Logger
+	firectx context.Context
+	logger  log.Logger
 
 	bucketReader fireobjstore.BucketReader
 
 	queriers     []*singleBlockQuerier
 	queriersLock sync.RWMutex
 }
 
-func NewBlockQuerier(logger log.Logger, bucketReader fireobjstore.BucketReader) *BlockQuerier {
-	if logger == nil {
-		logger = log.NewNopLogger()
-	}
+func NewBlockQuerier(firectx context.Context, bucketReader fireobjstore.BucketReader) *BlockQuerier {
 	return &BlockQuerier{
-		logger:       logger,
+		firectx: contextWithBlockMetrics(firectx,
+			newBlocksMetrics(
+				firecontext.Registry(firectx),
+			),
+		),
 		bucketReader: bucketReader,
 	}
 }
@@ -67,7 +71,7 @@ func (b *BlockQuerier) reconstructMetaFromBlock(ctx context.Context, ulid ulid.U
 	fakeMeta := block.NewMeta()
 	fakeMeta.ULID = ulid
 
-	q := newSingleBlockQuerierFromMeta(b.logger, b.bucketReader, fakeMeta)
+	q := newSingleBlockQuerierFromMeta(b.firectx, b.bucketReader, fakeMeta)
 	defer q.Close()
 
 	meta, err := q.reconstructMeta(ctx)
@@ -188,7 +192,7 @@ func (b *BlockQuerier) Sync(ctx context.Context) error {
 			continue
 		}
 
-		b.queriers[pos] = newSingleBlockQuerierFromMeta(b.logger, b.bucketReader, m)
+		b.queriers[pos] = newSingleBlockQuerierFromMeta(b.firectx, b.bucketReader, m)
 	}
 	// ensure queriers are in ascending order.
 	sort.Slice(b.queriers, func(i, j int) bool {
@@ -265,7 +269,9 @@ func (mm *minMax) InRange(start, end model.Time) bool {
 }
 
 type singleBlockQuerier struct {
-	logger       log.Logger
+	logger  log.Logger
+	metrics *blocksMetrics
+
 	bucketReader fireobjstore.BucketReader
 	meta         *block.Meta
 
@@ -282,9 +288,11 @@ type singleBlockQuerier struct {
 	profiles              parquetReader[*schemav1.Profile, *schemav1.ProfilePersister]
 }
 
-func newSingleBlockQuerierFromMeta(logger log.Logger, bucketReader fireobjstore.BucketReader, meta *block.Meta) *singleBlockQuerier {
+func newSingleBlockQuerierFromMeta(firectx context.Context, bucketReader fireobjstore.BucketReader, meta *block.Meta) *singleBlockQuerier {
 	q := &singleBlockQuerier{
-		logger:       logger,
+		logger:  firecontext.Logger(firectx),
+		metrics: contextBlockMetrics(firectx),
+
 		bucketReader: fireobjstore.BucketReaderWithPrefix(bucketReader, meta.ULID.String()),
 		meta:         meta,
 	}
@@ -727,8 +735,10 @@ func newByteSliceFromBucketReader(bucketReader fireobjstore.BucketReader, path s
 }
 
 func (q *singleBlockQuerier) open(ctx context.Context) error {
+	start := time.Now()
 	sp, ctx := opentracing.StartSpanFromContext(ctx, "BlockQuerier - open")
 	defer func() {
+		q.metrics.blockOpeningLatency.Observe(time.Since(start).Seconds())
 		sp.LogFields(
 			otlog.String("block_ulid", q.meta.ULID.String()),
 		)
@@ -755,7 +765,7 @@ func (q *singleBlockQuerier) open(ctx context.Context) error {
 
 	// open parquet files
 	for _, x := range q.tables {
-		if err := x.open(ctx, q.bucketReader); err != nil {
+		if err := x.open(contextWithBlockMetrics(ctx, q.metrics), q.bucketReader); err != nil {
 			return err
 		}
 	}
@@ -767,9 +777,11 @@ type parquetReader[M Models, P schemav1.PersisterName] struct {
 	persister P
 	file      *parquet.File
 	reader    fireobjstore.ReaderAt
+	metrics   *blocksMetrics
 }
 
 func (r *parquetReader[M, P]) open(ctx context.Context, bucketReader fireobjstore.BucketReader) error {
+	r.metrics = contextBlockMetrics(ctx)
 	filePath := r.persister.Name() + block.ParquetSuffix
 
 	ra, err := bucketReader.ReaderAt(ctx, filePath)
@@ -823,6 +835,7 @@ func (r *parquetReader[M, P]) columnIter(ctx context.Context, columnName string,
 	if index == -1 {
 		return query.NewErrIterator(fmt.Errorf("column '%s' not found in parquet file '%s'", columnName, r.relPath()))
 	}
+	ctx = query.AddMetricsToContext(ctx, r.metrics.query)
 	return query.NewColumnIterator(ctx, r.file.RowGroups(), index, columnName, 1000, predicate, alias)
 }
 

pkg/firedb/firedb.go

@@ -25,6 +25,7 @@ import (
 	"github.com/samber/lo"
 	"golang.org/x/sync/errgroup"
 
+	firecontext "github.com/grafana/fire/pkg/fire/context"
 	"github.com/grafana/fire/pkg/firedb/block"
 	commonv1 "github.com/grafana/fire/pkg/gen/common/v1"
 	ingestv1 "github.com/grafana/fire/pkg/gen/ingester/v1"
@@ -55,29 +56,32 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 type FireDB struct {
 	services.Service
 
+	logger  log.Logger
+	firectx context.Context
+
 	cfg    Config
-	reg    prometheus.Registerer
-	logger log.Logger
 	stopCh chan struct{}
 	wg     sync.WaitGroup
 
-	headLock    sync.RWMutex
-	head        *Head
-	headMetrics *headMetrics
+	headLock sync.RWMutex
+	head     *Head
 
 	headFlushTimer time.Timer
 
 	blockQuerier *BlockQuerier
 }
 
-func New(cfg Config, logger log.Logger, reg prometheus.Registerer) (*FireDB, error) {
-	headMetrics := newHeadMetrics(reg)
+func New(firectx context.Context, cfg Config) (*FireDB, error) {
+	reg := firecontext.Registry(firectx)
+
+	// ensure head metrics are registered early so they are reused for the new head
+	firectx = contextWithHeadMetrics(firectx, newHeadMetrics(reg))
+
 	f := &FireDB{
-		cfg:         cfg,
-		reg:         reg,
-		logger:      logger,
-		stopCh:      make(chan struct{}, 0),
-		headMetrics: headMetrics,
+		cfg:     cfg,
+		logger:  firecontext.Logger(firectx),
+		firectx: firectx,
+		stopCh:  make(chan struct{}, 0),
 	}
 	if _, err := f.initHead(); err != nil {
 		return nil, err
@@ -97,7 +101,7 @@ func New(cfg Config, logger log.Logger, reg prometheus.Registerer) (*FireDB, err
 		return nil, fmt.Errorf("mkdir %s: %w", f.LocalDataPath(), err)
 	}
 
-	f.blockQuerier = NewBlockQuerier(logger, bucketReader)
+	f.blockQuerier = NewBlockQuerier(firectx, bucketReader)
 
 	// do an initial querier sync
 	ctx := context.Background()
@@ -460,7 +464,7 @@ func (f *FireDB) initHead() (oldHead *Head, err error) {
 	f.headLock.Lock()
 	defer f.headLock.Unlock()
 	oldHead = f.head
-	f.head, err = NewHead(f.cfg, headWithMetrics(f.headMetrics), HeadWithLogger(f.logger))
+	f.head, err = NewHead(f.firectx, f.cfg)
 	if err != nil {
 		return oldHead, err
 	}

pkg/firedb/firedb_test.go

@@ -10,7 +10,6 @@ import (
 	"time"
 
 	"github.com/bufbuild/connect-go"
-	"github.com/go-kit/log"
 	"github.com/google/uuid"
 	"github.com/pkg/errors"
 	"github.com/prometheus/common/model"
@@ -33,16 +32,16 @@ func TestCreateLocalDir(t *testing.T) {
 	dataPath := t.TempDir()
 	localFile := dataPath + "/local"
 	require.NoError(t, ioutil.WriteFile(localFile, []byte("d"), 0o644))
-	_, err := New(Config{
+	_, err := New(context.Background(), Config{
 		DataPath:         dataPath,
 		MaxBlockDuration: 30 * time.Minute,
-	}, log.NewNopLogger(), nil)
+	})
 	require.Error(t, err)
 	require.NoError(t, os.Remove(localFile))
-	_, err = New(Config{
+	_, err = New(context.Background(), Config{
 		DataPath:         dataPath,
 		MaxBlockDuration: 30 * time.Minute,
-	}, log.NewNopLogger(), nil)
+	})
 	require.NoError(t, err)
 }
 
@@ -133,10 +132,10 @@ func TestMergeProfilesStacktraces(t *testing.T) {
 		step    = 15 * time.Second
 	)
 
-	db, err := New(Config{
+	db, err := New(context.Background(), Config{
 		DataPath:         testDir,
 		MaxBlockDuration: time.Duration(100000) * time.Minute, // we will manually flush
-	}, log.NewNopLogger(), nil)
+	})
 	require.NoError(t, err)
 	defer require.NoError(t, db.Close())