github.com/newrelic/go-agent@v3.26.0+incompatible/_integrations/nrgrpc/nrgrpc_server.go (about)

     1  // Copyright 2020 New Relic Corporation. All rights reserved.
     2  // SPDX-License-Identifier: Apache-2.0
     3  
     4  package nrgrpc
     5  
     6  import (
     7  	"context"
     8  	"net/http"
     9  	"strings"
    10  
    11  	newrelic "github.com/newrelic/go-agent"
    12  	"google.golang.org/grpc"
    13  	"google.golang.org/grpc/metadata"
    14  	"google.golang.org/grpc/status"
    15  )
    16  
    17  func startTransaction(ctx context.Context, app newrelic.Application, fullMethod string) newrelic.Transaction {
    18  	method := strings.TrimPrefix(fullMethod, "/")
    19  
    20  	var hdrs http.Header
    21  	if md, ok := metadata.FromIncomingContext(ctx); ok {
    22  		hdrs = make(http.Header, len(md))
    23  		for k, vs := range md {
    24  			for _, v := range vs {
    25  				hdrs.Add(k, v)
    26  			}
    27  		}
    28  	}
    29  
    30  	target := hdrs.Get(":authority")
    31  	url := getURL(method, target)
    32  
    33  	webReq := newrelic.NewStaticWebRequest(hdrs, url, method, newrelic.TransportHTTP)
    34  	txn := app.StartTransaction(method, nil, nil)
    35  	txn.SetWebRequest(webReq)
    36  
    37  	return txn
    38  }
    39  
    40  // UnaryServerInterceptor instruments server unary RPCs.
    41  //
    42  // Use this function with grpc.UnaryInterceptor and a newrelic.Application to
    43  // create a grpc.ServerOption to pass to grpc.NewServer.  This interceptor
    44  // records each unary call with a transaction.  You must use both
    45  // UnaryServerInterceptor and StreamServerInterceptor to instrument unary and
    46  // streaming calls.
    47  //
    48  // Example:
    49  //
    50  //	cfg := newrelic.NewConfig("gRPC Server", os.Getenv("NEW_RELIC_LICENSE_KEY"))
    51  //	app, _ := newrelic.NewApplication(cfg)
    52  //	server := grpc.NewServer(
    53  //		grpc.UnaryInterceptor(nrgrpc.UnaryServerInterceptor(app)),
    54  //		grpc.StreamInterceptor(nrgrpc.StreamServerInterceptor(app)),
    55  //	)
    56  //
    57  // These interceptors add the transaction to the call context so it may be
    58  // accessed in your method handlers using newrelic.FromContext.
    59  //
    60  // Full example:
    61  // https://github.com/newrelic/go-agent/blob/master/_integrations/nrgrpc/example/server/server.go
    62  //
    63  func UnaryServerInterceptor(app newrelic.Application) grpc.UnaryServerInterceptor {
    64  	if nil == app {
    65  		return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    66  			return handler(ctx, req)
    67  		}
    68  	}
    69  
    70  	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
    71  		txn := startTransaction(ctx, app, info.FullMethod)
    72  		defer txn.End()
    73  
    74  		ctx = newrelic.NewContext(ctx, txn)
    75  		resp, err = handler(ctx, req)
    76  		txn.WriteHeader(int(status.Code(err)))
    77  		return
    78  	}
    79  }
    80  
    81  type wrappedServerStream struct {
    82  	grpc.ServerStream
    83  	txn newrelic.Transaction
    84  }
    85  
    86  func (s wrappedServerStream) Context() context.Context {
    87  	ctx := s.ServerStream.Context()
    88  	return newrelic.NewContext(ctx, s.txn)
    89  }
    90  
    91  func newWrappedServerStream(stream grpc.ServerStream, txn newrelic.Transaction) grpc.ServerStream {
    92  	return wrappedServerStream{
    93  		ServerStream: stream,
    94  		txn:          txn,
    95  	}
    96  }
    97  
    98  // StreamServerInterceptor instruments server streaming RPCs.
    99  //
   100  // Use this function with grpc.StreamInterceptor and a newrelic.Application to
   101  // create a grpc.ServerOption to pass to grpc.NewServer.  This interceptor
   102  // records each streaming call with a transaction.  You must use both
   103  // UnaryServerInterceptor and StreamServerInterceptor to instrument unary and
   104  // streaming calls.
   105  //
   106  // Example:
   107  //
   108  //	cfg := newrelic.NewConfig("gRPC Server", os.Getenv("NEW_RELIC_LICENSE_KEY"))
   109  //	app, _ := newrelic.NewApplication(cfg)
   110  //	server := grpc.NewServer(
   111  //		grpc.UnaryInterceptor(nrgrpc.UnaryServerInterceptor(app)),
   112  //		grpc.StreamInterceptor(nrgrpc.StreamServerInterceptor(app)),
   113  //	)
   114  //
   115  // These interceptors add the transaction to the call context so it may be
   116  // accessed in your method handlers using newrelic.FromContext.
   117  //
   118  // Full example:
   119  // https://github.com/newrelic/go-agent/blob/master/_integrations/nrgrpc/example/server/server.go
   120  //
   121  func StreamServerInterceptor(app newrelic.Application) grpc.StreamServerInterceptor {
   122  	if nil == app {
   123  		return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
   124  			return handler(srv, ss)
   125  		}
   126  	}
   127  
   128  	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
   129  		txn := startTransaction(ss.Context(), app, info.FullMethod)
   130  		defer txn.End()
   131  
   132  		err := handler(srv, newWrappedServerStream(ss, txn))
   133  		txn.WriteHeader(int(status.Code(err)))
   134  		return err
   135  	}
   136  }