feat: k6 middleware

cmd/pyroscope/help-all.txt.tmpl

@@ -753,6 +753,8 @@ Usage of ./pyroscope:
     	When running in single binary (--target=all) Pyroscope will push (Go SDK) profiles to itself. Set to true to disable self-profiling.
   -self-profiling.mutex-profile-fraction int
     	 (default 5)
+  -self-profiling.use-k6-middleware
+    	Read k6 labels from request headers and set them as dynamic profile tags.
   -server.graceful-shutdown-timeout duration
     	Timeout for graceful shutdowns (default 30s)
   -server.grpc-conn-limit int

cmd/pyroscope/help.txt.tmpl

@@ -195,6 +195,8 @@ Usage of ./pyroscope:
     	When running in single binary (--target=all) Pyroscope will push (Go SDK) profiles to itself. Set to true to disable self-profiling.
   -self-profiling.mutex-profile-fraction int
     	 (default 5)
+  -self-profiling.use-k6-middleware
+    	Read k6 labels from request headers and set them as dynamic profile tags.
   -server.graceful-shutdown-timeout duration
     	Timeout for graceful shutdowns (default 30s)
   -server.grpc-conn-limit int

docs/sources/configure-server/reference-configuration-parameters/index.md

@@ -255,6 +255,10 @@ self_profiling:
   # CLI flag: -self-profiling.block-profile-rate
   [block_profile_rate: <int> | default = 5]
 
+  # Read k6 labels from request headers and set them as dynamic profile tags.
+  # CLI flag: -self-profiling.use-k6-middleware
+  [use_k6_middleware: <boolean> | default = false]
+
 # When set to true, incoming HTTP requests must specify tenant ID in HTTP
 # X-Scope-OrgId header. When set to false, tenant ID anonymous is used instead.
 # CLI flag: -auth.multitenancy-enabled

go.mod

