diff --git a/pkg/featureflags/client_capability.go b/pkg/featureflags/client_capability.go index f8f5c1c7f0..66e47698ef 100644 --- a/pkg/featureflags/client_capability.go +++ b/pkg/featureflags/client_capability.go @@ -9,16 +9,29 @@ import ( "connectrpc.com/connect" "github.com/go-kit/log/level" "github.com/grafana/dskit/middleware" + "github.com/grafana/pyroscope/pkg/tenant" "github.com/grafana/pyroscope/pkg/util" + "github.com/prometheus/client_golang/prometheus" "google.golang.org/grpc" "google.golang.org/grpc/metadata" ) const ( - // Capability names - update parseClientCapabilities below when new capabilities added + // Capability names - update parseClientCapabilities and trackEnabledCapabilities + // below when new capabilities added allowUtf8LabelNamesCapabilityName string = "allow-utf8-labelnames" ) +var ( + usage = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Name: "client_capability_enabled_total", + Help: "Total number of requests with client capabilities enabled", + }, + []string{"tenant", "capability_name"}, + ) +) + // Define a custom context key type to avoid collisions type contextKey struct{} @@ -35,7 +48,35 @@ func GetClientCapabilities(ctx context.Context) (ClientCapabilities, bool) { return value, ok } -func ClientCapabilitiesGRPCMiddleware() grpc.UnaryServerInterceptor { +type ClientCapabilityMiddleware struct { + usage *prometheus.CounterVec +} + +func NewClientCapabilityMiddleware(reg prometheus.Registerer) *ClientCapabilityMiddleware { + usage = util.RegisterOrGet(reg, usage) + + return &ClientCapabilityMiddleware{ + usage: usage, + } +} + +// trackEnabledCapabilities records metrics for each enabled capability +func (m *ClientCapabilityMiddleware) trackEnabledCapabilities(ctx context.Context, capabilities ClientCapabilities) { + tenantID, err := tenant.ExtractTenantIDFromContext(ctx) + if err != nil { + // Fall back if tenant cannot be extracted + tenantID = tenant.DefaultTenantID + } + + // Track each enabled capability + if capabilities.AllowUtf8LabelNames { + m.usage.WithLabelValues(tenantID, allowUtf8LabelNamesCapabilityName).Inc() + } +} + +// CreateGRPC creates gRPC middleware that extracts and parses the +// `Accept` header for capabilities the client supports +func (m *ClientCapabilityMiddleware) CreateGRPC() grpc.UnaryServerInterceptor { return func( ctx context.Context, req interface{}, @@ -62,13 +103,17 @@ func ClientCapabilitiesGRPCMiddleware() grpc.UnaryServerInterceptor { } enhancedCtx := WithClientCapabilities(ctx, clientCapabilities) + + // Track enabled capabilities for metrics + m.trackEnabledCapabilities(enhancedCtx, clientCapabilities) + return handler(enhancedCtx, req) } } -// ClientCapabilitiesHttpMiddleware creates middleware that extracts and parses the +// CreateHttp creates HTTP middleware that extracts and parses the // `Accept` header for capabilities the client supports -func ClientCapabilitiesHttpMiddleware() middleware.Interface { +func (m *ClientCapabilityMiddleware) CreateHttp() middleware.Interface { return middleware.Func(func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { clientCapabilities, err := parseClientCapabilities(r.Header) @@ -78,6 +123,10 @@ func ClientCapabilitiesHttpMiddleware() middleware.Interface { } ctx := WithClientCapabilities(r.Context(), clientCapabilities) + + // Track enabled capabilities for metrics + m.trackEnabledCapabilities(ctx, clientCapabilities) + next.ServeHTTP(w, r.WithContext(ctx)) }) }) diff --git a/pkg/pyroscope/modules.go b/pkg/pyroscope/modules.go index 0a2db5e030..51c219531c 100644 --- a/pkg/pyroscope/modules.go +++ b/pkg/pyroscope/modules.go @@ -465,9 +465,11 @@ func (f *Pyroscope) initServer() (services.Service, error) { // see https://github.com/grafana/pyroscope/issues/231 f.Cfg.Server.DoNotAddDefaultHTTPMiddleware = true f.Cfg.Server.ExcludeRequestInLog = true // gRPC-specific. + + clientCapabilitiesMiddleware := featureflags.NewClientCapabilityMiddleware(f.reg) f.Cfg.Server.GRPCMiddleware = append(f.Cfg.Server.GRPCMiddleware, util.RecoveryInterceptorGRPC, - featureflags.ClientCapabilitiesGRPCMiddleware(), + clientCapabilitiesMiddleware.CreateGRPC(), ) if f.Cfg.V2 { @@ -526,7 +528,7 @@ func (f *Pyroscope) initServer() (services.Service, error) { httpMetric, objstoreTracerMiddleware, httputil.K6Middleware(), - featureflags.ClientCapabilitiesHttpMiddleware(), + clientCapabilitiesMiddleware.CreateHttp(), } if f.Cfg.SelfProfiling.UseK6Middleware { defaultHTTPMiddleware = append(defaultHTTPMiddleware, httputil.K6Middleware())