move opentelemetry into trace package, and refactoring (#996)

* move opentelemetry into trace package, and refactoring

* rename rewritten package names

core/trace/opentelemetry/agent.go

@@ -0,0 +1,68 @@
+package opentelemetry
+
+import (
+	"sync"
+
+	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/syncx"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/exporters/jaeger"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/sdk/resource"
+	tracesdk "go.opentelemetry.io/otel/sdk/trace"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
+)
+
+var (
+	once    sync.Once
+	enabled syncx.AtomicBool
+)
+
+// Enabled returns if prometheus is enabled.
+func Enabled() bool {
+	return enabled.True()
+}
+
+// StartAgent starts a prometheus agent.
+func StartAgent(c Config) {
+	once.Do(func() {
+		if len(c.Endpoint) == 0 {
+			return
+		}
+
+		// Just support jaeger now
+		if c.Batcher != "jaeger" {
+			return
+		}
+
+		exp, err := jaeger.New(jaeger.WithCollectorEndpoint(jaeger.WithEndpoint(c.Endpoint)))
+		if err != nil {
+			logx.Error(err)
+			return
+		}
+
+		tp := tracesdk.NewTracerProvider(
+			// Set the sampling rate based on the parent span to 100%
+			tracesdk.WithSampler(tracesdk.ParentBased(tracesdk.TraceIDRatioBased(c.Sampler))),
+			// Always be sure to batch in production.
+			tracesdk.WithBatcher(exp),
+			// Record information about this application in an Resource.
+			tracesdk.WithResource(resource.NewSchemaless(semconv.ServiceNameKey.String(c.Name))),
+		)
+
+		otel.SetTracerProvider(tp)
+		otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{}))
+		otel.SetErrorHandler(otelErrHandler{})
+
+		enabled.Set(true)
+	})
+}
+
+// errHandler handing otel errors.
+type otelErrHandler struct{}
+
+var _ otel.ErrorHandler = otelErrHandler{}
+
+func (o otelErrHandler) Handle(err error) {
+	logx.Errorf("[otel] error: %v", err)
+}

core/trace/opentelemetry/attributes.go

@@ -0,0 +1,40 @@
+package opentelemetry
+
+import (
+	"go.opentelemetry.io/otel/attribute"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
+	gcodes "google.golang.org/grpc/codes"
+)
+
+const (
+	// GRPCStatusCodeKey is convention for numeric status code of a gRPC request.
+	GRPCStatusCodeKey = attribute.Key("rpc.grpc.status_code")
+	// RPCNameKey is the name of message transmitted or received.
+	RPCNameKey = attribute.Key("name")
+	// RPCMessageTypeKey is the type of message transmitted or received.
+	RPCMessageTypeKey = attribute.Key("message.type")
+	// RPCMessageIDKey is the identifier of message transmitted or received.
+	RPCMessageIDKey = attribute.Key("message.id")
+	// RPCMessageCompressedSizeKey is the compressed size of the message transmitted or received in bytes.
+	RPCMessageCompressedSizeKey = attribute.Key("message.compressed_size")
+	// RPCMessageUncompressedSizeKey is the uncompressed size of the message
+	// transmitted or received in bytes.
+	RPCMessageUncompressedSizeKey = attribute.Key("message.uncompressed_size")
+)
+
+// Semantic conventions for common RPC attributes.
+var (
+	// RPCSystemGRPC is the semantic convention for gRPC as the remoting system.
+	RPCSystemGRPC = semconv.RPCSystemKey.String("grpc")
+	// RPCNameMessage is the semantic convention for a message named message.
+	RPCNameMessage = RPCNameKey.String("message")
+	// RPCMessageTypeSent is the semantic conventions for sent RPC message types.
+	RPCMessageTypeSent = RPCMessageTypeKey.String("SENT")
+	// RPCMessageTypeReceived is the semantic conventions for the received RPC message types.
+	RPCMessageTypeReceived = RPCMessageTypeKey.String("RECEIVED")
+)
+
+// StatusCodeAttr returns a attribute.KeyValue that represents the give c.
+func StatusCodeAttr(c gcodes.Code) attribute.KeyValue {
+	return GRPCStatusCodeKey.Int64(int64(c))
+}

core/trace/opentelemetry/clientstream.go

