refactoring tracing interceptors. (#1009)

* refactor tracing interceptors

* add stream tracing interceptor

zrpc/internal/serverinterceptors/tracinginterceptor.go

@@ -27,3 +27,23 @@ func UnaryTracingInterceptor(serviceName string) grpc.UnaryServerInterceptor {
 		return handler(ctx, req)
 	}
 }
+
+func StreamTracingInterceptor(serviceName string) grpc.StreamServerInterceptor {
+	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo,
+		handler grpc.StreamHandler) error {
+		ctx := ss.Context()
+		md, ok := metadata.FromIncomingContext(ctx)
+		if !ok {
+			return handler(srv, ss)
+		}
+
+		carrier, err := trace.Extract(trace.GrpcFormat, md)
+		if err != nil {
+			return handler(srv, ss)
+		}
+
+		ctx, span := trace.StartServerSpan(ctx, carrier, serviceName, info.FullMethod)
+		defer span.Finish()
+		return handler(srv, ss)
+	}
+}

zrpc/internal/serverinterceptors/tracinginterceptor_test.go

@@ -46,3 +46,71 @@ func TestUnaryTracingInterceptor_GrpcFormat(t *testing.T) {
 	wg.Wait()
 	assert.Nil(t, err)
 }
+
+func TestStreamTracingInterceptor(t *testing.T) {
+	interceptor := StreamTracingInterceptor("foo")
+	var run int32
+	var wg sync.WaitGroup
+	wg.Add(1)
+	err := interceptor(nil, new(mockedServerStream), nil,
+		func(srv interface{}, stream grpc.ServerStream) error {
+			defer wg.Done()
+			atomic.AddInt32(&run, 1)
+			return nil
+		})
+	wg.Wait()
+	assert.Nil(t, err)
+	assert.Equal(t, int32(1), atomic.LoadInt32(&run))
+}
+
+func TestStreamTracingInterceptor_GrpcFormat(t *testing.T) {
+	interceptor := StreamTracingInterceptor("foo")
+	var run int32
+	var wg sync.WaitGroup
+	wg.Add(1)
+	var md metadata.MD
+	ctx := metadata.NewIncomingContext(context.Background(), md)
+	stream := mockedServerStream{ctx: ctx}
+	err := interceptor(nil, &stream, &grpc.StreamServerInfo{
+		FullMethod: "/foo",
+	}, func(srv interface{}, stream grpc.ServerStream) error {
+		defer wg.Done()
+		atomic.AddInt32(&run, 1)
+		return nil
+	})
+	wg.Wait()
+	assert.Nil(t, err)
+	assert.Equal(t, int32(1), atomic.LoadInt32(&run))
+}
+
+type mockedServerStream struct {
+	ctx context.Context
+}
+
+func (m *mockedServerStream) SetHeader(md metadata.MD) error {
+	panic("implement me")
+}
+
+func (m *mockedServerStream) SendHeader(md metadata.MD) error {
+	panic("implement me")
+}
+
+func (m *mockedServerStream) SetTrailer(md metadata.MD) {
+	panic("implement me")
+}
+
+func (m *mockedServerStream) Context() context.Context {
+	if m.ctx == nil {
+		return context.Background()
+	}
+
+	return m.ctx
+}
+
+func (m *mockedServerStream) SendMsg(v interface{}) error {
+	panic("implement me")
+}
+
+func (m *mockedServerStream) RecvMsg(v interface{}) error {
+	panic("implement me")
+}