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

     1  package interceptor
     2  
     3  import (
     4  	"bytes"
     5  	"container/ring"
     6  	"io"
     7  
     8  	sparta "github.com/mweagle/Sparta"
     9  	"github.com/sirupsen/logrus"
    10  )
    11  
    12  // XRay attributes
    13  const (
    14  	// XRayAttrBuildID is the XRay attribute associated with this
    15  	// service instance. See the official AWS docs at
    16  	// https://docs.aws.amazon.com/xray/latest/devguide/xray-sdk-go-segment.html#xray-sdk-go-segment-annotations
    17  	// for more information on XRay attributes
    18  	XRayAttrBuildID = "buildID"
    19  )
    20  
    21  // XRay metadata
    22  const (
    23  	// XRayMetadataErrValue is the metadata kayname used to store
    24  	// the error value when processing an event
    25  	XRayMetadataErrValue = "error"
    26  
    27  	// XRayMetadataErrEvent is the event associated with a lambda
    28  	// function that errors out
    29  	XRayMetadataErrEvent = "event"
    30  
    31  	// XRayMetadataRequestID is the AWS request ID that came along with the request
    32  	XRayMetadataRequestID = "reqID"
    33  
    34  	// XRayMetadataLogs is the key associated with the logfile entries. All log
    35  	// entries regardless of level will be included in the errLog value
    36  	XRayMetadataLogs = "log"
    37  )
    38  
    39  // XRayInterceptorMode represents the mode to use for the XRay interceptor
    40  type XRayInterceptorMode uint32
    41  
    42  const (
    43  	// XRayModeErrCaptureErrorValue = is the flag indicating to capture the error
    44  	// value iff it's non-empty
    45  	XRayModeErrCaptureErrorValue XRayInterceptorMode = 1 << iota
    46  	// XRayModeErrCaptureEvent is the flag indicating to capture the input event iff
    47  	// there was an error
    48  	XRayModeErrCaptureEvent
    49  	// XRayModeErrCaptureLogs is the flag indicating to capture all logs iff there
    50  	// was an error
    51  	XRayModeErrCaptureLogs
    52  	// XRayModeErrCaptureRequestID is the flag indicating to capture the request ID iff there
    53  	// was an error
    54  	XRayModeErrCaptureRequestID
    55  
    56  	// XRayAll is all options
    57  	XRayAll = XRayModeErrCaptureErrorValue |
    58  		XRayModeErrCaptureEvent |
    59  		XRayModeErrCaptureLogs |
    60  		XRayModeErrCaptureRequestID
    61  )
    62  
    63  var (
    64  	// devNullLogEntry is the reserved byte value that's returned by the
    65  	// filteringFormatter to instruct the Writer to throw away
    66  	// the serialized version.
    67  	//lint:ignore U1000 because it's actually used
    68  	devNullLogEntry = []byte("/dev/null")
    69  )
    70  
    71  // So we can turn up the level to max. Which means everything will go to the
    72  // serializer. If that happens, we need to know which entries
    73  // can be thrown away. A custom formatter can be used for that. So we can fake
    74  // this by returning a known string that tells the Out method to discard the data...
    75  // This isn't recommended, but it does tie things together
    76  //lint:ignore U1000 because it's used
    77  type filteringFormatter struct {
    78  	targetFormatter logrus.Formatter
    79  	originalLevel   logrus.Level
    80  	logRing         *ring.Ring
    81  }
    82  
    83  func (ff *filteringFormatter) Format(entry *logrus.Entry) ([]byte, error) {
    84  	ff.logRing.Value = entry
    85  	ff.logRing = ff.logRing.Next()
    86  
    87  	if ff.originalLevel >= entry.Level {
    88  		return ff.targetFormatter.Format(entry)
    89  	}
    90  	return devNullLogEntry, nil
    91  }
    92  
    93  // The filteringWriter works together with the filteringFormatter
    94  // to ignore formatted entries that are the /dev/null log entry values
    95  //lint:ignore U1000 because it's used
    96  type filteringWriter struct {
    97  	targetOutput io.Writer
    98  }
    99  
   100  func (fw *filteringWriter) Write(p []byte) (n int, err error) {
   101  	if !bytes.Equal(p, devNullLogEntry) {
   102  		return fw.targetOutput.Write(p)
   103  	}
   104  	return len(p), nil
   105  }
   106  
   107  // xrayInterceptor is an implementation of sparta.LambdaEventInterceptors that
   108  // handles tapping the event handling workflow and publishing a span with optional
   109  // request information on error.
   110  type xrayInterceptor struct {
   111  	mode XRayInterceptorMode
   112  	//lint:ignore U1000 because it's used
   113  	filteringFormatter *filteringFormatter
   114  	//lint:ignore U1000 because it's used
   115  	filteringWriter *filteringWriter
   116  }
   117  
   118  // RegisterXRayInterceptor handles pushing the tracing information into XRay
   119  func RegisterXRayInterceptor(handler *sparta.LambdaEventInterceptors,
   120  	mode XRayInterceptorMode) *sparta.LambdaEventInterceptors {
   121  	interceptor := &xrayInterceptor{
   122  		mode: mode,
   123  	}
   124  	if handler == nil {
   125  		handler = &sparta.LambdaEventInterceptors{}
   126  	}
   127  	return handler.Register(interceptor)
   128  }