github.com/songzhibin97/gkit@v1.2.13/trace/middle.go (about)

     1  package trace
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/songzhibin97/gkit/middleware"
     7  	"github.com/songzhibin97/gkit/options"
     8  	"go.opentelemetry.io/otel/trace"
     9  )
    10  
    11  // WithServer returns a new server middleware for OpenTelemetry.
    12  func WithServer(opts ...options.Option) middleware.MiddleWare {
    13  	tracer := NewTracer(trace.SpanKindServer, opts...)
    14  	return func(handler middleware.Endpoint) middleware.Endpoint {
    15  		return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
    16  			if tr, ok := FromServerTransportContext(ctx); ok {
    17  				var span trace.Span
    18  				ctx, span = tracer.Start(ctx, tr.Operation(), tr.RequestHeader())
    19  				setServerSpan(ctx, span, req)
    20  				defer func() { tracer.End(ctx, span, reply, err) }()
    21  			}
    22  			return handler(ctx, req)
    23  		}
    24  	}
    25  }
    26  
    27  // WithClient returns a new client middleware for OpenTelemetry.
    28  func WithClient(opts ...options.Option) middleware.MiddleWare {
    29  	tracer := NewTracer(trace.SpanKindClient, opts...)
    30  	return func(handler middleware.Endpoint) middleware.Endpoint {
    31  		return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
    32  			if tr, ok := FromClientTransportContext(ctx); ok {
    33  				var span trace.Span
    34  				ctx, span = tracer.Start(ctx, tr.Operation(), tr.RequestHeader())
    35  				setClientSpan(ctx, span, req)
    36  				defer func() { tracer.End(ctx, span, reply, err) }()
    37  			}
    38  			return handler(ctx, req)
    39  		}
    40  	}
    41  }