github.com/erda-project/erda-infra@v1.0.9/pkg/trace/inject/context/grpc/grpc_context.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package grpccontext
    16  
    17  import (
    18  	"context"
    19  
    20  	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
    21  	"google.golang.org/grpc"
    22  
    23  	injectcontext "github.com/erda-project/erda-infra/pkg/trace/inject/context"
    24  )
    25  
    26  // UnaryServerInterceptor .
    27  func UnaryServerInterceptor() grpc.UnaryServerInterceptor {
    28  	inter := otelgrpc.UnaryServerInterceptor()
    29  	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
    30  		return inter(ctx, req, info, func(ctx context.Context, req interface{}) (interface{}, error) {
    31  			injectcontext.SetContext(ctx)
    32  			defer injectcontext.ClearContext()
    33  			return handler(ctx, req)
    34  		})
    35  	}
    36  }
    37  
    38  // StreamServerInterceptor .
    39  func StreamServerInterceptor() grpc.StreamServerInterceptor {
    40  	inter := otelgrpc.StreamServerInterceptor()
    41  	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    42  		return inter(srv, ss, info, func(srv interface{}, stream grpc.ServerStream) error {
    43  			injectcontext.SetContext(stream.Context())
    44  			defer injectcontext.ClearContext()
    45  			return handler(srv, stream)
    46  		})
    47  	}
    48  }
    49  
    50  // UnaryClientInterceptor .
    51  func UnaryClientInterceptor() grpc.UnaryClientInterceptor {
    52  	inter := otelgrpc.UnaryClientInterceptor()
    53  	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    54  		ctx = injectcontext.ContextWithSpan(ctx)
    55  		return inter(ctx, method, req, reply, cc, invoker, opts...)
    56  	}
    57  }
    58  
    59  // StreamClientInterceptor .
    60  func StreamClientInterceptor() grpc.StreamClientInterceptor {
    61  	inter := otelgrpc.StreamClientInterceptor()
    62  	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
    63  		ctx = injectcontext.ContextWithSpan(ctx)
    64  		return inter(ctx, desc, cc, method, streamer, opts...)
    65  	}
    66  }