xray: Added a const TraceIDHeaderKey with x-amzn-trace-id. (#85)

aws · Feb 15, 2019 · 8c8a645 · 8c8a645
1 parent 1eb1439
commit 8c8a645
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 8 deletions.
diff --git a/xray/aws.go b/xray/aws.go
@@ -28,6 +28,7 @@ import (
 const RequestIDKey string = "request_id"
 const ExtendedRequestIDKey string = "id_2"
 const S3ExtendedRequestIDHeaderKey string = "x-amz-id-2"
+const TraceIDHeaderKey = "x-amzn-trace-id"
 
 type jsonMap struct {
 	object interface{}
@@ -63,7 +64,7 @@ var xRayBeforeValidateHandler = request.NamedHandler{
 		marshalctx, _ := BeginSubsegment(ctx, "marshal")
 
 		r.HTTPRequest = r.HTTPRequest.WithContext(marshalctx)
-		r.HTTPRequest.Header.Set("x-amzn-trace-id", opseg.DownstreamHeader().String())
+		r.HTTPRequest.Header.Set(TraceIDHeaderKey, opseg.DownstreamHeader().String())
 	},
 }
 

diff --git a/xray/client.go b/xray/client.go
@@ -88,7 +88,7 @@ func (rt *roundtripper) RoundTrip(r *http.Request) (*http.Response, error) {
 		seg.GetHTTP().GetRequest().Method = r.Method
 		seg.GetHTTP().GetRequest().URL = r.URL.String()
 
-		r.Header.Set("x-amzn-trace-id", seg.DownstreamHeader().String())
+		r.Header.Set(TraceIDHeaderKey, seg.DownstreamHeader().String())
 		seg.Unlock()
 
 		resp, err = rt.Base.RoundTrip(r)

diff --git a/xray/handler.go b/xray/handler.go
@@ -92,7 +92,7 @@ func HandlerWithContext(ctx context.Context, sn SegmentNamer, h http.Handler) ht
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		name := sn.Name(r.Host)
 
-		traceHeader := header.FromString(r.Header.Get("x-amzn-trace-id"))
+		traceHeader := header.FromString(r.Header.Get(TraceIDHeaderKey))
 
 		r = r.WithContext(ctx)
 		c, seg := NewSegmentFromHeader(r.Context(), name, traceHeader)
@@ -110,7 +110,7 @@ func Handler(sn SegmentNamer, h http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		name := sn.Name(r.Host)
 
-		traceHeader := header.FromString(r.Header.Get("x-amzn-trace-id"))
+		traceHeader := header.FromString(r.Header.Get(TraceIDHeaderKey))
 		ctx, seg := NewSegmentFromHeader(r.Context(), name, traceHeader)
 
 		r = r.WithContext(ctx)
@@ -155,7 +155,7 @@ func httpTrace(seg *Segment, h http.Handler, w http.ResponseWriter, r *http.Requ
 		respHeader.WriteString(strconv.Itoa(btoi(seg.Sampled)))
 	}
 
-	w.Header().Set("x-amzn-trace-id", respHeader.String())
+	w.Header().Set(TraceIDHeaderKey, respHeader.String())
 	seg.Unlock()
 
 	resp := &responseCapturer{w, 200, 0}

diff --git a/xray/handler_test.go b/xray/handler_test.go
@@ -91,7 +91,7 @@ func TestHandlerWithContextForNonRootHandler(t *testing.T) {
 	defer ts.Close()
 
 	req := httptest.NewRequest("DELETE", ts.URL, strings.NewReader(""))
-	req.Header.Set("x-amzn-trace-id", "Root=fakeid; Parent=reqid; Sampled=1")
+	req.Header.Set(TraceIDHeaderKey, "Root=fakeid; Parent=reqid; Sampled=1")
 
 	_, err := http.DefaultTransport.RoundTrip(req)
 	assert.NoError(t, err)
@@ -143,7 +143,7 @@ func TestNonRootHandler(t *testing.T) {
 	defer ts.Close()
 
 	req := httptest.NewRequest("DELETE", ts.URL, strings.NewReader(""))
-	req.Header.Set("x-amzn-trace-id", "Root=fakeid; Parent=reqid; Sampled=1")
+	req.Header.Set(TraceIDHeaderKey, "Root=fakeid; Parent=reqid; Sampled=1")
 
 	_, err := http.DefaultTransport.RoundTrip(req)
 	assert.NoError(t, err)

diff --git a/xray/util_test.go b/xray/util_test.go
@@ -91,7 +91,7 @@ type XRayHeaders struct {
 }
 
 func ParseHeadersForTest(h http.Header) XRayHeaders {
-	traceHeader := header.FromString(h.Get("x-amzn-trace-id"))
+	traceHeader := header.FromString(h.Get(TraceIDHeaderKey))
 
 	return XRayHeaders{
 		RootTraceID: traceHeader.TraceID,