github.com/thanos-io/thanos@v0.32.5/pkg/tracing/grpc.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package tracing
     5  
     6  import (
     7  	"context"
     8  
     9  	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
    10  	grpc_opentracing "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/tracing"
    11  	opentracing "github.com/opentracing/opentracing-go"
    12  	"google.golang.org/grpc"
    13  )
    14  
    15  // UnaryClientInterceptor returns a new unary client interceptor for OpenTracing.
    16  func UnaryClientInterceptor(tracer opentracing.Tracer) grpc.UnaryClientInterceptor {
    17  	return grpc_opentracing.UnaryClientInterceptor(grpc_opentracing.WithTracer(tracer))
    18  }
    19  
    20  // StreamClientInterceptor returns a new streaming client interceptor for OpenTracing.
    21  func StreamClientInterceptor(tracer opentracing.Tracer) grpc.StreamClientInterceptor {
    22  	return grpc_opentracing.StreamClientInterceptor(grpc_opentracing.WithTracer(tracer))
    23  }
    24  
    25  // UnaryServerInterceptor returns a new unary server interceptor for OpenTracing and injects given tracer.
    26  func UnaryServerInterceptor(tracer opentracing.Tracer) grpc.UnaryServerInterceptor {
    27  	interceptor := grpc_opentracing.UnaryServerInterceptor(grpc_opentracing.WithTracer(tracer))
    28  	return func(parentCtx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    29  		// Add our own tracer.
    30  		return interceptor(ContextWithTracer(parentCtx, tracer), req, info, handler)
    31  	}
    32  }
    33  
    34  // StreamServerInterceptor returns a new streaming server interceptor for OpenTracing and injects given tracer.
    35  func StreamServerInterceptor(tracer opentracing.Tracer) grpc.StreamServerInterceptor {
    36  	interceptor := grpc_opentracing.StreamServerInterceptor(grpc_opentracing.WithTracer(tracer))
    37  	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    38  		// Add our own tracer.
    39  		wrappedStream := grpc_middleware.WrapServerStream(stream)
    40  		wrappedStream.WrappedContext = ContextWithTracer(stream.Context(), tracer)
    41  
    42  		return interceptor(srv, wrappedStream, info, handler)
    43  	}
    44  }