Skip to content

Commit 92e7337

Browse files
jcchavezsjcchavezs
committed
feat(span-customizer): adds support for span customizer.
1 parent 1c762e0 commit 92e7337

File tree

10 files changed

+295
-42
lines changed

10 files changed

+295
-42
lines changed

‎middleware/grpc/client.go‎

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2019 The OpenZipkin Authors
1+
// Copyright 2020 The OpenZipkin Authors
22
//
33
// Licensed under the Apache License, Version 2.0 (the "License");
44
// you may not use this file except in compliance with the License.
@@ -42,25 +42,41 @@ func WithRemoteServiceName(name string) ClientOption {
4242
}
4343
}
4444

45+
// WithClientOutPayloadParser adds a parser for the stats.OutPayload to be able to access
46+
// the outgoing request payload
47+
func WithClientOutPayloadParser(parser func(*stats.OutPayload, zipkin.SpanCustomizer)) ClientOption {
48+
return func(h *clientHandler) {
49+
h.handleRPCParser.outPayload = parser
50+
}
51+
}
52+
53+
// WithClientOutHeaderParser adds a parser for the stats.OutHeader to be able to access
54+
// the outgoing request payload
55+
func WithClientOutHeaderParser(parser func(*stats.OutHeader, zipkin.SpanCustomizer)) ClientOption {
56+
return func(h *clientHandler) {
57+
h.handleRPCParser.outHeader = parser
58+
}
59+
}
60+
4561
// WithClientInPayloadParser adds a parser for the stats.InPayload to be able to access
46-
// the request payload
47-
func WithClientInPayloadParser(parser func(*stats.InPayload, zipkin.Span)) ClientOption {
62+
// the incoming response payload
63+
func WithClientInPayloadParser(parser func(*stats.InPayload, zipkin.SpanCustomizer)) ClientOption {
4864
return func(h *clientHandler) {
4965
h.handleRPCParser.inPayload = parser
5066
}
5167
}
5268

5369
// WithClientInTrailerParser adds a parser for the stats.InTrailer to be able to access
54-
// the request trailer
55-
func WithClientInTrailerParser(parser func(*stats.InTrailer, zipkin.Span)) ClientOption {
70+
// the incoming response trailer
71+
func WithClientInTrailerParser(parser func(*stats.InTrailer, zipkin.SpanCustomizer)) ClientOption {
5672
return func(h *clientHandler) {
5773
h.handleRPCParser.inTrailer = parser
5874
}
5975
}
6076

6177
// WithClientInHeaderParser adds a parser for the stats.InHeader to be able to access
62-
// the request payload
63-
func WithClientInHeaderParser(parser func(*stats.InHeader, zipkin.Span)) ClientOption {
78+
// the incoming response header
79+
func WithClientInHeaderParser(parser func(*stats.InHeader, zipkin.SpanCustomizer)) ClientOption {
6480
return func(h *clientHandler) {
6581
h.handleRPCParser.inHeader = parser
6682
}

‎middleware/grpc/client_parser_test.go‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
// Copyright 2020 The OpenZipkin Authors
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package grpc_test
16+
17+
import (
18+
"context"
19+
"testing"
20+
21+
"github.com/openzipkin/zipkin-go"
22+
zipkingrpc "github.com/openzipkin/zipkin-go/middleware/grpc"
23+
service "github.com/openzipkin/zipkin-go/proto/testing"
24+
"google.golang.org/grpc"
25+
"google.golang.org/grpc/metadata"
26+
"google.golang.org/grpc/stats"
27+
)
28+
29+
func TestGRPCClientCanAccessToPayloadAndMetadata(t *testing.T) {
30+
tracer, flusher := createTracer(false)
31+
32+
s := grpc.NewServer()
33+
defer s.Stop()
34+
35+
service.RegisterHelloServiceServer(s, &TestHelloService{
36+
responseHeader: metadata.Pairs("test_key", "test_value_1"),
37+
responseTrailer: metadata.Pairs("test_key", "test_value_2"),
38+
})
39+
40+
dialer := initListener(s)
41+
42+
ctx := context.Background()
43+
conn, err := grpc.DialContext(
44+
ctx,
45+
"bufnet",
46+
grpc.WithContextDialer(dialer),
47+
grpc.WithInsecure(),
48+
grpc.WithStatsHandler(zipkingrpc.NewClientHandler(
49+
tracer,
50+
zipkingrpc.WithClientOutPayloadParser(func(outPayload *stats.OutPayload, span zipkin.SpanCustomizer) {
51+
m, ok := outPayload.Payload.(*service.HelloRequest)
52+
if !ok {
53+
t.Fatal("failed to cast the payload as a service.HelloResponse")
54+
}
55+
if want, have := "Hello", m.Payload; want != have {
56+
t.Errorf("incorrect payload: want %q, have %q", want, have)
57+
}
58+
}),
59+
zipkingrpc.WithClientOutHeaderParser(func(outHeader *stats.OutHeader, span zipkin.SpanCustomizer) {
60+
if want, have := "test_value", outHeader.Header.Get("test_key")[0]; want != have {
61+
t.Errorf("incorrect header value, want %q, have %q", want, have)
62+
}
63+
}),
64+
zipkingrpc.WithClientInPayloadParser(func(inPayload *stats.InPayload, span zipkin.SpanCustomizer) {
65+
m, ok := inPayload.Payload.(*service.HelloResponse)
66+
if !ok {
67+
t.Fatal("failed to cast the payload as a service.HelloRequest")
68+
}
69+
if want, have := "World", m.Payload; want != have {
70+
t.Errorf("incorrect payload: want %q, have %q", want, have)
71+
}
72+
}),
73+
zipkingrpc.WithClientInHeaderParser(func(inHeader *stats.InHeader, span zipkin.SpanCustomizer) {
74+
if want, have := "test_value_1", inHeader.Header.Get("test_key")[0]; want != have {
75+
t.Errorf("incorrect header value, want %q, have %q", want, have)
76+
}
77+
}),
78+
zipkingrpc.WithClientInTrailerParser(func(inTrailer *stats.InTrailer, span zipkin.SpanCustomizer) {
79+
if want, have := "test_value_2", inTrailer.Trailer.Get("test_key")[0]; want != have {
80+
t.Errorf("incorrect header value, want %q, have %q", want, have)
81+
}
82+
}),
83+
)),
84+
)
85+
86+
if err != nil {
87+
t.Fatalf("Failed to dial bufnet: %v", err)
88+
}
89+
defer conn.Close()
90+
91+
client := service.NewHelloServiceClient(conn)
92+
93+
ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs("test_key", "test_value"))
94+
_, err = client.Hello(ctx, &service.HelloRequest{
95+
Payload: "Hello",
96+
})
97+
if err != nil {
98+
t.Fatalf("unexpected error: %v", err)
99+
}
100+
101+
spans := flusher()
102+
if want, have := 1, len(spans); want != have {
103+
t.Errorf("unexpected number of spans, want %d, have %d", want, have)
104+
}
105+
}

‎middleware/grpc/grpc_suite_test.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2019 The OpenZipkin Authors
1+
// Copyright 2020 The OpenZipkin Authors
22
//
33
// Licensed under the Apache License, Version 2.0 (the "License");
44
// you may not use this file except in compliance with the License.
@@ -125,6 +125,8 @@ func (g *sequentialIdGenerator) reset() {
125125

126126
type TestHelloService struct {
127127
service.UnimplementedHelloServiceServer
128+
responseHeader metadata.MD
129+
responseTrailer metadata.MD
128130
}
129131

130132
func (s *TestHelloService) Hello(ctx context.Context, req *service.HelloRequest) (*service.HelloResponse, error) {
@@ -158,6 +160,9 @@ func (s *TestHelloService) Hello(ctx context.Context, req *service.HelloRequest)
158160
}
159161
}
160162

163+
grpc.SetTrailer(ctx, s.responseTrailer)
164+
grpc.SendHeader(ctx, s.responseHeader)
165+
161166
return resp, nil
162167
}
163168

‎middleware/grpc/server.go‎

Lines changed: 27 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2019 The OpenZipkin Authors
1+
// Copyright 2020 The OpenZipkin Authors
22
//
33
// Licensed under the Apache License, Version 2.0 (the "License");
44
// you may not use this file except in compliance with the License.
@@ -41,26 +41,42 @@ func ServerTags(tags map[string]string) ServerOption {
4141
}
4242

4343
// WithServerInPayloadParser adds a parser for the stats.InPayload to be able to access
44-
// the request payload
45-
func WithServerInPayloadParser(parser func(*stats.InPayload, zipkin.Span)) ServerOption {
44+
// the incoming request payload
45+
func WithServerInPayloadParser(parser func(*stats.InPayload, zipkin.SpanCustomizer)) ServerOption {
4646
return func(h *serverHandler) {
4747
h.handleRPCParser.inPayload = parser
4848
}
4949
}
5050

51-
// WithserverInTrailerParser adds a parser for the stats.InTrailer to be able to access
52-
// the request trailer
53-
func WithserverInTrailerParser(parser func(*stats.InTrailer, zipkin.Span)) ServerOption {
51+
// WithServerInHeaderParser adds a parser for the stats.InHeader to be able to access
52+
// the incoming request header
53+
func WithServerInHeaderParser(parser func(*stats.InHeader, zipkin.SpanCustomizer)) ServerOption {
5454
return func(h *serverHandler) {
55-
h.handleRPCParser.inTrailer = parser
55+
h.handleRPCParser.inHeader = parser
5656
}
5757
}
5858

59-
// WithServerInHeaderParser adds a parser for the stats.InHeader to be able to access
60-
// the request payload
61-
func WithServerInHeaderParser(parser func(*stats.InHeader, zipkin.Span)) ServerOption {
59+
// WithServerOutPayloadParser adds a parser for the stats.OutPayload to be able to access
60+
// the outgoing response payload
61+
func WithServerOutPayloadParser(parser func(*stats.OutPayload, zipkin.SpanCustomizer)) ServerOption {
6262
return func(h *serverHandler) {
63-
h.handleRPCParser.inHeader = parser
63+
h.handleRPCParser.outPayload = parser
64+
}
65+
}
66+
67+
// WithServerOutTrailerParser adds a parser for the stats.OutTrailer to be able to access
68+
// the outgoing response trailer
69+
func WithServerOutTrailerParser(parser func(*stats.OutTrailer, zipkin.SpanCustomizer)) ServerOption {
70+
return func(h *serverHandler) {
71+
h.handleRPCParser.outTrailer = parser
72+
}
73+
}
74+
75+
// WithServerOutHeaderParser adds a parser for the stats.OutHeader to be able to access
76+
// the outgoing response payload
77+
func WithServerOutHeaderParser(parser func(*stats.OutHeader, zipkin.SpanCustomizer)) ServerOption {
78+
return func(h *serverHandler) {
79+
h.handleRPCParser.outHeader = parser
6480
}
6581
}
6682

‎middleware/grpc/server_parser_test.go‎

Lines changed: 35 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2019 The OpenZipkin Authors
1+
// Copyright 2020 The OpenZipkin Authors
22
//
33
// Licensed under the Apache License, Version 2.0 (the "License");
44
// you may not use this file except in compliance with the License.
@@ -77,30 +77,56 @@ func TestGRPCServerCreatesASpanAndContext(t *testing.T) {
7777
}
7878
}
7979

80-
func TestGRPCServerCanAccessToHeaders(t *testing.T) {
80+
func TestGRPCServerCanAccessToPayloadAndMetadata(t *testing.T) {
8181
tracer, flusher := createTracer(false)
8282

8383
s := grpc.NewServer(
8484
grpc.StatsHandler(
8585
zipkingrpc.NewServerHandler(
8686
tracer,
8787
zipkingrpc.ServerTags(map[string]string{"default": "tag"}),
88-
zipkingrpc.WithServerInHeaderParser(func(inHeader *stats.InHeader, span zipkin.Span) {
88+
zipkingrpc.WithServerInPayloadParser(func(inPayload *stats.InPayload, span zipkin.SpanCustomizer) {
89+
m, ok := inPayload.Payload.(*service.HelloRequest)
90+
if !ok {
91+
t.Fatal("failed to cast the payload as a service.HelloRequest")
92+
}
93+
if want, have := "Hello", m.Payload; want != have {
94+
t.Errorf("incorrect payload: want %q, have %q", want, have)
95+
}
96+
}),
97+
zipkingrpc.WithServerInHeaderParser(func(inHeader *stats.InHeader, span zipkin.SpanCustomizer) {
8998
if want, have := "test_value", inHeader.Header.Get("test_key")[0]; want != have {
90-
t.Errorf("unexpected metadata value in header, want: %q, have %q", want, have)
99+
t.Errorf("incorrect header value, want %q, have %q", want, have)
100+
}
101+
}),
102+
zipkingrpc.WithServerOutPayloadParser(func(outPayload *stats.OutPayload, span zipkin.SpanCustomizer) {
103+
m, ok := outPayload.Payload.(*service.HelloResponse)
104+
if !ok {
105+
t.Fatal("failed to cast the payload as a service.HelloResponse")
106+
}
107+
if want, have := "World", m.Payload; want != have {
108+
t.Errorf("incorrect payload: want %q, have %q", want, have)
109+
}
110+
}),
111+
zipkingrpc.WithServerOutHeaderParser(func(outHeader *stats.OutHeader, span zipkin.SpanCustomizer) {
112+
if want, have := "test_value_1", outHeader.Header.Get("test_key")[0]; want != have {
113+
t.Errorf("incorrect header value, want %q, have %q", want, have)
91114
}
92115
}),
93-
zipkingrpc.WithServerInTrailerParser(func(inTrailer *stats.InTrailer, span zipkin.Span) {
94-
if want, have := "test_value", inTrailer.Trailer.Get("test_key")[0]; want != have {
95-
t.Errorf("unexpected metadata value in header, want: %q, have %q", want, have)
116+
zipkingrpc.WithServerOutTrailerParser(func(outTrailer *stats.OutTrailer, span zipkin.SpanCustomizer) {
117+
if want, have := "test_value_2", outTrailer.Trailer.Get("test_key")[0]; want != have {
118+
t.Errorf("incorrect trailer value, want %q, have %q", want, have)
96119
}
97120
}),
98121
),
99122
),
100123
)
101124
defer s.Stop()
102125

103-
service.RegisterHelloServiceServer(s, &TestHelloService{})
126+
service.RegisterHelloServiceServer(s, &TestHelloService{
127+
responseHeader: metadata.Pairs("test_key", "test_value_1"),
128+
responseTrailer: metadata.Pairs("test_key", "test_value_2"),
129+
})
104130

105131
dialer := initListener(s)
106132

@@ -118,7 +144,7 @@ func TestGRPCServerCanAccessToHeaders(t *testing.T) {
118144

119145
client := service.NewHelloServiceClient(conn)
120146

121-
ctx = metadata.AppendToOutgoingContext(ctx, "test_key", "test_value")
147+
ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs("test_key", "test_value"))
122148
_, err = client.Hello(ctx, &service.HelloRequest{
123149
Payload: "Hello",
124150
})
@@ -130,9 +156,4 @@ func TestGRPCServerCanAccessToHeaders(t *testing.T) {
130156
if want, have := 1, len(spans); want != have {
131157
t.Errorf("unexpected number of spans, want %d, have %d", want, have)
132158
}
133-
134-
span := spans[0]
135-
if want, have := model.Server, span.Kind; want != have {
136-
t.Errorf("unexpected kind, want %q, have %q", want, have)
137-
}
138159
}

0 commit comments

Comments
 (0)