Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: k6 middleware #3580

Merged
merged 23 commits into from
Oct 3, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions cmd/pyroscope/help-all.txt.tmpl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions cmd/pyroscope/help.txt.tmpl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
13 changes: 7 additions & 6 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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.2.0
github.com/grafana/pyroscope-go/godeltaprof v0.1.8
github.com/grafana/pyroscope-go/x/k6 v0.0.0-20241003203156-a917cea171d3
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
Expand All @@ -44,7 +45,7 @@ require (
github.com/iancoleman/strcase v0.3.0
github.com/json-iterator/go v1.1.12
github.com/k0kubun/pp/v3 v3.2.0
github.com/klauspost/compress v1.17.9
github.com/klauspost/compress v1.17.10
github.com/kubescape/go-git-url v0.0.27
github.com/mattn/go-isatty v0.0.19
github.com/minio/minio-go/v7 v7.0.72
Expand Down Expand Up @@ -74,6 +75,7 @@ require (
github.com/valyala/bytebufferpool v1.0.0
github.com/xlab/treeprint v1.2.0
go.etcd.io/bbolt v1.3.10
go.opentelemetry.io/otel v1.30.0
go.opentelemetry.io/proto/otlp v1.1.0
go.uber.org/atomic v1.11.0
go.uber.org/goleak v1.3.0
Expand Down Expand Up @@ -141,7 +143,7 @@ require (
github.com/efficientgo/e2e v0.14.1-0.20230710114240-c316eb95ae5b // indirect
github.com/facette/natsort v0.0.0-20181210072756-2cd4dd1e2dcb // indirect
github.com/go-logfmt/logfmt v0.6.0 // indirect
github.com/go-logr/logr v1.4.1 // indirect
github.com/go-logr/logr v1.4.2 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/go-openapi/analysis v0.22.2 // indirect
github.com/go-openapi/errors v0.21.1 // indirect
Expand Down Expand Up @@ -224,9 +226,8 @@ require (
go.opentelemetry.io/collector/semconv v0.96.0 // indirect
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0 // indirect
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 // indirect
go.opentelemetry.io/otel v1.24.0 // indirect
go.opentelemetry.io/otel/metric v1.24.0 // indirect
go.opentelemetry.io/otel/trace v1.24.0 // indirect
go.opentelemetry.io/otel/metric v1.30.0 // indirect
go.opentelemetry.io/otel/trace v1.30.0 // indirect
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.24.0 // indirect
golang.org/x/crypto v0.24.0 // indirect
Expand Down
26 changes: 14 additions & 12 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -245,8 +245,8 @@ github.com/go-logfmt/logfmt v0.5.0/go.mod h1:wCYkCAKZfumFQihp8CzCvQ3paCTfi41vtzG
github.com/go-logfmt/logfmt v0.6.0 h1:wGYYu3uicYdqXVgoYbvnkrPVXkuLM1p1ifugDMEdRi4=
github.com/go-logfmt/logfmt v0.6.0/go.mod h1:WYhtIu8zTZfxdn5+rREduYbwxfcBr/Vr6KEVveWlfTs=
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
github.com/go-logr/logr v1.4.1 h1:pKouT5E8xu9zeFC39JXRDukb6JFQPXM5p5I91188VAQ=
github.com/go-logr/logr v1.4.1/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
github.com/go-openapi/analysis v0.22.2 h1:ZBmNoP2h5omLKr/srIC9bfqrUGzT6g6gNv03HE9Vpj0=
Expand Down Expand Up @@ -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.2.0 h1:aILLKjTj8CS8f/24OPMGPewQSYlhmdQMBmol1d3KGj8=
github.com/grafana/pyroscope-go v1.2.0/go.mod h1:2GHr28Nr05bg2pElS+dDsc98f3JTUh2f6Fz1hWXrqwk=
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-20241003203156-a917cea171d3 h1:GtwQDlBz8aJHMy2Ko28UDRGgGzi7v4Vf20+ZyXaGy7M=
github.com/grafana/pyroscope-go/x/k6 v0.0.0-20241003203156-a917cea171d3/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=
Expand Down Expand Up @@ -499,8 +501,8 @@ github.com/k0kubun/pp/v3 v3.2.0 h1:h33hNTZ9nVFNP3u2Fsgz8JXiF5JINoZfFq4SvKJwNcs=
github.com/k0kubun/pp/v3 v3.2.0/go.mod h1:ODtJQbQcIRfAD3N+theGCV1m/CBxweERz2dapdz1EwA=
github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA=
github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
github.com/klauspost/compress v1.17.10 h1:oXAz+Vh0PMUvJczoi+flxpnBEPxoER1IaAnU/NMPtT0=
github.com/klauspost/compress v1.17.10/go.mod h1:pMDklpSncoRMuLFrf1W9Ss9KT+0rH90U12bZKk7uwG0=
github.com/klauspost/cpuid/v2 v2.0.1/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
github.com/klauspost/cpuid/v2 v2.2.8 h1:+StwCXwm9PdpiEkPyzBXIy+M9KUb4ODm0Zarf1kS5BM=
github.com/klauspost/cpuid/v2 v2.2.8/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
Expand Down Expand Up @@ -760,14 +762,14 @@ go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.4
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0/go.mod h1:Mjt1i1INqiaoZOMGR1RIUJN+i3ChKoFRqzrRQhlkbs0=
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 h1:jq9TW8u3so/bN+JPT166wjOI6/vQPF6Xe7nMNIltagk=
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0/go.mod h1:p8pYQP+m5XfbZm9fxtSKAbM6oIllS7s2AfxrChvc7iw=
go.opentelemetry.io/otel v1.24.0 h1:0LAOdjNmQeSTzGBzduGe/rU4tZhMwL5rWgtp9Ku5Jfo=
go.opentelemetry.io/otel v1.24.0/go.mod h1:W7b9Ozg4nkF5tWI5zsXkaKKDjdVjpD4oAt9Qi/MArHo=
go.opentelemetry.io/otel/metric v1.24.0 h1:6EhoGWWK28x1fbpA4tYTOWBkPefTDQnb8WSGXlc88kI=
go.opentelemetry.io/otel/metric v1.24.0/go.mod h1:VYhLe1rFfxuTXLgj4CBiyz+9WYBA8pNGJgDcSFRKBco=
go.opentelemetry.io/otel v1.30.0 h1:F2t8sK4qf1fAmY9ua4ohFS/K+FUuOPemHUIXHtktrts=
go.opentelemetry.io/otel v1.30.0/go.mod h1:tFw4Br9b7fOS+uEao81PJjVMjW/5fvNCbpsDIXqP0pc=
go.opentelemetry.io/otel/metric v1.30.0 h1:4xNulvn9gjzo4hjg+wzIKG7iNFEaBMX00Qd4QIZs7+w=
go.opentelemetry.io/otel/metric v1.30.0/go.mod h1:aXTfST94tswhWEb+5QjlSqG+cZlmyXy/u8jFpor3WqQ=
go.opentelemetry.io/otel/sdk v1.24.0 h1:YMPPDNymmQN3ZgczicBY3B6sf9n62Dlj9pWD3ucgoDw=
go.opentelemetry.io/otel/sdk v1.24.0/go.mod h1:KVrIYw6tEubO9E96HQpcmpTKDVn9gdv35HoYiQWGDFg=
go.opentelemetry.io/otel/trace v1.24.0 h1:CsKnnL4dUAr/0llH9FKuc698G04IrpWV0MQA/Y1YELI=
go.opentelemetry.io/otel/trace v1.24.0/go.mod h1:HPc3Xr/cOApsBI154IU0OI0HJexz+aw5uPdbs3UCjNU=
go.opentelemetry.io/otel/trace v1.30.0 h1:7UBkkYzeg3C7kQX8VAidWh2biiQbtAKjyIML8dQ9wmc=
go.opentelemetry.io/otel/trace v1.30.0/go.mod h1:5EyKqTzzmyqB9bwtCCq6pDLktPK6fmGf/Dph+8VI02o=
go.opentelemetry.io/proto/otlp v1.1.0 h1:2Di21piLrCqJ3U3eXGCTPHE9R8Nh+0uglSnOyxikMeI=
go.opentelemetry.io/proto/otlp v1.1.0/go.mod h1:GpBHCBWiqvVLDqmHZsoMM3C5ySeKTC7ej/RNTae6MdY=
go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
Expand Down
5 changes: 3 additions & 2 deletions go.work.sum
Original file line number Diff line number Diff line change
Expand Up @@ -689,6 +689,7 @@ github.com/go-logfmt/logfmt v0.5.1/go.mod h1:WYhtIu8zTZfxdn5+rREduYbwxfcBr/Vr6KE
github.com/go-logr/logr v0.1.0/go.mod h1:ixOQHD9gLJUVQQ2ZOR7zLEifBX6tGkNJF4QyIY7sIas=
github.com/go-logr/logr v1.2.0/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
github.com/go-logr/logr v1.3.0/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
github.com/go-openapi/analysis v0.22.0/go.mod h1:acDnkkCI2QxIo8sSIPgmp1wUlRohV7vfGtAIVae73b0=
github.com/go-openapi/errors v0.21.0/go.mod h1:jxNTMUxRCKj65yb/okJGEtahVd7uvWnuWfj53bse4ho=
github.com/go-openapi/runtime v0.26.0 h1:HYOFtG00FM1UvqrcxbEJg/SwvDRvYLQKGhw2zaQjTcc=
Expand Down Expand Up @@ -1124,11 +1125,13 @@ go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.24.0/go.mod h
go.opentelemetry.io/otel/metric v1.17.0/go.mod h1:h4skoxdZI17AxwITdmdZjjYJQH5nzijUUjm+wtPph5o=
go.opentelemetry.io/otel/metric v1.21.0/go.mod h1:o1p3CA8nNHW8j5yuQLdc1eeqEaPfzug24uvsyIEJRWM=
go.opentelemetry.io/otel/metric v1.22.0/go.mod h1:evJGjVpZv0mQ5QBRJoBF64yMuOf4xCWdXjK8pzFvliY=
go.opentelemetry.io/otel/metric v1.30.0 h1:4xNulvn9gjzo4hjg+wzIKG7iNFEaBMX00Qd4QIZs7+w=
go.opentelemetry.io/otel/sdk v1.19.0/go.mod h1:NedEbbS4w3C6zElbLdPJKOpJQOrGUJ+GfzpjUvI0v1A=
go.opentelemetry.io/otel/sdk v1.21.0/go.mod h1:Nna6Yv7PWTdgJHVRD9hIYywQBRx7pbox6nwBnZIxl/E=
go.opentelemetry.io/otel/trace v1.17.0/go.mod h1:I/4vKTgFclIsXRVucpH25X0mpFSczM7aHeaz0ZBLWjY=
go.opentelemetry.io/otel/trace v1.21.0/go.mod h1:LGbsEB0f9LGjN+OZaQQ26sohbOmiMR+BaslueVtS/qQ=
go.opentelemetry.io/otel/trace v1.22.0/go.mod h1:RbbHXVqKES9QhzZq/fE5UnOSILqRt40a21sPw2He1xo=
go.opentelemetry.io/otel/trace v1.30.0 h1:7UBkkYzeg3C7kQX8VAidWh2biiQbtAKjyIML8dQ9wmc=
go.opentelemetry.io/proto/otlp v0.7.0/go.mod h1:PqfVotwruBrMGOCsRd/89rSnXhoiJIqeYNgFYFoEGnI=
go.opentelemetry.io/proto/otlp v1.0.0/go.mod h1:Sy6pihPLfYHkr3NkUbEhGHFhINUSI/v80hjKIs5JXpM=
go.uber.org/automaxprocs v1.5.3 h1:kWazyxZUrS3Gs4qUpbwo5kEIMGe/DAvi5Z4tl2NW4j8=
Expand Down Expand Up @@ -1214,8 +1217,6 @@ golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.2.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
golang.org/x/sync v0.4.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
golang.org/x/sys v0.25.0 h1:r+8e+loiHxRqhXVl6ML1nO3l1+oFoWbnlu2Ehimmi34=
golang.org/x/sys v0.25.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2 h1:IRJeR9r1pYWsHKTRe/IInb7lYvbBVIqOgsX/u0mbOWY=
golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
golang.org/x/term v0.13.0/go.mod h1:LTmsnFJwVN6bCy1rVCoS+qHT1HhALEFxKncY3WNNh4U=
Expand Down
2 changes: 1 addition & 1 deletion pkg/clientpool/ingester_client_pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ func (f *ingesterPoolFactory) FromInstance(inst ring.InstanceDesc) (ring_client.
return nil, err
}

httpClient := util.InstrumentedDefaultHTTPClient(util.WithTracingTransport())
httpClient := util.InstrumentedDefaultHTTPClient(util.WithTracingTransport(), util.WithBaggageTransport())
return &ingesterPoolClient{
IngesterServiceClient: ingesterv1connect.NewIngesterServiceClient(httpClient, "http://"+inst.Addr, f.options...),
HealthClient: grpc_health_v1.NewHealthClient(conn),
Expand Down
2 changes: 1 addition & 1 deletion pkg/clientpool/store_gateway_client_pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ func (f *storeGatewayPoolFactory) FromInstance(inst ring.InstanceDesc) (ring_cli
return nil, err
}

httpClient := util.InstrumentedDefaultHTTPClient(util.WithTracingTransport())
httpClient := util.InstrumentedDefaultHTTPClient(util.WithTracingTransport(), util.WithBaggageTransport())
return &storeGatewayPoolClient{
StoreGatewayServiceClient: storegatewayv1connect.NewStoreGatewayServiceClient(httpClient, "http://"+inst.Addr, f.options...),
HealthClient: grpc_health_v1.NewHealthClient(conn),
Expand Down
14 changes: 13 additions & 1 deletion pkg/phlare/modules.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions pkg/phlare/phlare.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
9 changes: 8 additions & 1 deletion pkg/querier/grpc_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
6 changes: 4 additions & 2 deletions pkg/querier/ingester_querier.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,13 +58,15 @@ func forAllIngesters[T any](ctx context.Context, ingesterQuerier *IngesterQuerie
if err != nil {
return nil, err
}
return forGivenReplicationSet(ctx, func(addr string) (IngesterQueryClient, error) {

clientFactoryFn := func(addr string) (IngesterQueryClient, error) {
client, err := ingesterQuerier.pool.GetClientFor(addr)
if err != nil {
return nil, err
}
return client.(IngesterQueryClient), nil
}, replicationSet, f)
}
return forGivenReplicationSet(ctx, clientFactoryFn, replicationSet, f)
}

// forAllPlannedIngesters runs f, in parallel, for all ingesters part of the plan
Expand Down
7 changes: 6 additions & 1 deletion pkg/querier/replication.go
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,12 @@ func forGivenReplicationSet[Result any, Querier any](ctx context.Context, client
}

// forGivenPlan runs f, in parallel, for given plan.
func forGivenPlan[Result any, Querier any](ctx context.Context, plan map[string]*blockPlanEntry, clientFactory func(string) (Querier, error), replicationSet ring.ReplicationSet, f QueryReplicaWithHintsFn[Result, Querier]) ([]ResponseFromReplica[Result], error) {
func forGivenPlan[Result any, Querier any](
ctx context.Context,
plan map[string]*blockPlanEntry,
clientFactory func(string) (Querier, error),
replicationSet ring.ReplicationSet, f QueryReplicaWithHintsFn[Result, Querier],
) ([]ResponseFromReplica[Result], error) {
g, _ := errgroup.WithContext(ctx)

var (
Expand Down
2 changes: 1 addition & 1 deletion pkg/storegateway/clientpool/client_pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ func (f *poolFactory) FromInstance(inst ring.InstanceDesc) (ring_client.PoolClie
return nil, err
}

httpClient := util.InstrumentedDefaultHTTPClient(util.WithTracingTransport())
httpClient := util.InstrumentedDefaultHTTPClient(util.WithTracingTransport(), util.WithBaggageTransport())
return &storeGatewayPoolClient{
StoreGatewayServiceClient: storegatewayv1connect.NewStoreGatewayServiceClient(httpClient, "http://"+inst.Addr, f.options...),
HealthClient: grpc_health_v1.NewHealthClient(conn),
Expand Down
22 changes: 22 additions & 0 deletions pkg/util/http.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import (
"time"

"github.com/grafana/dskit/instrument"
"go.opentelemetry.io/otel/baggage"

"github.com/dustin/go-humanize"
"github.com/felixge/httpsnoop"
Expand Down Expand Up @@ -84,6 +85,27 @@ func WithTracingTransport() RoundTripperInstrumentFunc {
}
}

// WithBaggageTransport will set the Baggage header on the request if there is
// any baggage in the context and it was not already set.
func WithBaggageTransport() RoundTripperInstrumentFunc {
return func(next http.RoundTripper) http.RoundTripper {
return RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
_, ok := req.Header["Baggage"]
if ok {
return next.RoundTrip(req)
}

b := baggage.FromContext(req.Context())
if b.Len() == 0 {
return next.RoundTrip(req)
}

req.Header.Set("Baggage", b.String())
return next.RoundTrip(req)
})
}
}

// WriteYAMLResponse writes some YAML as a HTTP response.
func WriteYAMLResponse(w http.ResponseWriter, v interface{}) {
// There is not standardised content-type for YAML, text/plain ensures the
Expand Down
16 changes: 16 additions & 0 deletions pkg/util/http/middleware.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
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 {
return k6.LabelsFromBaggageHandler(h)
})
}
Loading
Loading