github.com/zak-blake/goa@v1.4.1/middleware/xray/middleware.go (about)

     1  package xray
     2  
     3  import (
     4  	"context"
     5  	"crypto/rand"
     6  	"fmt"
     7  	"net"
     8  	"net/http"
     9  	"sync"
    10  	"time"
    11  
    12  	"github.com/goadesign/goa"
    13  	"github.com/goadesign/goa/middleware"
    14  )
    15  
    16  const (
    17  	// segKey is the key used to store the segments in the context.
    18  	segKey key = iota + 1
    19  )
    20  
    21  // New returns a middleware that sends AWS X-Ray segments to the daemon running
    22  // at the given address.
    23  //
    24  // service is the name of the service reported to X-Ray. daemon is the hostname
    25  // (including port) of the X-Ray daemon collecting the segments.
    26  //
    27  // The middleware works by extracting the trace information from the context
    28  // using the tracing middleware package. The tracing middleware must be mounted
    29  // first on the service.
    30  //
    31  // The middleware stores the request segment in the context. Use ContextSegment
    32  // to retrieve it. User code can further configure the segment for example to set
    33  // a service version or record an error.
    34  //
    35  // User code may create child segments using the Segment NewSubsegment method
    36  // for tracing requests to external services. Such segments should be closed via
    37  // the Close method once the request completes. The middleware takes care of
    38  // closing the top level segment. Typical usage:
    39  //
    40  //     segment := xray.ContextSegment(ctx)
    41  //     sub := segment.NewSubsegment("external-service")
    42  //     defer sub.Close()
    43  //     err := client.MakeRequest()
    44  //     if err != nil {
    45  //         sub.Error = xray.Wrap(err)
    46  //     }
    47  //     return
    48  //
    49  func New(service, daemon string) (goa.Middleware, error) {
    50  	connection, err := periodicallyRedialingConn(context.Background(), time.Minute, func() (net.Conn, error) {
    51  		return net.Dial("udp", daemon)
    52  	})
    53  	if err != nil {
    54  		return nil, fmt.Errorf("xray: failed to connect to daemon - %s", err)
    55  	}
    56  	return func(h goa.Handler) goa.Handler {
    57  		return func(ctx context.Context, rw http.ResponseWriter, req *http.Request) error {
    58  			var (
    59  				err     error
    60  				traceID = middleware.ContextTraceID(ctx)
    61  			)
    62  			if traceID == "" {
    63  				// No tracing
    64  				return h(ctx, rw, req)
    65  			}
    66  
    67  			s := newSegment(ctx, traceID, service, req, connection())
    68  			ctx = WithSegment(ctx, s)
    69  
    70  			defer func() {
    71  				go func() {
    72  					defer s.Close()
    73  
    74  					s.RecordContextResponse(ctx)
    75  					if err != nil {
    76  						s.RecordError(err)
    77  					}
    78  				}()
    79  			}()
    80  
    81  			err = h(ctx, rw, req)
    82  
    83  			return err
    84  		}
    85  	}, nil
    86  }
    87  
    88  // NewID is a span ID creation algorithm which produces values that are
    89  // compatible with AWS X-Ray.
    90  func NewID() string {
    91  	b := make([]byte, 8)
    92  	rand.Read(b)
    93  	return fmt.Sprintf("%x", b)
    94  }
    95  
    96  // NewTraceID is a trace ID creation algorithm which produces values that are
    97  // compatible with AWS X-Ray.
    98  func NewTraceID() string {
    99  	b := make([]byte, 12)
   100  	rand.Read(b)
   101  	return fmt.Sprintf("%d-%x-%s", 1, time.Now().Unix(), fmt.Sprintf("%x", b))
   102  }
   103  
   104  // WithSegment creates a context containing the given segment. Use ContextSegment
   105  // to retrieve it.
   106  func WithSegment(ctx context.Context, s *Segment) context.Context {
   107  	return context.WithValue(ctx, segKey, s)
   108  }
   109  
   110  // ContextSegment extracts the segment set in the context with WithSegment.
   111  func ContextSegment(ctx context.Context) *Segment {
   112  	if s := ctx.Value(segKey); s != nil {
   113  		return s.(*Segment)
   114  	}
   115  	return nil
   116  }
   117  
   118  // newSegment creates a new segment for the incoming request.
   119  func newSegment(ctx context.Context, traceID, name string, req *http.Request, c net.Conn) *Segment {
   120  	var (
   121  		spanID   = middleware.ContextSpanID(ctx)
   122  		parentID = middleware.ContextParentSpanID(ctx)
   123  	)
   124  
   125  	s := NewSegment(name, traceID, spanID, c)
   126  	s.RecordRequest(req, "")
   127  
   128  	if parentID != "" {
   129  		s.ParentID = parentID
   130  	}
   131  
   132  	return s
   133  }
   134  
   135  // now returns the current time as a float appropriate for X-Ray processing.
   136  func now() float64 {
   137  	return float64(time.Now().Truncate(time.Millisecond).UnixNano()) / 1e9
   138  }
   139  
   140  // periodicallyRedialingConn creates a goroutine to periodically re-dial a connection, so the hostname can be
   141  // re-resolved if the IP changes.
   142  // Returns a func that provides the latest Conn value.
   143  func periodicallyRedialingConn(ctx context.Context, renewPeriod time.Duration, dial func() (net.Conn, error)) (func() net.Conn, error) {
   144  	var (
   145  		err error
   146  
   147  		// guard access to c
   148  		mu sync.RWMutex
   149  		c  net.Conn
   150  	)
   151  
   152  	// get an initial connection
   153  	if c, err = dial(); err != nil {
   154  		return nil, err
   155  	}
   156  
   157  	// periodically re-dial
   158  	go func() {
   159  		ticker := time.NewTicker(renewPeriod)
   160  		for {
   161  			select {
   162  			case <-ticker.C:
   163  				newConn, err := dial()
   164  				if err != nil {
   165  					continue // we don't have anything better to replace `c` with
   166  				}
   167  				mu.Lock()
   168  				c = newConn
   169  				mu.Unlock()
   170  			case <-ctx.Done():
   171  				return
   172  			}
   173  		}
   174  	}()
   175  
   176  	return func() net.Conn {
   177  		mu.RLock()
   178  		defer mu.RUnlock()
   179  		return c
   180  	}, nil
   181  }