add the opentelemetry tracing (#908)

* add the opentelemetry tracing

* fix the error sampler config

* 添加stream的链路跟踪

* fix the error field name

zrpc/internal/client.go

@@ -72,6 +72,10 @@ func (c *client) buildDialOptions(opts ...ClientOption) []grpc.DialOption {
 			clientinterceptors.PrometheusInterceptor,
 			clientinterceptors.BreakerInterceptor,
 			clientinterceptors.TimeoutInterceptor(cliOpts.Timeout),
+			clientinterceptors.OpenTracingInterceptor(),
+		),
+		WithStreamClientInterceptors(
+			clientinterceptors.StreamOpenTracingInterceptor(),
 		),
 	}
 

zrpc/internal/clientinterceptors/opentracinginterceptor.go

@@ -0,0 +1,108 @@
+package clientinterceptors
+
+import (
+	"context"
+
+	"github.com/tal-tech/go-zero/core/opentelemetry"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/trace"
+	"google.golang.org/grpc"
+	grpc_codes "google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+func OpenTracingInterceptor() 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 {
+		if !opentelemetry.Enabled() {
+			return invoker(ctx, method, req, reply, cc, opts...)
+		}
+
+		requestMetadata, _ := metadata.FromOutgoingContext(ctx)
+		metadataCopy := requestMetadata.Copy()
+
+		tr := otel.Tracer(opentelemetry.TraceName)
+		name, attr := opentelemetry.SpanInfo(method, cc.Target())
+
+		var span trace.Span
+		ctx, span = tr.Start(ctx,
+			name,
+			trace.WithSpanKind(trace.SpanKindClient),
+			trace.WithAttributes(attr...),
+		)
+		defer span.End()
+
+		opentelemetry.Inject(ctx, propagator, &metadataCopy)
+		ctx = metadata.NewOutgoingContext(ctx, metadataCopy)
+
+		opentelemetry.MessageSent.Event(ctx, 1, req)
+
+		err := invoker(ctx, method, req, reply, cc, opts...)
+
+		opentelemetry.MessageReceived.Event(ctx, 1, reply)
+
+		if err != nil {
+			s, _ := status.FromError(err)
+			span.SetStatus(codes.Error, s.Message())
+			span.SetAttributes(opentelemetry.StatusCodeAttr(s.Code()))
+		} else {
+			span.SetAttributes(opentelemetry.StatusCodeAttr(grpc_codes.OK))
+		}
+
+		return err
+	}
+}
+
+func StreamOpenTracingInterceptor() grpc.StreamClientInterceptor {
+	propagator := otel.GetTextMapPropagator()
+	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+		if !opentelemetry.Enabled() {
+			return streamer(ctx, desc, cc, method, opts...)
+		}
+
+		requestMetadata, _ := metadata.FromOutgoingContext(ctx)
+		metadataCopy := requestMetadata.Copy()
+
+		tr := otel.Tracer("ecoplants")
+
+		name, attr := opentelemetry.SpanInfo(method, cc.Target())
+		var span trace.Span
+		ctx, span = tr.Start(
+			ctx,
+			name,
+			trace.WithSpanKind(trace.SpanKindClient),
+			trace.WithAttributes(attr...),
+		)
+
+		opentelemetry.Inject(ctx, propagator, &metadataCopy)
+		ctx = metadata.NewOutgoingContext(ctx, metadataCopy)
+
+		s, err := streamer(ctx, desc, cc, method, opts...)
+		if err != nil {
+			grpcStatus, _ := status.FromError(err)
+			span.SetStatus(codes.Error, grpcStatus.Message())
+			span.SetAttributes(opentelemetry.StatusCodeAttr(grpcStatus.Code()))
+			span.End()
+			return s, err
+		}
+		stream := opentelemetry.WrapClientStream(ctx, s, desc)
+
+		go func() {
+			err := <-stream.Finished
+
+			if err != nil {
+				s, _ := status.FromError(err)
+				span.SetStatus(codes.Error, s.Message())
+				span.SetAttributes(opentelemetry.StatusCodeAttr(s.Code()))
+			} else {
+				span.SetAttributes(opentelemetry.StatusCodeAttr(grpc_codes.OK))
+			}
+
+			span.End()
+		}()
+
+		return stream, nil
+	}
+}

zrpc/internal/clientinterceptors/opentracinginterceptor_test.go

@@ -0,0 +1,27 @@
+package clientinterceptors
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/opentelemetry"
+	"google.golang.org/grpc"
+)
+
+func TestOpenTracingInterceptor(t *testing.T) {
+	opentelemetry.StartAgent(opentelemetry.Config{
+		Name:     "go-zero-test",
+		Endpoint: "http://localhost:14268/api/traces",
+		Batcher:  "jaeger",
+		Sampler:  1.0,
+	})
+
+	cc := new(grpc.ClientConn)
+	err := OpenTracingInterceptor()(context.Background(), "/ListUser", nil, nil, cc,
+		func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
+			opts ...grpc.CallOption) error {
+			return nil
+		})
+	assert.Nil(t, err)
+}

zrpc/internal/rpcserver.go

