Skip to content

Commit 3ae4e48

Browse files
devgianludevgianlu
committed
feat: add k6 gRPC stream middleware
1 parent fe8f520 commit 3ae4e48

9 files changed

+415
-30
lines changed

‎go.mod

+1-1
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ replace github.com/grafana/pyroscope-go/godeltaprof => ./godeltaprof
66

77
require (
88
github.com/grafana/pyroscope-go/godeltaprof v0.1.8
9-
github.com/stretchr/testify v1.9.0
9+
github.com/stretchr/testify v1.10.0
1010
)
1111

1212
require (

‎go.sum

+1-2
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,7 @@ github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/
2424
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
2525
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
2626
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
27-
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
28-
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
27+
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
2928
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
3029
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
3130
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=

‎go.work.sum

+1
Original file line numberDiff line numberDiff line change
@@ -303,6 +303,7 @@ github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJ
303303
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e h1:aoZm08cpOy4WuID//EZDgcC4zIxODThtZNPirFr42+A=
304304
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo=
305305
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8=
306+
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
306307
github.com/yuin/goldmark v1.4.13 h1:fVcFKWvrslecOb/tg+Cc05dkeYx540o0FuFt3nUVDoE=
307308
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
308309
go.opentelemetry.io/otel v1.13.0 h1:1ZAKnNQKwBBxFtww/GwxNUyTf0AxkZzrukO8MeXqe4Y=

‎x/k6/baggage.go

+27
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package k6
22

33
import (
44
"context"
5+
grpcmiddleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
56
"net/http"
67
"runtime/pprof"
78
"strings"
@@ -43,6 +44,32 @@ func LabelsFromBaggageUnaryInterceptor(ctx context.Context, req interface{}, _ *
4344
return handler(ctx, req)
4445
}
4546

47+
func LabelsFromBaggageStreamInterceptor(srv interface{}, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
48+
ctx := ss.Context()
49+
50+
var found bool
51+
ctx, found = setBaggageContextFromMetadata(ctx)
52+
if !found {
53+
return handler(srv, ss)
54+
}
55+
56+
labels := getBaggageLabelsFromContext(ctx)
57+
if labels == nil {
58+
return handler(srv, ss)
59+
}
60+
61+
// Inlined version of pyroscope.TagWrapper and pprof.Do to reduce noise in
62+
// the stack trace.
63+
defer pprof.SetGoroutineLabels(ctx)
64+
ctx = pprof.WithLabels(ctx, *labels)
65+
pprof.SetGoroutineLabels(ctx)
66+
67+
return handler(srv, &grpcmiddleware.WrappedServerStream{
68+
ServerStream: ss,
69+
WrappedContext: ctx,
70+
})
71+
}
72+
4673
type labelHandler struct {
4774
innerHandler http.Handler
4875
}

‎x/k6/baggage_test.go

+78
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package k6
22

33
import (
44
"context"
5+
grpcmiddleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
56
"net/http"
67
"net/http/httptest"
78
"runtime/pprof"
@@ -153,6 +154,83 @@ func TestLabelsFromBaggageUnaryInterceptor(t *testing.T) {
153154
})
154155
}
155156

157+
func TestLabelsFromBaggageStreamInterceptor(t *testing.T) {
158+
info := &grpc.StreamServerInfo{
159+
FullMethod: "/example.ExampleService/Test",
160+
}
161+
162+
t.Run("adds_k6_labels_from_grpc_baggage", func(t *testing.T) {
163+
testCtx := testAddBaggageToGRPCRequest(t, context.Background(),
164+
"k6.test_run_id", "123",
165+
"not_k6.some_other_key", "value",
166+
)
167+
168+
handler := func(srv any, stream grpc.ServerStream) error {
169+
ctx := stream.Context()
170+
171+
b := baggage.FromContext(ctx)
172+
require.NotNil(t, b)
173+
testAssertEqualMembers(t, b.Members(),
174+
"k6.test_run_id", "123",
175+
"not_k6.some_other_key", "value",
176+
)
177+
178+
val, ok := pprof.Label(ctx, "k6_test_run_id")
179+
require.True(t, ok)
180+
require.Equal(t, "123", val)
181+
182+
_, ok = pprof.Label(ctx, "not_k6_some_other_key")
183+
require.False(t, ok)
184+
185+
return nil
186+
}
187+
188+
err := LabelsFromBaggageStreamInterceptor(nil, &grpcmiddleware.WrappedServerStream{WrappedContext: testCtx}, info, handler)
189+
require.NoError(t, err)
190+
})
191+
192+
t.Run("passthrough_requests_with_no_baggage", func(t *testing.T) {
193+
testCtx := testAddBaggageToGRPCRequest(t, context.Background())
194+
195+
handler := func(srv any, stream grpc.ServerStream) error {
196+
ctx := stream.Context()
197+
198+
b := baggage.FromContext(ctx)
199+
require.NotNil(t, b)
200+
require.Equal(t, 0, b.Len())
201+
202+
return nil
203+
}
204+
205+
err := LabelsFromBaggageStreamInterceptor(nil, &grpcmiddleware.WrappedServerStream{WrappedContext: testCtx}, info, handler)
206+
require.NoError(t, err)
207+
})
208+
209+
t.Run("passthrough_requests_with_no_k6_baggage", func(t *testing.T) {
210+
testCtx := testAddBaggageToGRPCRequest(t, context.Background(),
211+
"not_k6.some_other_key", "value",
212+
)
213+
214+
handler := func(srv any, stream grpc.ServerStream) error {
215+
ctx := stream.Context()
216+
217+
b := baggage.FromContext(ctx)
218+
require.NotNil(t, b)
219+
testAssertEqualMembers(t, b.Members(),
220+
"not_k6.some_other_key", "value",
221+
)
222+
223+
_, ok := pprof.Label(ctx, "not_k6_some_other_key")
224+
require.False(t, ok)
225+
226+
return nil
227+
}
228+
229+
err := LabelsFromBaggageStreamInterceptor(nil, &grpcmiddleware.WrappedServerStream{WrappedContext: testCtx}, info, handler)
230+
require.NoError(t, err)
231+
})
232+
}
233+
156234
func Test_setBaggageContextFromHeader(t *testing.T) {
157235
t.Run("sets_baggage_context", func(t *testing.T) {
158236
req := httptest.NewRequest("GET", "http://example.com", nil)

‎x/k6/go.mod

+10-7
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
module github.com/grafana/pyroscope-go/x/k6
22

3-
go 1.22
3+
go 1.23
4+
5+
toolchain go1.23.4
46

57
replace github.com/grafana/pyroscope-go => ../../
68

79
require (
810
github.com/grafana/pyroscope-go v1.1.1
9-
github.com/stretchr/testify v1.9.0
11+
github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.0
12+
github.com/stretchr/testify v1.10.0
1013
go.opentelemetry.io/otel v1.24.0
1114
google.golang.org/grpc v1.67.1
1215
)
@@ -18,11 +21,11 @@ require (
1821
github.com/kr/text v0.2.0 // indirect
1922
github.com/pmezard/go-difflib v1.0.0 // indirect
2023
github.com/rogpeppe/go-internal v1.9.0 // indirect
21-
golang.org/x/net v0.28.0 // indirect
22-
golang.org/x/sys v0.28.0 // indirect
23-
golang.org/x/text v0.21.0 // indirect
24-
google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 // indirect
25-
google.golang.org/protobuf v1.34.2 // indirect
24+
golang.org/x/net v0.35.0 // indirect
25+
golang.org/x/sys v0.30.0 // indirect
26+
golang.org/x/text v0.22.0 // indirect
27+
google.golang.org/genproto/googleapis/rpc v0.0.0-20240826202546-f6391c0de4c7 // indirect
28+
google.golang.org/protobuf v1.36.4 // indirect
2629
gopkg.in/yaml.v3 v3.0.1 // indirect
2730
)
2831

‎x/k6/go.sum

+12-10
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
55
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
66
github.com/grafana/pyroscope-go/godeltaprof v0.1.8 h1:iwOtYXeeVSAeYefJNaxDytgjKtUuKQbJqgAIjlnicKg=
77
github.com/grafana/pyroscope-go/godeltaprof v0.1.8/go.mod h1:2+l7K7twW49Ct4wFluZD3tZ6e0SjanjcUUBPVD/UuGU=
8+
github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.0 h1:FbSCl+KggFl+Ocym490i/EyXF4lPgLoUtcSWquBM0Rs=
9+
github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.0/go.mod h1:qOchhhIlmRcqk/O9uCo/puJlyo07YINaIqdZfZG3Jkc=
810
github.com/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0NAMnU=
911
github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
1012
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
@@ -17,22 +19,22 @@ github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZV
1719
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
1820
github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
1921
github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
20-
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
21-
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
22+
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
23+
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
2224
go.opentelemetry.io/otel v1.24.0 h1:0LAOdjNmQeSTzGBzduGe/rU4tZhMwL5rWgtp9Ku5Jfo=
2325
go.opentelemetry.io/otel v1.24.0/go.mod h1:W7b9Ozg4nkF5tWI5zsXkaKKDjdVjpD4oAt9Qi/MArHo=
2426
golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
2527
golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
26-
golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA=
27-
golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
28-
golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
29-
golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
30-
google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 h1:e7S5W7MGGLaSu8j3YjdezkZ+m1/Nm0uRVRMEMGk26Xs=
31-
google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142/go.mod h1:UqMtugtsSgubUsoxbuAoiCXvqvErP7Gf0so0mK9tHxU=
28+
golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
29+
golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
30+
golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM=
31+
golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY=
32+
google.golang.org/genproto/googleapis/rpc v0.0.0-20240826202546-f6391c0de4c7 h1:2035KHhUv+EpyB+hWgJnaWKJOdX1E95w2S8Rr4uWKTs=
33+
google.golang.org/genproto/googleapis/rpc v0.0.0-20240826202546-f6391c0de4c7/go.mod h1:UqMtugtsSgubUsoxbuAoiCXvqvErP7Gf0so0mK9tHxU=
3234
google.golang.org/grpc v1.67.1 h1:zWnc1Vrcno+lHZCOofnIMvycFcc0QRGIzm9dhnDX68E=
3335
google.golang.org/grpc v1.67.1/go.mod h1:1gLDyUQU7CTLJI90u3nXZ9ekeghjeM7pTDZlqFNg2AA=
34-
google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
35-
google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
36+
google.golang.org/protobuf v1.36.4 h1:6A3ZDJHn/eNqc1i+IdefRzy/9PokBTPvcqMySR7NNIM=
37+
google.golang.org/protobuf v1.36.4/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
3638
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
3739
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
3840
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=

‎x/k6/go.work

+3-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
go 1.22
1+
go 1.23
2+
3+
toolchain go1.23.4
24

35
use (
46
.

0 commit comments

Comments
 (0)