github.com/blend/go-sdk@v1.20220411.3/grpcutil/logged.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package grpcutil
     9  
    10  import (
    11  	"context"
    12  	"io"
    13  	"time"
    14  
    15  	"google.golang.org/grpc"
    16  	"google.golang.org/grpc/metadata"
    17  
    18  	"github.com/blend/go-sdk/ex"
    19  	"github.com/blend/go-sdk/logger"
    20  )
    21  
    22  // LoggedServerUnary returns a unary server interceptor.
    23  func LoggedServerUnary(log logger.Triggerable) grpc.UnaryServerInterceptor {
    24  	return func(ctx context.Context, args interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    25  		startTime := time.Now().UTC()
    26  		result, err := handler(ctx, args)
    27  		if log != nil {
    28  			event := NewRPCEvent(info.FullMethod, time.Now().UTC().Sub(startTime))
    29  			event.Err = err
    30  			if md, ok := metadata.FromIncomingContext(ctx); ok {
    31  				event.Authority = MetaValue(md, MetaTagAuthority)
    32  				event.UserAgent = MetaValue(md, MetaTagUserAgent)
    33  				event.ContentType = MetaValue(md, MetaTagContentType)
    34  			}
    35  			log.TriggerContext(ctx, event)
    36  		}
    37  		return result, err
    38  	}
    39  }
    40  
    41  // LoggedClientUnary returns a unary client interceptor.
    42  func LoggedClientUnary(log logger.Triggerable) grpc.UnaryClientInterceptor {
    43  	return func(ctx context.Context, method string, req interface{}, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    44  		startTime := time.Now().UTC()
    45  		err := invoker(ctx, method, req, reply, cc, opts...)
    46  		if log != nil {
    47  			event := NewRPCEvent(method, time.Now().UTC().Sub(startTime))
    48  			event.Err = err
    49  			if md, ok := metadata.FromOutgoingContext(ctx); ok {
    50  				event.Authority = MetaValue(md, MetaTagAuthority)
    51  				event.UserAgent = MetaValue(md, MetaTagUserAgent)
    52  				event.ContentType = MetaValue(md, MetaTagContentType)
    53  			}
    54  			log.TriggerContext(ctx, event)
    55  		}
    56  		return err
    57  	}
    58  }
    59  
    60  // LoggedServerStream returns a stream server interceptor.
    61  func LoggedServerStream(log logger.Triggerable) grpc.StreamServerInterceptor {
    62  	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) (err error) {
    63  		startTime := time.Now().UTC()
    64  		err = handler(srv, instrumentedServerStream{ServerStream: stream, Method: info.FullMethod, Log: log})
    65  		if log != nil {
    66  			event := NewRPCEvent(info.FullMethod, time.Now().UTC().Sub(startTime))
    67  			event.Err = err
    68  			if md, ok := metadata.FromIncomingContext(stream.Context()); ok {
    69  				event.Authority = MetaValue(md, MetaTagAuthority)
    70  				event.UserAgent = MetaValue(md, MetaTagUserAgent)
    71  				event.ContentType = MetaValue(md, MetaTagContentType)
    72  			}
    73  			log.TriggerContext(stream.Context(), event)
    74  		}
    75  		return err
    76  	}
    77  }
    78  
    79  // LoggedClientStream returns a stream server interceptor.
    80  func LoggedClientStream(log logger.Triggerable) grpc.StreamClientInterceptor {
    81  	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
    82  		startTime := time.Now().UTC()
    83  		clientStreamer, err := streamer(ctx, desc, cc, method, opts...)
    84  		if log != nil {
    85  			event := NewRPCEvent(method, time.Now().UTC().Sub(startTime))
    86  			event.Err = err
    87  			if md, ok := metadata.FromOutgoingContext(ctx); ok {
    88  				event.Authority = MetaValue(md, MetaTagAuthority)
    89  				event.UserAgent = MetaValue(md, MetaTagUserAgent)
    90  				event.ContentType = MetaValue(md, MetaTagContentType)
    91  			}
    92  			log.TriggerContext(ctx, event)
    93  		}
    94  		return clientStreamer, err
    95  	}
    96  }
    97  
    98  type instrumentedServerStream struct {
    99  	grpc.ServerStream
   100  	Method string
   101  	Log    logger.Triggerable
   102  }
   103  
   104  // RecvMessage overrides the underlying RecvMsg method.
   105  func (iss instrumentedServerStream) RecvMsg(m interface{}) (err error) {
   106  	if iss.Log != nil {
   107  		startTime := time.Now().UTC()
   108  		defer func() {
   109  			if ex.Is(err, context.Canceled) || ex.Is(err, io.EOF) {
   110  				return
   111  			}
   112  			event := NewRPCStreamMessageEvent(iss.Method, StreamMessageDirectionReceive, time.Now().UTC().Sub(startTime))
   113  			event.Err = err
   114  			if md, ok := metadata.FromIncomingContext(iss.ServerStream.Context()); ok {
   115  				event.Authority = MetaValue(md, MetaTagAuthority)
   116  				event.UserAgent = MetaValue(md, MetaTagUserAgent)
   117  				event.ContentType = MetaValue(md, MetaTagContentType)
   118  			}
   119  			iss.Log.TriggerContext(iss.ServerStream.Context(), event)
   120  		}()
   121  	}
   122  	return iss.ServerStream.RecvMsg(m)
   123  }
   124  
   125  // SendMsg overrides the underlying SendMsg method.
   126  func (iss instrumentedServerStream) SendMsg(m interface{}) (err error) {
   127  	if iss.Log != nil {
   128  		startTime := time.Now().UTC()
   129  		defer func() {
   130  			if ex.Is(err, context.Canceled) || ex.Is(err, io.EOF) {
   131  				return
   132  			}
   133  			event := NewRPCStreamMessageEvent(iss.Method, StreamMessageDirectionSend, time.Now().UTC().Sub(startTime))
   134  			event.Err = err
   135  			if md, ok := metadata.FromIncomingContext(iss.ServerStream.Context()); ok {
   136  				event.Authority = MetaValue(md, MetaTagAuthority)
   137  				event.UserAgent = MetaValue(md, MetaTagUserAgent)
   138  				event.ContentType = MetaValue(md, MetaTagContentType)
   139  			}
   140  			iss.Log.TriggerContext(iss.ServerStream.Context(), event)
   141  		}()
   142  	}
   143  	return iss.ServerStream.SendMsg(m)
   144  }