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 }