@@ -0,0 +1,121 @@
+package opentelemetry
+
+import (
+	"context"
+	"io"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+const (
+	receiveEndEvent streamEventType = iota
+	errorEvent
+)
+
+type (
+	streamEventType int
+
+	streamEvent struct {
+		Type streamEventType
+		Err  error
+	}
+
+	clientStream struct {
+		grpc.ClientStream
+		Finished          chan error
+		desc              *grpc.StreamDesc
+		events            chan streamEvent
+		eventsDone        chan struct{}
+		receivedMessageID int
+		sentMessageID     int
+	}
+)
+
+func (w *clientStream) RecvMsg(m interface{}) error {
+	err := w.ClientStream.RecvMsg(m)
+	if err == nil && !w.desc.ServerStreams {
+		w.sendStreamEvent(receiveEndEvent, nil)
+	} else if err == io.EOF {
+		w.sendStreamEvent(receiveEndEvent, nil)
+	} else if err != nil {
+		w.sendStreamEvent(errorEvent, err)
+	} else {
+		w.receivedMessageID++
+		MessageReceived.Event(w.Context(), w.receivedMessageID, m)
+	}
+
+	return err
+}
+
+func (w *clientStream) SendMsg(m interface{}) error {
+	err := w.ClientStream.SendMsg(m)
+	w.sentMessageID++
+	MessageSent.Event(w.Context(), w.sentMessageID, m)
+	if err != nil {
+		w.sendStreamEvent(errorEvent, err)
+	}
+
+	return err
+}
+
+func (w *clientStream) Header() (metadata.MD, error) {
+	md, err := w.ClientStream.Header()
+	if err != nil {
+		w.sendStreamEvent(errorEvent, err)
+	}
+
+	return md, err
+}
+
+func (w *clientStream) CloseSend() error {
+	err := w.ClientStream.CloseSend()
+	if err != nil {
+		w.sendStreamEvent(errorEvent, err)
+	}
+
+	return err
+}
+
+func (w *clientStream) sendStreamEvent(eventType streamEventType, err error) {
+	select {
+	case <-w.eventsDone:
+	case w.events <- streamEvent{Type: eventType, Err: err}:
+	}
+}
+
+// WrapClientStream wraps s with given ctx and desc.
+func WrapClientStream(ctx context.Context, s grpc.ClientStream, desc *grpc.StreamDesc) *clientStream {
+	events := make(chan streamEvent)
+	eventsDone := make(chan struct{})
+	finished := make(chan error)
+
+	go func() {
+		defer close(eventsDone)
+
+		for {
+			select {
+			case event := <-events:
+				switch event.Type {
+				case receiveEndEvent:
+					finished <- nil
+					return
+				case errorEvent:
+					finished <- event.Err
+					return
+				}
+			case <-ctx.Done():
+				finished <- ctx.Err()
+				return
+			}
+		}
+	}()
+
+	return &clientStream{
+		ClientStream: s,
+		desc:         desc,
+		events:       events,
+		eventsDone:   eventsDone,
+		Finished:     finished,
+	}
+}

core/trace/opentelemetry/config.go

@@ -0,0 +1,12 @@
+package opentelemetry
+
+// TraceName represents the tracing name.
+const TraceName = "go-zero"
+
+// A Config is a opentelemetry config.
+type Config struct {
+	Name     string  `json:",optional"`
+	Endpoint string  `json:",optional"`
+	Sampler  float64 `json:",default=1.0"`
+	Batcher  string  `json:",default=jaeger"`
+}

core/trace/opentelemetry/message.go

@@ -0,0 +1,38 @@
+package opentelemetry
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
+	"google.golang.org/protobuf/proto"
+)
+
+const messageEvent = "message"
+
+var (
+	// MessageSent is the type of sent messages.
+	MessageSent = messageType(RPCMessageTypeSent)
+	// MessageReceived is the type of received messages.
+	MessageReceived = messageType(RPCMessageTypeReceived)
+)
+
+type messageType attribute.KeyValue
+
+// Event adds an event of the messageType to the span associated with the
+// passed context with id and size (if message is a proto message).
+func (m messageType) Event(ctx context.Context, id int, message interface{}) {
+	span := trace.SpanFromContext(ctx)
+	if p, ok := message.(proto.Message); ok {
+		span.AddEvent(messageEvent, trace.WithAttributes(
+			attribute.KeyValue(m),
+			RPCMessageIDKey.Int(id),
+			RPCMessageUncompressedSizeKey.Int(proto.Size(p)),
+		))
+	} else {
+		span.AddEvent(messageEvent, trace.WithAttributes(
+			attribute.KeyValue(m),
+			RPCMessageIDKey.Int(id),
+		))
+	}
+}

core/trace/opentelemetry/serverstream.go

