From 733af535ea3aae51bc961e7b5d36910234243ec3 Mon Sep 17 00:00:00 2001 From: Bryan Huhta <32787160+bryanhuhta@users.noreply.github.com> Date: Thu, 3 Oct 2024 15:50:59 -0500 Subject: [PATCH] feat: Add k6 middleware (#3580) --- cmd/pyroscope/help-all.txt.tmpl | 2 ++ cmd/pyroscope/help.txt.tmpl | 2 ++ .../index.md | 4 +++ go.mod | 13 +++++----- go.sum | 26 ++++++++++--------- go.work.sum | 5 ++-- pkg/clientpool/ingester_client_pool.go | 2 +- pkg/clientpool/store_gateway_client_pool.go | 2 +- pkg/phlare/modules.go | 14 +++++++++- pkg/phlare/phlare.go | 2 ++ pkg/querier/grpc_handler.go | 9 ++++++- pkg/querier/ingester_querier.go | 6 +++-- pkg/querier/replication.go | 7 ++++- pkg/storegateway/clientpool/client_pool.go | 2 +- pkg/util/http.go | 22 ++++++++++++++++ pkg/util/http/middleware.go | 16 ++++++++++++ tools/k6/lib/request.js | 6 ++++- tools/k6/run.sh | 2 +- tools/k6/tests/reads.js | 2 +- 19 files changed, 113 insertions(+), 31 deletions(-) create mode 100644 pkg/util/http/middleware.go diff --git a/cmd/pyroscope/help-all.txt.tmpl b/cmd/pyroscope/help-all.txt.tmpl index 33d4f6f967..018c23051b 100644 --- a/cmd/pyroscope/help-all.txt.tmpl +++ b/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 diff --git a/cmd/pyroscope/help.txt.tmpl b/cmd/pyroscope/help.txt.tmpl index c133f3e8d3..260b1fd361 100644 --- a/cmd/pyroscope/help.txt.tmpl +++ b/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 diff --git a/docs/sources/configure-server/reference-configuration-parameters/index.md b/docs/sources/configure-server/reference-configuration-parameters/index.md index f78698d032..69cedaaa66 100644 --- a/docs/sources/configure-server/reference-configuration-parameters/index.md +++ b/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: | 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: | 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 diff --git a/go.mod b/go.mod index a6f6aebc2d..2bcb6e894e 100644 --- a/go.mod +++ b/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.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 @@ -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 @@ -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 @@ -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 @@ -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 diff --git a/go.sum b/go.sum index 2f6cd0f7a4..a1f59cda6a 100644 --- a/go.sum +++ b/go.sum @@ -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= @@ -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= @@ -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= @@ -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= diff --git a/go.work.sum b/go.work.sum index 535a18aba3..902a5cce1a 100644 --- a/go.work.sum +++ b/go.work.sum @@ -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= @@ -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= @@ -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= diff --git a/pkg/clientpool/ingester_client_pool.go b/pkg/clientpool/ingester_client_pool.go index e5511fea5b..15e290acbf 100644 --- a/pkg/clientpool/ingester_client_pool.go +++ b/pkg/clientpool/ingester_client_pool.go @@ -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), diff --git a/pkg/clientpool/store_gateway_client_pool.go b/pkg/clientpool/store_gateway_client_pool.go index a4a22c6856..64436e394c 100644 --- a/pkg/clientpool/store_gateway_client_pool.go +++ b/pkg/clientpool/store_gateway_client_pool.go @@ -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), diff --git a/pkg/phlare/modules.go b/pkg/phlare/modules.go index 4ae954f72d..22bb1785ef 100644 --- a/pkg/phlare/modules.go +++ b/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) diff --git a/pkg/phlare/phlare.go b/pkg/phlare/phlare.go index 386a47eec5..d164a8fe49 100644 --- a/pkg/phlare/phlare.go +++ b/pkg/phlare/phlare.go @@ -134,6 +134,7 @@ 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) { @@ -141,6 +142,7 @@ func (c *SelfProfilingConfig) RegisterFlags(f *flag.FlagSet) { 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) { diff --git a/pkg/querier/grpc_handler.go b/pkg/querier/grpc_handler.go index be32c05282..42d49f43f5 100644 --- a/pkg/querier/grpc_handler.go +++ b/pkg/querier/grpc_handler.go @@ -7,6 +7,7 @@ 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 { @@ -14,9 +15,15 @@ type QuerierSvc interface { 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) } diff --git a/pkg/querier/ingester_querier.go b/pkg/querier/ingester_querier.go index 739ad2765b..03b2d4b907 100644 --- a/pkg/querier/ingester_querier.go +++ b/pkg/querier/ingester_querier.go @@ -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 diff --git a/pkg/querier/replication.go b/pkg/querier/replication.go index fce35368ce..fdcf1dfd4f 100644 --- a/pkg/querier/replication.go +++ b/pkg/querier/replication.go @@ -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 ( diff --git a/pkg/storegateway/clientpool/client_pool.go b/pkg/storegateway/clientpool/client_pool.go index d8400eb478..a43e515475 100644 --- a/pkg/storegateway/clientpool/client_pool.go +++ b/pkg/storegateway/clientpool/client_pool.go @@ -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), diff --git a/pkg/util/http.go b/pkg/util/http.go index 3149f5df42..c55f6eb8fb 100644 --- a/pkg/util/http.go +++ b/pkg/util/http.go @@ -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" @@ -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 diff --git a/pkg/util/http/middleware.go b/pkg/util/http/middleware.go new file mode 100644 index 0000000000..69b247f5d9 --- /dev/null +++ b/pkg/util/http/middleware.go @@ -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) + }) +} diff --git a/tools/k6/lib/request.js b/tools/k6/lib/request.js index 1d631cfe12..36136e2840 100644 --- a/tools/k6/lib/request.js +++ b/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, { diff --git a/tools/k6/run.sh b/tools/k6/run.sh index ec0ff92fc9..596b8fd118 100755 --- a/tools/k6/run.sh +++ b/tools/k6/run.sh @@ -51,7 +51,7 @@ DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" ; pwd -P) source "${DIR}"/.env if [ -n "${IS_CLOUD}" ]; then - k6 cloud tools/k6/tests/"${TEST}" \ + k6 cloud run tools/k6/tests/"${TEST}" \ -e "K6_BASE_URL=${K6_BASE_URL}" \ -e "K6_READ_TOKEN=${K6_READ_TOKEN}" \ -e "K6_TENANT_ID=${K6_TENANT_ID}" \ diff --git a/tools/k6/tests/reads.js b/tools/k6/tests/reads.js index 2b5e157a6f..09cba8ddde 100644 --- a/tools/k6/tests/reads.js +++ b/tools/k6/tests/reads.js @@ -13,7 +13,7 @@ import { export const options = { ext: { loadimpact: { - projectID: 3700226, + projectID: 16425, name: 'reads', }, },