Refactor metrics

Author: aleks-p

pkg/metastore/fsm/fsm.go

@@ -275,8 +275,6 @@ func (fsm *FSM) applyCommand(cmd *raft.Log) any {
 		return errResponse(cmd, fmt.Errorf("unknown command type: %d", e.Type))
 	}
 
-	fsm.metrics.contextRegistrySize.Set(float64(fsm.contextRegistry.Size()))
-
 	// Apply is never called concurrently with Restore, so we don't need
 	// to lock the FSM: db.boltdb is guaranteed to be in a consistent state.
 	rawTx, err := fsm.db.boltdb.Begin(true)

pkg/metastore/fsm/metrics.go

@@ -15,7 +15,6 @@ type metrics struct {
 	fsmRestoreSnapshotDuration    prometheus.Histogram
 	fsmApplyCommandSize           *prometheus.HistogramVec
 	fsmApplyCommandDuration       *prometheus.HistogramVec
-	contextRegistrySize           prometheus.Gauge
 }
 
 func newMetrics(reg prometheus.Registerer) *metrics {
@@ -64,11 +63,6 @@ func newMetrics(reg prometheus.Registerer) *metrics {
 			NativeHistogramMaxBucketNumber:  50,
 			NativeHistogramMinResetDuration: time.Hour,
 		}, []string{"command"}),
-
-		contextRegistrySize: prometheus.NewGauge(prometheus.GaugeOpts{
-			Name: "fsm_context_registry_size",
-			Help: "Number of contexts currently stored in the registry for tracing propagation",
-		}),
 	}
 	if reg != nil {
 		util.RegisterOrGet(reg, m.boltDBPersistSnapshotSize)
@@ -77,7 +71,6 @@ func newMetrics(reg prometheus.Registerer) *metrics {
 		util.RegisterOrGet(reg, m.fsmRestoreSnapshotDuration)
 		util.RegisterOrGet(reg, m.fsmApplyCommandSize)
 		util.RegisterOrGet(reg, m.fsmApplyCommandDuration)
-		util.RegisterOrGet(reg, m.contextRegistrySize)
 	}
 	return m
 }

pkg/metastore/metastore.go

