zeromicro/go-zero
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactoring tracing interceptors. (#1009)

Browse Source 
* refactor tracing interceptors

* add stream tracing interceptor
This commit is contained in:
Kevin Wan
2021-09-07 17:58:22 +08:00
committed by GitHub
parent 96a35ecf1a
commit c837dc21bb
9 changed files with 160 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
zrpc/internal/clientinterceptors/opentracinginterceptor.go
View File

@@ -13,8 +13,8 @@ import (
"google.golang.org/grpc/status"
)
// OpenTracingInterceptor returns a grpc.UnaryClientInterceptor for opentelemetry.
func OpenTracingInterceptor() grpc.UnaryClientInterceptor {
// UnaryOpenTracingInterceptor returns a grpc.UnaryClientInterceptor for opentelemetry.
func UnaryOpenTracingInterceptor() grpc.UnaryClientInterceptor {
propagator := otel.GetTextMapPropagator()
return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {

2
zrpc/internal/clientinterceptors/opentracinginterceptor_test.go
View File

@@ -18,7 +18,7 @@ func TestOpenTracingInterceptor(t *testing.T) {
})
cc := new(grpc.ClientConn)
err := OpenTracingInterceptor()(context.Background(), "/ListUser", nil, nil, cc,
err := UnaryOpenTracingInterceptor()(context.Background(), "/ListUser", nil, nil, cc,
func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
opts ...grpc.CallOption) error {
return nil

20
zrpc/internal/clientinterceptors/tracinginterceptor.go
View File

@@ -8,8 +8,8 @@ import (
"google.golang.org/grpc/metadata"
)
// TracingInterceptor is an interceptor that handles tracing.
func TracingInterceptor(ctx context.Context, method string, req, reply interface{},
// UnaryTracingInterceptor is an interceptor that handles tracing.
func UnaryTracingInterceptor(ctx context.Context, method string, req, reply interface{},
cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
ctx, span := trace.StartClientSpan(ctx, cc.Target(), method)
defer span.Finish()
@@ -23,3 +23,19 @@ func TracingInterceptor(ctx context.Context, method string, req, reply interface
return invoker(ctx, method, req, reply, cc, opts...)
}
// StreamTracingInterceptor is an interceptor that handles tracing for stream requests.
func StreamTracingInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn,
method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
ctx, span := trace.StartClientSpan(ctx, cc.Target(), method)
defer span.Finish()
var pairs []string
span.Visit(func(key, val string) bool {
pairs = append(pairs, key, val)
return true
})
ctx = metadata.AppendToOutgoingContext(ctx, pairs...)
return streamer(ctx, desc, cc, method, opts...)
}

48
zrpc/internal/clientinterceptors/tracinginterceptor_test.go
View File

@@ -12,12 +12,12 @@ import (
"google.golang.org/grpc/metadata"
)
func TestTracingInterceptor(t *testing.T) {
func TestUnaryTracingInterceptor(t *testing.T) {
var run int32
var wg sync.WaitGroup
wg.Add(1)
cc := new(grpc.ClientConn)
err := TracingInterceptor(context.Background(), "/foo", nil, nil, cc,
err := UnaryTracingInterceptor(context.Background(), "/foo", nil, nil, cc,
func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
opts ...grpc.CallOption) error {
defer wg.Done()
@@ -29,7 +29,24 @@ func TestTracingInterceptor(t *testing.T) {
assert.Equal(t, int32(1), atomic.LoadInt32(&run))
}
func TestTracingInterceptor_GrpcFormat(t *testing.T) {
func TestStreamTracingInterceptor(t *testing.T) {
var run int32
var wg sync.WaitGroup
wg.Add(1)
cc := new(grpc.ClientConn)
_, err := StreamTracingInterceptor(context.Background(), nil, cc, "/foo",
func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string,
opts ...grpc.CallOption) (grpc.ClientStream, error) {
defer wg.Done()
atomic.AddInt32(&run, 1)
return nil, nil
})
wg.Wait()
assert.Nil(t, err)
assert.Equal(t, int32(1), atomic.LoadInt32(&run))
}
func TestUnaryTracingInterceptor_GrpcFormat(t *testing.T) {
var run int32
var wg sync.WaitGroup
wg.Add(1)
@@ -40,7 +57,7 @@ func TestTracingInterceptor_GrpcFormat(t *testing.T) {
assert.Nil(t, err)
ctx, _ := trace.StartServerSpan(context.Background(), carrier, "user", "/foo")
cc := new(grpc.ClientConn)
err = TracingInterceptor(ctx, "/foo", nil, nil, cc,
err = UnaryTracingInterceptor(ctx, "/foo", nil, nil, cc,
func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
opts ...grpc.CallOption) error {
defer wg.Done()
@@ -51,3 +68,26 @@ func TestTracingInterceptor_GrpcFormat(t *testing.T) {
assert.Nil(t, err)
assert.Equal(t, int32(1), atomic.LoadInt32(&run))
}
func TestStreamTracingInterceptor_GrpcFormat(t *testing.T) {
var run int32
var wg sync.WaitGroup
wg.Add(1)
md := metadata.New(map[string]string{
"foo": "bar",
})
carrier, err := trace.Inject(trace.GrpcFormat, md)
assert.Nil(t, err)
ctx, _ := trace.StartServerSpan(context.Background(), carrier, "user", "/foo")
cc := new(grpc.ClientConn)
_, err = StreamTracingInterceptor(ctx, nil, cc, "/foo",
func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string,
opts ...grpc.CallOption) (grpc.ClientStream, error) {
defer wg.Done()
atomic.AddInt32(&run, 1)
return nil, nil
})
wg.Wait()
assert.Nil(t, err)
assert.Equal(t, int32(1), atomic.LoadInt32(&run))
}