github.com/lingyao2333/mo-zero@v1.4.1/zrpc/internal/serverinterceptors/tracinginterceptor.go (about)

     1  package serverinterceptors
     2  
     3  import (
     4  	"context"
     5  
     6  	ztrace "github.com/lingyao2333/mo-zero/core/trace"
     7  	"go.opentelemetry.io/otel"
     8  	"go.opentelemetry.io/otel/baggage"
     9  	"go.opentelemetry.io/otel/codes"
    10  	"go.opentelemetry.io/otel/trace"
    11  	"google.golang.org/grpc"
    12  	gcodes "google.golang.org/grpc/codes"
    13  	"google.golang.org/grpc/metadata"
    14  	"google.golang.org/grpc/status"
    15  )
    16  
    17  // UnaryTracingInterceptor is a grpc.UnaryServerInterceptor for opentelemetry.
    18  func UnaryTracingInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
    19  	handler grpc.UnaryHandler) (interface{}, error) {
    20  	ctx, span := startSpan(ctx, info.FullMethod)
    21  	defer span.End()
    22  
    23  	ztrace.MessageReceived.Event(ctx, 1, req)
    24  	resp, err := handler(ctx, req)
    25  	if err != nil {
    26  		s, ok := status.FromError(err)
    27  		if ok {
    28  			span.SetStatus(codes.Error, s.Message())
    29  			span.SetAttributes(ztrace.StatusCodeAttr(s.Code()))
    30  			ztrace.MessageSent.Event(ctx, 1, s.Proto())
    31  		} else {
    32  			span.SetStatus(codes.Error, err.Error())
    33  		}
    34  		return nil, err
    35  	}
    36  
    37  	span.SetAttributes(ztrace.StatusCodeAttr(gcodes.OK))
    38  	ztrace.MessageSent.Event(ctx, 1, resp)
    39  
    40  	return resp, nil
    41  }
    42  
    43  // StreamTracingInterceptor returns a grpc.StreamServerInterceptor for opentelemetry.
    44  func StreamTracingInterceptor(svr interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo,
    45  	handler grpc.StreamHandler) error {
    46  	ctx, span := startSpan(ss.Context(), info.FullMethod)
    47  	defer span.End()
    48  
    49  	if err := handler(svr, wrapServerStream(ctx, ss)); err != nil {
    50  		s, ok := status.FromError(err)
    51  		if ok {
    52  			span.SetStatus(codes.Error, s.Message())
    53  			span.SetAttributes(ztrace.StatusCodeAttr(s.Code()))
    54  		} else {
    55  			span.SetStatus(codes.Error, err.Error())
    56  		}
    57  		return err
    58  	}
    59  
    60  	span.SetAttributes(ztrace.StatusCodeAttr(gcodes.OK))
    61  	return nil
    62  }
    63  
    64  // serverStream wraps around the embedded grpc.ServerStream,
    65  // and intercepts the RecvMsg and SendMsg method call.
    66  type serverStream struct {
    67  	grpc.ServerStream
    68  	ctx               context.Context
    69  	receivedMessageID int
    70  	sentMessageID     int
    71  }
    72  
    73  func (w *serverStream) Context() context.Context {
    74  	return w.ctx
    75  }
    76  
    77  func (w *serverStream) RecvMsg(m interface{}) error {
    78  	err := w.ServerStream.RecvMsg(m)
    79  	if err == nil {
    80  		w.receivedMessageID++
    81  		ztrace.MessageReceived.Event(w.Context(), w.receivedMessageID, m)
    82  	}
    83  
    84  	return err
    85  }
    86  
    87  func (w *serverStream) SendMsg(m interface{}) error {
    88  	err := w.ServerStream.SendMsg(m)
    89  	w.sentMessageID++
    90  	ztrace.MessageSent.Event(w.Context(), w.sentMessageID, m)
    91  
    92  	return err
    93  }
    94  
    95  func startSpan(ctx context.Context, method string) (context.Context, trace.Span) {
    96  	md, ok := metadata.FromIncomingContext(ctx)
    97  	if !ok {
    98  		md = metadata.MD{}
    99  	}
   100  	bags, spanCtx := ztrace.Extract(ctx, otel.GetTextMapPropagator(), &md)
   101  	ctx = baggage.ContextWithBaggage(ctx, bags)
   102  	tr := otel.Tracer(ztrace.TraceName)
   103  	name, attr := ztrace.SpanInfo(method, ztrace.PeerFromCtx(ctx))
   104  
   105  	return tr.Start(trace.ContextWithRemoteSpanContext(ctx, spanCtx), name,
   106  		trace.WithSpanKind(trace.SpanKindServer), trace.WithAttributes(attr...))
   107  }
   108  
   109  // wrapServerStream wraps the given grpc.ServerStream with the given context.
   110  func wrapServerStream(ctx context.Context, ss grpc.ServerStream) *serverStream {
   111  	return &serverStream{
   112  		ServerStream: ss,
   113  		ctx:          ctx,
   114  	}
   115  }