@@ -59,11 +59,13 @@ func (s *rpcServer) Start(register RegisterFn) error {
 		serverinterceptors.UnaryStatInterceptor(s.metrics),
 		serverinterceptors.UnaryPrometheusInterceptor(),
 		serverinterceptors.UnaryBreakerInterceptor(),
+		serverinterceptors.UnaryOpenTracingInterceptor(),
 	}
 	unaryInterceptors = append(unaryInterceptors, s.unaryInterceptors...)
 	streamInterceptors := []grpc.StreamServerInterceptor{
 		serverinterceptors.StreamCrashInterceptor,
 		serverinterceptors.StreamBreakerInterceptor,
+		serverinterceptors.StreamOpenTracingInterceptor(),
 	}
 	streamInterceptors = append(streamInterceptors, s.streamInterceptors...)
 	options := append(s.options, WithUnaryServerInterceptors(unaryInterceptors...),

zrpc/internal/serverinterceptors/opentracinginterceptor.go

@@ -0,0 +1,96 @@
+package serverinterceptors
+
+import (
+	"context"
+
+	"github.com/tal-tech/go-zero/core/opentelemetry"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/baggage"
+	"go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/trace"
+	"google.golang.org/grpc"
+	grpc_codes "google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+func UnaryOpenTracingInterceptor() grpc.UnaryServerInterceptor {
+	propagator := otel.GetTextMapPropagator()
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+		if !opentelemetry.Enabled() {
+			return handler(ctx, req)
+		}
+
+		requestMetadata, _ := metadata.FromIncomingContext(ctx)
+		metadataCopy := requestMetadata.Copy()
+
+		bags, spanCtx := opentelemetry.Extract(ctx, propagator, &metadataCopy)
+		ctx = baggage.ContextWithBaggage(ctx, bags)
+
+		tr := otel.Tracer(opentelemetry.TraceName)
+		name, attr := opentelemetry.SpanInfo(info.FullMethod, opentelemetry.PeerFromCtx(ctx))
+
+		var span trace.Span
+		ctx, span = tr.Start(
+			trace.ContextWithRemoteSpanContext(ctx, spanCtx),
+			name,
+			trace.WithSpanKind(trace.SpanKindServer),
+			trace.WithAttributes(attr...),
+		)
+		defer span.End()
+
+		opentelemetry.MessageReceived.Event(ctx, 1, req)
+
+		resp, err := handler(ctx, req)
+
+		if err != nil {
+			s, _ := status.FromError(err)
+			span.SetStatus(codes.Error, s.Message())
+			span.SetAttributes(opentelemetry.StatusCodeAttr(s.Code()))
+			opentelemetry.MessageSent.Event(ctx, 1, s.Proto())
+		} else {
+			span.SetAttributes(opentelemetry.StatusCodeAttr(grpc_codes.OK))
+			opentelemetry.MessageSent.Event(ctx, 1, resp)
+		}
+
+		return resp, err
+	}
+}
+
+func StreamOpenTracingInterceptor() grpc.StreamServerInterceptor {
+	propagator := otel.GetTextMapPropagator()
+	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		ctx := ss.Context()
+		if !opentelemetry.Enabled() {
+			return handler(srv, opentelemetry.WrapServerStream(ctx, ss))
+		}
+
+		requestMetadata, _ := metadata.FromIncomingContext(ctx)
+		metadataCopy := requestMetadata.Copy()
+
+		bags, spanCtx := opentelemetry.Extract(ctx, propagator, &metadataCopy)
+		ctx = baggage.ContextWithBaggage(ctx, bags)
+
+		tr := otel.Tracer(opentelemetry.TraceName)
+		name, attr := opentelemetry.SpanInfo(info.FullMethod, opentelemetry.PeerFromCtx(ctx))
+		ctx, span := tr.Start(
+			trace.ContextWithRemoteSpanContext(ctx, spanCtx),
+			name,
+			trace.WithSpanKind(trace.SpanKindServer),
+			trace.WithAttributes(attr...),
+		)
+		defer span.End()
+
+		err := handler(srv, opentelemetry.WrapServerStream(ctx, ss))
+
+		if err != nil {
+			s, _ := status.FromError(err)
+			span.SetStatus(codes.Error, s.Message())
+			span.SetAttributes(opentelemetry.StatusCodeAttr(s.Code()))
+		} else {
+			span.SetAttributes(opentelemetry.StatusCodeAttr(grpc_codes.OK))
+		}
+
+		return err
+	}
+}

zrpc/internal/serverinterceptors/opentracinginterceptor_test.go

@@ -0,0 +1,36 @@
+package serverinterceptors
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/opentelemetry"
+	"google.golang.org/grpc"
+)
+
+func TestUnaryOpenTracingInterceptor_Disable(t *testing.T) {
+	interceptor := UnaryOpenTracingInterceptor()
+	_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+		FullMethod: "/",
+	}, func(ctx context.Context, req interface{}) (interface{}, error) {
+		return nil, nil
+	})
+	assert.Nil(t, err)
+}
+
+func TestUnaryOpenTracingInterceptor_Enabled(t *testing.T) {
+	opentelemetry.StartAgent(opentelemetry.Config{
+		Name:     "go-zero-test",
+		Endpoint: "http://localhost:14268/api/traces",
+		Batcher:  "jaeger",
+		Sampler:  1.0,
+	})
+	interceptor := UnaryOpenTracingInterceptor()
+	_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+		FullMethod: "/package.TestService.GetUser",
+	}, func(ctx context.Context, req interface{}) (interface{}, error) {
+		return nil, nil
+	})
+	assert.Nil(t, err)
+}