@@ -31,8 +31,9 @@ require (
 	github.com/grafana/alloy/syntax v0.1.0
 	github.com/grafana/dskit v0.0.0-20231221015914-de83901bf4d6
 	github.com/grafana/jfr-parser/pprof v0.0.0-20240228024232-8abcb81c304c
-	github.com/grafana/pyroscope-go v1.0.3
+	github.com/grafana/pyroscope-go v1.1.1
 	github.com/grafana/pyroscope-go/godeltaprof v0.1.8
+	github.com/grafana/pyroscope-go/x/k6 v0.0.0-20240920080526-46ebf824bf43
 	github.com/grafana/pyroscope/api v0.4.0
 	github.com/grafana/regexp v0.0.0-20221123153739-15dc172cd2db
 	github.com/grpc-ecosystem/go-grpc-middleware v1.4.0

go.sum

@@ -393,10 +393,12 @@ github.com/grafana/jfr-parser/pprof v0.0.0-20240228024232-8abcb81c304c h1:tGu1DT
 github.com/grafana/jfr-parser/pprof v0.0.0-20240228024232-8abcb81c304c/go.mod h1:P5406BrWxjahTzVF6aCSumNI1KPlZJc0zO0v+zKZ4gc=
 github.com/grafana/memberlist v0.3.1-0.20220708130638-bd88e10a3d91 h1:/NipyHnOmvRsVzj81j2qE0VxsvsqhOB0f4vJIhk2qCQ=
 github.com/grafana/memberlist v0.3.1-0.20220708130638-bd88e10a3d91/go.mod h1:MS2lj3INKhZjWNqd3N0m3J+Jxf3DAOnAH9VT3Sh9MUE=
-github.com/grafana/pyroscope-go v1.0.3 h1:8WWmItzLfg4m8G+j//ElSjMeMr88Y6Lvblar6qeTyKk=
-github.com/grafana/pyroscope-go v1.0.3/go.mod h1:0d7ftwSMBV/Awm7CCiYmHQEG8Y44Ma3YSjt+nWcWztY=
+github.com/grafana/pyroscope-go v1.1.1 h1:PQoUU9oWtO3ve/fgIiklYuGilvsm8qaGhlY4Vw6MAcQ=
+github.com/grafana/pyroscope-go v1.1.1/go.mod h1:Mw26jU7jsL/KStNSGGuuVYdUq7Qghem5P8aXYXSXG88=
 github.com/grafana/pyroscope-go/godeltaprof v0.1.8 h1:iwOtYXeeVSAeYefJNaxDytgjKtUuKQbJqgAIjlnicKg=
 github.com/grafana/pyroscope-go/godeltaprof v0.1.8/go.mod h1:2+l7K7twW49Ct4wFluZD3tZ6e0SjanjcUUBPVD/UuGU=
+github.com/grafana/pyroscope-go/x/k6 v0.0.0-20240920080526-46ebf824bf43 h1:/ThZsTNBh/LfiJEtICIy8ht6syUyE7+STQU3nwjDleU=
+github.com/grafana/pyroscope-go/x/k6 v0.0.0-20240920080526-46ebf824bf43/go.mod h1:nfbW6/4ke3ywlqLb+Zgr9t1z9Zv3m+2ImUp+vbkzHpc=
 github.com/grafana/regexp v0.0.0-20221123153739-15dc172cd2db h1:7aN5cccjIqCLTzedH7MZzRZt5/lsAHch6Z3L2ZGn5FA=
 github.com/grafana/regexp v0.0.0-20221123153739-15dc172cd2db/go.mod h1:M5qHK+eWfAv8VR/265dIuEpL3fNfeC21tXXp9itM24A=
 github.com/grpc-ecosystem/go-grpc-middleware v1.4.0 h1:UH//fgunKIs4JdUbpDl1VZCDaL56wXCB/5+wF6uHfaI=

pkg/phlare/modules.go

@@ -56,6 +56,7 @@ import (
 	"github.com/grafana/pyroscope/pkg/usagestats"
 	"github.com/grafana/pyroscope/pkg/util"
 	"github.com/grafana/pyroscope/pkg/util/build"
+	httputil "github.com/grafana/pyroscope/pkg/util/http"
 	"github.com/grafana/pyroscope/pkg/validation"
 	"github.com/grafana/pyroscope/pkg/validation/exporter"
 )
@@ -314,7 +315,13 @@ func (f *Phlare) initQuerier() (services.Service, error) {
 		f.API.RegisterQuerierServiceHandler(querierSvc)
 		f.API.RegisterVCSServiceHandler(querierSvc)
 	}
-	qWorker, err := worker.NewQuerierWorker(f.Cfg.Worker, querier.NewGRPCHandler(querierSvc), log.With(f.logger, "component", "querier-worker"), f.reg)
+
+	qWorker, err := worker.NewQuerierWorker(
+		f.Cfg.Worker,
+		querier.NewGRPCHandler(querierSvc, f.Cfg.SelfProfiling.UseK6Middleware),
+		log.With(f.logger, "component", "querier-worker"),
+		f.reg,
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -534,7 +541,12 @@ func (f *Phlare) initServer() (services.Service, error) {
 		},
 		httpMetric,
 		objstoreTracerMiddleware,
+		httputil.K6Middleware(),
 	}
+	if f.Cfg.SelfProfiling.UseK6Middleware {
+		defaultHTTPMiddleware = append(defaultHTTPMiddleware, httputil.K6Middleware())
+	}
+
 	f.Server.HTTPServer.Handler = middleware.Merge(defaultHTTPMiddleware...).Wrap(f.Server.HTTP)
 
 	s := NewServerService(f.Server, servicesToWaitFor, f.logger)

pkg/phlare/phlare.go

@@ -134,13 +134,15 @@ type SelfProfilingConfig struct {
 	DisablePush          bool `yaml:"disable_push,omitempty"`
 	MutexProfileFraction int  `yaml:"mutex_profile_fraction,omitempty"`
 	BlockProfileRate     int  `yaml:"block_profile_rate,omitempty"`
+	UseK6Middleware      bool `yaml:"use_k6_middleware,omitempty"`
 }
 
 func (c *SelfProfilingConfig) RegisterFlags(f *flag.FlagSet) {
 	// these are values that worked well in OG Pyroscope Cloud without adding much overhead
 	f.IntVar(&c.MutexProfileFraction, "self-profiling.mutex-profile-fraction", 5, "")
 	f.IntVar(&c.BlockProfileRate, "self-profiling.block-profile-rate", 5, "")
 	f.BoolVar(&c.DisablePush, "self-profiling.disable-push", false, "When running in single binary (--target=all) Pyroscope will push (Go SDK) profiles to itself. Set to true to disable self-profiling.")
+	f.BoolVar(&c.UseK6Middleware, "self-profiling.use-k6-middleware", false, "Read k6 labels from request headers and set them as dynamic profile tags.")
 }
 
 func (c *Config) RegisterFlags(f *flag.FlagSet) {

pkg/querier/grpc_handler.go

@@ -7,16 +7,23 @@ import (
 	vcsv1connect "github.com/grafana/pyroscope/api/gen/proto/go/vcs/v1/vcsv1connect"
 	connectapi "github.com/grafana/pyroscope/pkg/api/connect"
 	"github.com/grafana/pyroscope/pkg/util/connectgrpc"
+	httputil "github.com/grafana/pyroscope/pkg/util/http"
 )
 
 type QuerierSvc interface {
 	querierv1connect.QuerierServiceHandler
 	vcsv1connect.VCSServiceHandler
 }
 
-func NewGRPCHandler(svc QuerierSvc) connectgrpc.GRPCHandler {
+func NewGRPCHandler(svc QuerierSvc, useK6Middleware bool) connectgrpc.GRPCHandler {
 	mux := http.NewServeMux()
 	mux.Handle(querierv1connect.NewQuerierServiceHandler(svc, connectapi.DefaultHandlerOptions()...))
 	mux.Handle(vcsv1connect.NewVCSServiceHandler(svc, connectapi.DefaultHandlerOptions()...))
+
+	if useK6Middleware {
+		httpMiddleware := httputil.K6Middleware()
+		return connectgrpc.NewHandler(httpMiddleware.Wrap(mux))
+	}
+
 	return connectgrpc.NewHandler(mux)
 }

pkg/util/http/middleware.go

@@ -0,0 +1,19 @@
+package http
+
+import (
+	"net/http"
+
+	"github.com/grafana/dskit/middleware"
+	"github.com/grafana/pyroscope-go/x/k6"
+)
+
+// K6Middleware creates a middleware that extracts k6 load test labels from the
+// request baggage and adds them as dynamic profiling labels.
+func K6Middleware() middleware.Interface {
+	return middleware.Func(func(h http.Handler) http.Handler {
+		next := k6.LabelsFromBaggageHandler(h)
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			next.ServeHTTP(w, r)
+		})
+	})
+}

tools/k6/lib/request.js

@@ -9,6 +9,7 @@ import { READ_TOKEN, TENANT_ID, BASE_URL } from './env.js';
 export function doSelectMergeProfileRequest(body, headers) {
   const res = http.post(`${BASE_URL}/querier.v1.QuerierService/SelectMergeProfile`, JSON.stringify(body), {
     headers: withHeaders(headers),
+    tags: { name: "querier.v1.QuerierService/SelectMergeProfile" }
   });
 
   check(res, {
@@ -35,6 +36,7 @@ export function doRenderRequest(body, headers) {
 export function doSelectMergeStacktracesRequest(body, headers) {
   const res = http.post(`${BASE_URL}/querier.v1.QuerierService/SelectMergeStacktraces`, JSON.stringify(body), {
     headers: withHeaders(headers),
+    tags: { name: "querier.v1.QuerierService/SelectMergeStacktraces" }
   });
 
   check(res, {
@@ -45,6 +47,7 @@ export function doSelectMergeStacktracesRequest(body, headers) {
 export function doLabelNamesRequest(body, headers) {
   const res = http.post(`${BASE_URL}/querier.v1.QuerierService/LabelNames`, JSON.stringify(body), {
     headers: withHeaders(headers),
+    tags: { name: "querier.v1.QuerierService/LabelNames" }
   });
 
   check(res, {
@@ -55,6 +58,7 @@ export function doLabelNamesRequest(body, headers) {
 export function doSeriesRequest(body, headers) {
   const res = http.post(`${BASE_URL}/querier.v1.QuerierService/Series`, JSON.stringify(body), {
     headers: withHeaders(headers),
+    tags: { name: "querier.v1.QuerierService/Series" }
   });
 
   check(res, {
@@ -70,7 +74,7 @@ export function doRenderDiffRequest(body, headers) {
 
   const res = http.get(params.toString(), {
     headers: withHeaders(headers),
-    tags: { name: '/render' },
+    tags: { name: '/render-diff' },
   });
 
   check(res, {

tools/k6/tests/reads.js

@@ -13,7 +13,7 @@ import {
 export const options = {
   ext: {
     loadimpact: {
-      projectID: 3700226,
+      projectID: 16425,
       name: 'reads',
     },
   },