@@ -119,19 +119,18 @@ func New(
 	placementMgr *placement.Manager,
 ) (*Metastore, error) {
 	m := &Metastore{
-		config:         config,
-		overrides:      overrides,
-		logger:         logger,
-		reg:            reg,
-		health:         healthService,
-		bucket:         bucket,
-		placement:      placementMgr,
-		raftNodeClient: client,
+		config:          config,
+		overrides:       overrides,
+		logger:          logger,
+		reg:             reg,
+		health:          healthService,
+		bucket:          bucket,
+		placement:       placementMgr,
+		raftNodeClient:  client,
+		contextRegistry: tracing.NewContextRegistry(reg),
 	}
 
 	var err error
-	// Create the context registry that will be shared between FSM and Node.
-	m.contextRegistry = tracing.NewContextRegistry()
 	if m.fsm, err = fsm.New(m.logger, m.reg, m.config.FSM, m.contextRegistry); err != nil {
 		return nil, fmt.Errorf("failed to initialize store: %w", err)
 	}

pkg/metastore/raftnode/node.go

@@ -313,7 +313,7 @@ func (n *Node) TransferLeadership() (err error) {
 // Propose makes an attempt to apply the given command to the FSM.
 // The function returns an error if node is not the leader.
 func (n *Node) Propose(ctx context.Context, t fsm.RaftLogEntryType, m proto.Message) (resp proto.Message, err error) {
-	span, ctx := opentracing.StartSpanFromContext(ctx, "FSM.Propose")
+	span, ctx := opentracing.StartSpanFromContext(ctx, "node.Propose")
 	defer func() {
 		if err != nil {
 			ext.LogError(span, err)

pkg/metastore/tracing/context_registry.go

@@ -4,6 +4,10 @@ import (
 	"context"
 	"sync"
 	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+
+	"github.com/grafana/pyroscope/pkg/util"
 )
 
 // ContextRegistry maintains a mapping of IDs to contexts for tracing purposes.
@@ -21,6 +25,9 @@ type ContextRegistry struct {
 	entryTTL time.Duration
 	stop     chan struct{}
 	done     chan struct{}
+
+	// sizeMetric tracks the number of entries in the registry
+	sizeMetric prometheus.Gauge
 }
 
 type contextEntry struct {
@@ -34,25 +41,34 @@ const (
 )
 
 // NewContextRegistry creates a new context registry with background cleanup.
-func NewContextRegistry() *ContextRegistry {
-	return newContextRegistry(defaultCleanupInterval, defaultEntryTTL)
+func NewContextRegistry(reg prometheus.Registerer) *ContextRegistry {
+	return newContextRegistry(defaultCleanupInterval, defaultEntryTTL, reg)
 }
 
-// NewContextRegistry creates a new context registry with background cleanup.
-func newContextRegistry(cleanupInterval, entryTTL time.Duration) *ContextRegistry {
+// newContextRegistry creates a new context registry with background cleanup.
+func newContextRegistry(cleanupInterval, entryTTL time.Duration, reg prometheus.Registerer) *ContextRegistry {
 	if cleanupInterval <= 0 {
 		cleanupInterval = defaultCleanupInterval
 	}
 	if entryTTL <= 0 {
 		entryTTL = defaultEntryTTL
 	}
 
+	sizeMetric := prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "context_registry_size",
+		Help: "Number of contexts currently stored in the registry for tracing propagation",
+	})
+	if reg != nil {
+		util.RegisterOrGet(reg, sizeMetric)
+	}
+
 	r := &ContextRegistry{
 		entries:         make(map[string]*contextEntry),
 		cleanupInterval: cleanupInterval,
 		entryTTL:        entryTTL,
 		stop:            make(chan struct{}),
 		done:            make(chan struct{}),
+		sizeMetric:      sizeMetric,
 	}
 
 	go r.cleanupLoop()
@@ -113,6 +129,10 @@ func (r *ContextRegistry) cleanup() {
 			delete(r.entries, id)
 		}
 	}
+
+	if r.sizeMetric != nil {
+		r.sizeMetric.Set(float64(len(r.entries)))
+	}
 }
 
 // Shutdown stops the cleanup loop.

pkg/metastore/tracing/context_registry_test.go

@@ -13,7 +13,7 @@ import (
 type contextKey string
 
 func TestContextRegistry_StoreAndRetrieve(t *testing.T) {
-	r := newContextRegistry(1*time.Second, 5*time.Second)
+	r := newContextRegistry(1*time.Second, 5*time.Second, nil)
 	defer r.Shutdown()
 
 	ctx := context.WithValue(context.Background(), contextKey("key"), "value")
@@ -25,7 +25,7 @@ func TestContextRegistry_StoreAndRetrieve(t *testing.T) {
 }
 
 func TestContextRegistry_RetrieveNotFound(t *testing.T) {
-	r := newContextRegistry(1*time.Second, 5*time.Second)
+	r := newContextRegistry(1*time.Second, 5*time.Second, nil)
 	defer r.Shutdown()
 
 	retrieved, found := r.Retrieve("nonexistent-id")
@@ -34,7 +34,7 @@ func TestContextRegistry_RetrieveNotFound(t *testing.T) {
 }
 
 func TestContextRegistry_Delete(t *testing.T) {
-	r := newContextRegistry(1*time.Second, 5*time.Second)
+	r := newContextRegistry(1*time.Second, 5*time.Second, nil)
 	defer r.Shutdown()
 
 	ctx := context.WithValue(context.Background(), contextKey("key"), "value")
@@ -51,7 +51,7 @@ func TestContextRegistry_Delete(t *testing.T) {
 
 func TestContextRegistry_Cleanup(t *testing.T) {
 	// Use short TTL for a faster test
-	r := newContextRegistry(100*time.Millisecond, 200*time.Millisecond)
+	r := newContextRegistry(100*time.Millisecond, 200*time.Millisecond, nil)
 	defer r.Shutdown()
 
 	ctx := context.WithValue(context.Background(), contextKey("key"), "value")
@@ -69,7 +69,7 @@ func TestContextRegistry_Cleanup(t *testing.T) {
 }
 
 func TestContextRegistry_Size(t *testing.T) {
-	r := newContextRegistry(1*time.Second, 5*time.Second)
+	r := newContextRegistry(1*time.Second, 5*time.Second, nil)
 	defer r.Shutdown()
 
 	assert.Equal(t, 0, r.Size())
@@ -86,7 +86,7 @@ func TestContextRegistry_Size(t *testing.T) {
 }
 
 func TestContextRegistry_ConcurrentAccess(t *testing.T) {
-	r := newContextRegistry(1*time.Second, 5*time.Second)
+	r := newContextRegistry(1*time.Second, 5*time.Second, nil)
 	defer r.Shutdown()
 
 	done := make(chan bool)