github.com/microsoft/moc@v0.17.1/pkg/diagnostics/context.go (about)

     1  // Copyright (c) Microsoft Corporation.
     2  // Licensed under the Apache v2.0 license.
     3  
     4  package diagnostics
     5  
     6  import (
     7  	"context"
     8  
     9  	"google.golang.org/grpc/metadata"
    10  )
    11  
    12  type context_key int
    13  
    14  const CorrelationIdcontext_key context_key = 0
    15  const CorrelationIdMetadataKey string = "correlation-id-4b54b6c9-647a-4929-be87-481ba63fc04d"
    16  
    17  // Replace the context by using this method before making gRPC calls to MOC cloud/node agent, the
    18  // specified correlation-id will also be passed to server.
    19  func NewContextWithCorrelationId(parent context.Context, correlationId string) context.Context {
    20  	ctx := parent
    21  	if ctx == nil {
    22  		ctx = context.Background()
    23  	}
    24  
    25  	// Adds to outgoing context
    26  	md, ok := metadata.FromOutgoingContext(ctx)
    27  	if ok {
    28  		md.Set(CorrelationIdMetadataKey, correlationId)
    29  	} else {
    30  		md = metadata.Pairs(CorrelationIdMetadataKey, correlationId)
    31  	}
    32  	ctx = metadata.NewOutgoingContext(ctx, md)
    33  
    34  	// Adds to runtime context
    35  	return context.WithValue(ctx, CorrelationIdcontext_key, correlationId)
    36  }
    37  
    38  // For server-side MOC agent to pick the correlation-id which is previously set by gRPC client by
    39  // using the NewContextWithCorrelationId function
    40  func GetCorrelationIdFromIncomingContext(ctx context.Context) string {
    41  	md, ok := metadata.FromIncomingContext(ctx)
    42  	if ok {
    43  		values := md.Get(CorrelationIdMetadataKey)
    44  		if len(values) > 0 {
    45  			return values[0]
    46  		}
    47  	}
    48  	return ""
    49  }
    50  
    51  // Fetch the correlation-id from context, that was either previously set in the current call-stack
    52  // by using the NewContextWithCorrelationId function, or set by the client before making gPRC call.
    53  func GetCorrelationIdFromContext(ctx context.Context) string {
    54  	value, ok := ctx.Value(CorrelationIdcontext_key).(string)
    55  	if ok {
    56  		return value
    57  	}
    58  	return GetCorrelationIdFromIncomingContext(ctx)
    59  }