github.com/brownsys/tracing-framework-go@v0.0.0-20161210174012-0542a62412fe/xtrace/grpc/grpc.go (about)

     1  // Package grpc provides wrappers around certain calls
     2  // in the standard grpc package that automate the propagation
     3  // of X-Trace Task and Event IDs.
     4  package grpc
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"strconv"
    10  
    11  	"github.com/brownsys/tracing-framework-go/xtrace/client"
    12  	"google.golang.org/grpc"
    13  	"google.golang.org/grpc/metadata"
    14  )
    15  
    16  var (
    17  	// EventIDMetadataKey is the key used to store an
    18  	// X-Trace Event ID in gRPC metadata.
    19  	EventIDMetadataKey = "xtrace_event_id"
    20  	// TaskIDMetadataKey is the key used to store an
    21  	// X-Trace Task ID in gRPC metadata.
    22  	TaskIDMetadataKey = "xtrace_task_id"
    23  )
    24  
    25  // Invoke wraps grpc's Invoke function, adding the calling goroutine's
    26  // current Task and Event IDs if possible.
    27  func Invoke(ctx context.Context, method string, args, reply interface{},
    28  	cc *grpc.ClientConn, opts ...grpc.CallOption) error {
    29  	eid := client.GetEventID()
    30  	tid := client.GetTaskID()
    31  
    32  	var pairs []string
    33  	if eid != 0 {
    34  		pairs = append(pairs, EventIDMetadataKey, fmt.Sprint(eid))
    35  	}
    36  	if tid != 0 {
    37  		pairs = append(pairs, TaskIDMetadataKey, fmt.Sprint(tid))
    38  	}
    39  	if eid != 0 || tid != 0 {
    40  		md := metadata.Pairs(pairs...)
    41  		ctx = metadata.NewContext(ctx, md)
    42  	}
    43  	return grpc.Invoke(ctx, method, args, reply, cc, opts...)
    44  }
    45  
    46  // ExtractIDs attempts to extract a Task ID and an
    47  // Event ID from ctx, and sets it as the calling
    48  // goroutine's current IDs. It should be called at
    49  // the beginning of any gRPC handler. An error will
    50  // be returned if either ID was not found, or if it
    51  // could not be parsed. Note that if one ID was found
    52  // and parsed successfully, but the other was not,
    53  // neither will be set.
    54  func ExtractIDs(ctx context.Context) error {
    55  	md, ok := metadata.FromContext(ctx)
    56  	if !ok {
    57  		return fmt.Errorf("no metadata found in context")
    58  	}
    59  
    60  	eid, err := getIDFromMetadata(md, EventIDMetadataKey, "Event ID")
    61  	if err != nil {
    62  		return err
    63  	}
    64  	tid, err := getIDFromMetadata(md, TaskIDMetadataKey, "Task ID")
    65  	if err != nil {
    66  		return err
    67  	}
    68  
    69  	client.SetEventID(eid)
    70  	client.SetTaskID(tid)
    71  	return nil
    72  }
    73  
    74  // get the given ID from md; name should be "Task ID" or "Event ID",
    75  // and getIDFromMetadata will produce the proper error messages that
    76  // can be returned directly without wrapping.
    77  func getIDFromMetadata(md metadata.MD, key, name string) (int64, error) {
    78  	m := make(map[string]string)
    79  	for k, v := range md {
    80  		for _, v := range v {
    81  			kk, vv, err := metadata.DecodeKeyValue(k, v)
    82  			if err != nil {
    83  				return 0, fmt.Errorf("malformed metadata: %v", err)
    84  			}
    85  			m[kk] = vv
    86  		}
    87  	}
    88  	str, ok := m[key]
    89  	if !ok {
    90  		return 0, fmt.Errorf("no %v found in metadata", name)
    91  	}
    92  	// use ParseUint because IDs should never be
    93  	// negative (even though they're stored as
    94  	// int64s)
    95  	id, err := strconv.ParseUint(str, 10, 63)
    96  	if err != nil {
    97  		return 0, fmt.Errorf("parse metadata for key %q (value %q): %v", key, str, err)
    98  	}
    99  	return int64(id), nil
   100  }