@@ -0,0 +1,47 @@
+package opentelemetry
+
+import (
+	"context"
+
+	"google.golang.org/grpc"
+)
+
+// serverStream wraps around the embedded grpc.ServerStream, and intercepts the RecvMsg and
+// SendMsg method call.
+type serverStream struct {
+	grpc.ServerStream
+	ctx context.Context
+
+	receivedMessageID int
+	sentMessageID     int
+}
+
+func (w *serverStream) Context() context.Context {
+	return w.ctx
+}
+
+func (w *serverStream) RecvMsg(m interface{}) error {
+	err := w.ServerStream.RecvMsg(m)
+	if err == nil {
+		w.receivedMessageID++
+		MessageReceived.Event(w.Context(), w.receivedMessageID, m)
+	}
+
+	return err
+}
+
+func (w *serverStream) SendMsg(m interface{}) error {
+	err := w.ServerStream.SendMsg(m)
+	w.sentMessageID++
+	MessageSent.Event(w.Context(), w.sentMessageID, m)
+
+	return err
+}
+
+// WrapServerStream wraps the given grpc.ServerStream with the given context.
+func WrapServerStream(ctx context.Context, ss grpc.ServerStream) *serverStream {
+	return &serverStream{
+		ServerStream: ss,
+		ctx:          ctx,
+	}
+}

core/trace/opentelemetry/tracer.go

@@ -0,0 +1,55 @@
+package opentelemetry
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel/baggage"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/trace"
+	"google.golang.org/grpc/metadata"
+)
+
+// assert that metadataSupplier implements the TextMapCarrier interface
+var _ propagation.TextMapCarrier = new(metadataSupplier)
+
+type metadataSupplier struct {
+	metadata *metadata.MD
+}
+
+func (s *metadataSupplier) Get(key string) string {
+	values := s.metadata.Get(key)
+	if len(values) == 0 {
+		return ""
+	}
+
+	return values[0]
+}
+
+func (s *metadataSupplier) Set(key, value string) {
+	s.metadata.Set(key, value)
+}
+
+func (s *metadataSupplier) Keys() []string {
+	out := make([]string, 0, len(*s.metadata))
+	for key := range *s.metadata {
+		out = append(out, key)
+	}
+
+	return out
+}
+
+// Inject injects the metadata into ctx.
+func Inject(ctx context.Context, p propagation.TextMapPropagator, metadata *metadata.MD) {
+	p.Inject(ctx, &metadataSupplier{
+		metadata: metadata,
+	})
+}
+
+// Extract extracts the metadata from ctx.
+func Extract(ctx context.Context, p propagation.TextMapPropagator, metadata *metadata.MD) (baggage.Baggage, trace.SpanContext) {
+	ctx = p.Extract(ctx, &metadataSupplier{
+		metadata: metadata,
+	})
+
+	return baggage.FromContext(ctx), trace.SpanContextFromContext(ctx)
+}

core/trace/opentelemetry/utils.go

@@ -0,0 +1,69 @@
+package opentelemetry
+
+import (
+	"context"
+	"net"
+	"strings"
+
+	"go.opentelemetry.io/otel/attribute"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
+	"google.golang.org/grpc/peer"
+)
+
+const localhost = "127.0.0.1"
+
+// PeerFromCtx returns the peer from ctx.
+func PeerFromCtx(ctx context.Context) string {
+	p, ok := peer.FromContext(ctx)
+	if !ok {
+		return ""
+	}
+
+	return p.Addr.String()
+}
+
+// SpanInfo returns the span info.
+func SpanInfo(fullMethod, peerAddress string) (string, []attribute.KeyValue) {
+	attrs := []attribute.KeyValue{RPCSystemGRPC}
+	name, mAttrs := ParseFullMethod(fullMethod)
+	attrs = append(attrs, mAttrs...)
+	attrs = append(attrs, PeerAttr(peerAddress)...)
+	return name, attrs
+}
+
+// ParseFullMethod returns the method name and attributes.
+func ParseFullMethod(fullMethod string) (string, []attribute.KeyValue) {
+	name := strings.TrimLeft(fullMethod, "/")
+	parts := strings.SplitN(name, "/", 2)
+	if len(parts) != 2 {
+		// Invalid format, does not follow `/package.service/method`.
+		return name, []attribute.KeyValue(nil)
+	}
+
+	var attrs []attribute.KeyValue
+	if service := parts[0]; service != "" {
+		attrs = append(attrs, semconv.RPCServiceKey.String(service))
+	}
+	if method := parts[1]; method != "" {
+		attrs = append(attrs, semconv.RPCMethodKey.String(method))
+	}
+
+	return name, attrs
+}
+
+// PeerAttr returns the peer attributes.
+func PeerAttr(addr string) []attribute.KeyValue {
+	host, port, err := net.SplitHostPort(addr)
+	if err != nil {
+		return []attribute.KeyValue(nil)
+	}
+
+	if len(host) == 0 {
+		host = localhost
+	}
+
+	return []attribute.KeyValue{
+		semconv.NetPeerIPKey.String(host),
+		semconv.NetPeerPortKey.String(port),
+	}
+}