github.com/mweagle/Sparta@v1.15.0/interceptor/xray_awsbinary.go (about)

     1  // +build lambdabinary
     2  
     3  package interceptor
     4  
     5  import (
     6  	"container/ring"
     7  	"context"
     8  	"encoding/json"
     9  	"log"
    10  
    11  	"github.com/aws/aws-lambda-go/lambdacontext"
    12  	"github.com/aws/aws-xray-sdk-go/xray"
    13  	sparta "github.com/mweagle/Sparta"
    14  	"github.com/sirupsen/logrus"
    15  )
    16  
    17  const (
    18  	logRingSize = 1024
    19  )
    20  
    21  func (xri *xrayInterceptor) Begin(ctx context.Context, msg json.RawMessage) context.Context {
    22  
    23  	// Put this into the context...
    24  	segmentCtx, segment := xray.BeginSubsegment(ctx, "Sparta")
    25  
    26  	// Add some xray annotations to help track this version
    27  	// of the service
    28  	// https://docs.aws.amazon.com/xray/latest/devguide/xray-sdk-go-segment.html
    29  	if segment != nil {
    30  		errAddAnnotation := segment.AddAnnotation(XRayAttrBuildID, sparta.StampedBuildID)
    31  		if errAddAnnotation != nil {
    32  			log.Printf("Failed to update segment context: %s", errAddAnnotation)
    33  		}
    34  		segmentCtx = context.WithValue(segmentCtx, contextKeySegment, segment)
    35  	}
    36  	return segmentCtx
    37  }
    38  
    39  func (xri *xrayInterceptor) BeforeSetup(ctx context.Context, msg json.RawMessage) context.Context {
    40  	return ctx
    41  }
    42  
    43  func (xri *xrayInterceptor) AfterSetup(ctx context.Context, msg json.RawMessage) context.Context {
    44  	if xri.mode&XRayModeErrCaptureLogs != 0 {
    45  		logger, loggerOk := ctx.Value(sparta.ContextKeyLogger).(*logrus.Logger)
    46  		if loggerOk {
    47  			// So we need a loggerWrapper that has the debug level turned on.
    48  			// This filtering formatter will put everything in a logring and
    49  			// flush iff there is an error
    50  			xri.filteringFormatter = &filteringFormatter{
    51  				targetFormatter: logger.Formatter,
    52  				originalLevel:   logger.Level,
    53  				logRing:         ring.New(logRingSize),
    54  			}
    55  			xri.filteringWriter = &filteringWriter{
    56  				targetOutput: logger.Out,
    57  			}
    58  			logger.SetLevel(logrus.TraceLevel)
    59  			logger.SetFormatter(xri.filteringFormatter)
    60  			logger.SetOutput(xri.filteringWriter)
    61  		} else {
    62  			log.Printf("WARNING: Failed to get logger from context\n")
    63  		}
    64  	}
    65  	return ctx
    66  }
    67  
    68  func (xri *xrayInterceptor) BeforeDispatch(ctx context.Context, msg json.RawMessage) context.Context {
    69  	return ctx
    70  }
    71  func (xri *xrayInterceptor) AfterDispatch(ctx context.Context, msg json.RawMessage) context.Context {
    72  	return ctx
    73  }
    74  
    75  func (xri *xrayInterceptor) Complete(ctx context.Context, msg json.RawMessage) context.Context {
    76  	segmentVal := ctx.Value(contextKeySegment)
    77  	if segmentVal != nil {
    78  		segment, segmentOk := segmentVal.(*xray.Segment)
    79  		if segmentOk {
    80  			errValue, errValueOk := ctx.Value(sparta.ContextKeyLambdaError).(error)
    81  			if errValueOk && errValue != nil {
    82  
    83  				// Include the error value?
    84  				if xri.mode&XRayModeErrCaptureErrorValue != 0 {
    85  					metadataEventErr := segment.AddMetadataToNamespace(sparta.ProperName, XRayMetadataErrValue, errValue.Error())
    86  					if metadataEventErr != nil {
    87  						log.Printf("Failed to set event %s metadata: %s", XRayMetadataErrValue, metadataEventErr)
    88  					}
    89  				}
    90  
    91  				// Include the event?
    92  				if xri.mode&XRayModeErrCaptureEvent != 0 {
    93  					metadataEventErr = segment.AddMetadataToNamespace(sparta.ProperName, XRayMetadataErrEvent, msg)
    94  					if metadataEventErr != nil {
    95  						log.Printf("Failed to set event %s metadata: %s", XRayMetadataErrEvent, metadataEventErr)
    96  					}
    97  				}
    98  
    99  				// Include the request ID?
   100  				if xri.mode&XRayModeErrCaptureRequestID != 0 {
   101  					awsContext, _ := lambdacontext.FromContext(ctx)
   102  					if awsContext != nil {
   103  						metadataEventErr = segment.AddMetadataToNamespace(sparta.ProperName, XRayMetadataRequestID, awsContext.AwsRequestID)
   104  						if metadataEventErr != nil {
   105  							log.Printf("Failed to set event %s metadata: %s", XRayMetadataRequestID, metadataEventErr)
   106  						}
   107  					}
   108  				}
   109  
   110  				// Include the cached log info?
   111  				if xri.mode&XRayModeErrCaptureLogs != 0 {
   112  					logFormatter := &logrus.JSONFormatter{}
   113  					logMessages := make([]map[string]*json.RawMessage, 0)
   114  					xri.filteringFormatter.logRing.Do(func(eachLogEntry interface{}) {
   115  						if eachLogEntry != nil {
   116  							// Format each one to text?
   117  							eachTypedEntry, eachTypedEntryOk := eachLogEntry.(*logrus.Entry)
   118  							if eachTypedEntryOk {
   119  								formatted, formattedErr := logFormatter.Format(eachTypedEntry)
   120  								if formattedErr == nil {
   121  									var jsonData map[string]*json.RawMessage
   122  									unmarshalErr := json.Unmarshal(formatted, &jsonData)
   123  									if unmarshalErr == nil {
   124  										logMessages = append(logMessages, jsonData)
   125  									}
   126  								}
   127  							}
   128  						}
   129  					})
   130  					metadataEventErr = segment.AddMetadataToNamespace(sparta.ProperName, XRayMetadataLogs, logMessages)
   131  					if metadataEventErr != nil {
   132  						log.Printf("Failed to set event %s metadata: %s", XRayMetadataLogs, metadataEventErr)
   133  					}
   134  					// Either way, clear out the ring...
   135  					xri.filteringFormatter.logRing = ring.New(logRingSize)
   136  				}
   137  			}
   138  			segment.Close(errValue)
   139  		}
   140  	}
   141  	return ctx
   142  }