gitlab.com/gitlab-org/labkit@v1.21.0/correlation/grpc/client_interceptors.go (about)

     1  package grpccorrelation
     2  
     3  import (
     4  	"context"
     5  
     6  	"gitlab.com/gitlab-org/labkit/correlation"
     7  	"google.golang.org/grpc"
     8  	"google.golang.org/grpc/metadata"
     9  )
    10  
    11  func appendToOutgoingContext(ctx context.Context, clientName string) context.Context {
    12  	correlationID := correlation.ExtractFromContext(ctx)
    13  	if correlationID != "" {
    14  		ctx = metadata.AppendToOutgoingContext(ctx, metadataCorrelatorKey, correlationID)
    15  	}
    16  	if clientName != "" {
    17  		ctx = metadata.AppendToOutgoingContext(ctx, metadataClientNameKey, clientName)
    18  	}
    19  
    20  	return ctx
    21  }
    22  
    23  // UnaryClientCorrelationInterceptor propagates Correlation-IDs downstream.
    24  func UnaryClientCorrelationInterceptor(opts ...ClientCorrelationInterceptorOption) grpc.UnaryClientInterceptor {
    25  	config := applyClientCorrelationInterceptorOptions(opts)
    26  
    27  	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    28  		ctx = appendToOutgoingContext(ctx, config.clientName)
    29  		return invoker(ctx, method, req, reply, cc, opts...)
    30  	}
    31  }
    32  
    33  // StreamClientCorrelationInterceptor propagates Correlation-IDs downstream.
    34  func StreamClientCorrelationInterceptor(opts ...ClientCorrelationInterceptorOption) grpc.StreamClientInterceptor {
    35  	config := applyClientCorrelationInterceptorOptions(opts)
    36  
    37  	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn,
    38  		method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
    39  		ctx = appendToOutgoingContext(ctx, config.clientName)
    40  		return streamer(ctx, desc, cc, method, opts...)
    41  	}
    42  }
    43  
    44  // InjectToOutgoingContext will inject the correlation ID into the
    45  // outgoing context metadata. Repeat calls will overwrite any existing
    46  // correlation IDs.
    47  func InjectToOutgoingContext(ctx context.Context, correlationID string) context.Context {
    48  	return metadata.AppendToOutgoingContext(ctx, metadataCorrelatorKey, correlationID)
    49  }