github.com/prysmaticlabs/prysm@v1.4.4/shared/grpcutils/grpcutils.go (about)

     1  package grpcutils
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"strings"
     8  	"time"
     9  
    10  	"github.com/sirupsen/logrus"
    11  	"google.golang.org/grpc"
    12  	"google.golang.org/grpc/metadata"
    13  )
    14  
    15  // LogRequests logs the gRPC backend as well as request duration when the log level is set to debug
    16  // or higher.
    17  func LogRequests(
    18  	ctx context.Context,
    19  	method string, req,
    20  	reply interface{},
    21  	cc *grpc.ClientConn,
    22  	invoker grpc.UnaryInvoker,
    23  	opts ...grpc.CallOption,
    24  ) error {
    25  	// Shortcut when debug logging is not enabled.
    26  	if logrus.GetLevel() < logrus.DebugLevel {
    27  		return invoker(ctx, method, req, reply, cc, opts...)
    28  	}
    29  
    30  	var header metadata.MD
    31  	opts = append(
    32  		opts,
    33  		grpc.Header(&header),
    34  	)
    35  	start := time.Now()
    36  	err := invoker(ctx, method, req, reply, cc, opts...)
    37  	logrus.WithField("backend", header["x-backend"]).
    38  		WithField("method", method).WithField("duration", time.Since(start)).
    39  		Debug("gRPC request finished.")
    40  	return err
    41  }
    42  
    43  // LogStream prints the method at DEBUG level at the start of the stream.
    44  func LogStream(
    45  	ctx context.Context,
    46  	sd *grpc.StreamDesc,
    47  	conn *grpc.ClientConn,
    48  	method string,
    49  	streamer grpc.Streamer,
    50  	opts ...grpc.CallOption,
    51  ) (grpc.ClientStream, error) {
    52  	// Shortcut when debug logging is not enabled.
    53  	if logrus.GetLevel() < logrus.DebugLevel {
    54  		return streamer(ctx, sd, conn, method, opts...)
    55  	}
    56  
    57  	var header metadata.MD
    58  	opts = append(
    59  		opts,
    60  		grpc.Header(&header),
    61  	)
    62  	strm, err := streamer(ctx, sd, conn, method, opts...)
    63  	logrus.WithField("backend", header["x-backend"]).
    64  		WithField("method", method).
    65  		Debug("gRPC stream started.")
    66  	return strm, err
    67  }
    68  
    69  // AppendHeaders parses the provided GRPC headers
    70  // and attaches them to the provided context.
    71  func AppendHeaders(parent context.Context, headers []string) context.Context {
    72  	for _, h := range headers {
    73  		if h != "" {
    74  			keyValue := strings.Split(h, "=")
    75  			if len(keyValue) < 2 {
    76  				logrus.Warnf("Incorrect gRPC header flag format. Skipping %v", keyValue[0])
    77  				continue
    78  			}
    79  			parent = metadata.AppendToOutgoingContext(parent, keyValue[0], strings.Join(keyValue[1:], "="))
    80  		}
    81  	}
    82  	return parent
    83  }
    84  
    85  // AppendCustomErrorHeader sets a CustomErrorMetadataKey gRPC header on the passed in context,
    86  // using the passed in error data as the header's value. The data is serialized as JSON.
    87  func AppendCustomErrorHeader(ctx context.Context, errorData interface{}) error {
    88  	j, err := json.Marshal(errorData)
    89  	if err != nil {
    90  		return fmt.Errorf("could not marshal error data into JSON: %w", err)
    91  	}
    92  	if err := grpc.SetHeader(ctx, metadata.Pairs(CustomErrorMetadataKey, string(j))); err != nil {
    93  		return fmt.Errorf("could not set custom error header: %w", err)
    94  	}
    95  	return nil
    96  }