[ffapi] [ffresty] Allow for Configuring Request Header for Distributed Trace Logging

Signed-off-by: hfuss <[email protected]>

Author: onelapahead

pkg/ffapi/handler.go

@@ -39,7 +39,24 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-const FFRequestIDHeader = "X-FireFly-Request-ID"
+const DefaultRequestIDHeader = "X-FireFly-Request-ID"
+
+var (
+	requestIDHeader string = DefaultRequestIDHeader
+)
+
+// RequestIDHeader returns the header name used to pass the request ID for
+// ffapi server and ffresty client requests.
+func RequestIDHeader() string {
+	return requestIDHeader
+}
+
+// SetRequestIDHeader configures the header used to pass the request ID between
+// ffapi server and ffresty client requests. Should be set at initialization of
+// a process.
+func SetRequestIDHeader(header string) {
+	requestIDHeader = header
+}
 
 type (
 	CtxHeadersKey     struct{}
@@ -343,7 +360,7 @@ func (hs *HandlerFactory) APIWrapper(handler HandlerFunction) http.HandlerFunc {
 
 		reqTimeout := hs.getTimeout(req)
 		ctx, cancel := context.WithTimeout(req.Context(), reqTimeout)
-		httpReqID := req.Header.Get(FFRequestIDHeader)
+		httpReqID := req.Header.Get(RequestIDHeader())
 		if httpReqID == "" {
 			httpReqID = fftypes.ShortID()
 		}

pkg/ffapi/handler_test.go

@@ -108,13 +108,23 @@ func TestRouteServePOST201WithParams(t *testing.T) {
 			assert.Equal(t, "true", r.QP["param2"])
 			assert.Equal(t, "false", r.QP["param3"])
 			assert.Equal(t, []string{"x", "y"}, r.QAP["param4"])
+			assert.Equal(t, "12345", r.Req.Context().Value(CtxFFRequestIDKey{}).(string))
+			assert.Equal(t, "custom-value", r.Req.Context().Value(CtxHeadersKey{}).(http.Header).Get("X-Custom-Header"))
 			return map[string]interface{}{"output1": "value2"}, nil
 		},
 	}}, "", nil)
 	defer done()
+	SetRequestIDHeader("x-unittest-req-id") // tests custom req header
 
 	b, _ := json.Marshal(map[string]interface{}{"input1": "value1"})
-	res, err := http.Post(fmt.Sprintf("http://%s/test/stuff?param1=thing&param2&param3=false&param4=x&param4=y", s.Addr()), "application/json", bytes.NewReader(b))
+
+	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("http://%s/test/stuff?param1=thing&param2&param3=false&param4=x&param4=y", s.Addr()), bytes.NewReader(b))
+	require.NoError(t, err)
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("X-Custom-Header", "custom-value") // tests custom header
+	req.Header.Set("x-unittest-req-id", "12345")      // tests custom req header
+
+	res, err := http.DefaultClient.Do(req)
 	assert.NoError(t, err)
 	assert.Equal(t, 201, res.StatusCode)
 	var resJSON map[string]interface{}
@@ -134,13 +144,19 @@ func TestRouteServePOST201WithParamsLargeNumber(t *testing.T) {
 		JSONOutputCodes: []int{201},
 		JSONHandler: func(r *APIRequest) (output interface{}, err error) {
 			assert.Equal(t, r.Input, map[string]interface{}{"largeNumberParam": json.Number("10000000000000000000000000001")})
+			assert.Equal(t, "12345", r.Req.Context().Value(CtxFFRequestIDKey{}).(string))
 			// Echo the input back as the response
 			return r.Input, nil
 		},
 	}}, "", nil)
 	defer done()
 
-	res, err := http.Post(fmt.Sprintf("http://%s/test/stuff", s.Addr()), "application/json", bytes.NewReader([]byte(largeParamLiteral)))
+	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("http://%s/test/stuff", s.Addr()), bytes.NewReader([]byte(largeParamLiteral)))
+	require.NoError(t, err)
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set(DefaultRequestIDHeader, "12345") // tests client setting req header
+
+	res, err := http.DefaultClient.Do(req)
 	assert.NoError(t, err)
 	assert.Equal(t, 201, res.StatusCode)
 	var resJSON map[string]interface{}

pkg/ffresty/ffresty.go

@@ -305,7 +305,7 @@ func NewWithConfig(ctx context.Context, ffrestyConfig Config) (client *resty.Cli
 		// If an X-FireFlyRequestID was set on the context, pass that header on this request too
 		ffRequestID := rCtx.Value(ffapi.CtxFFRequestIDKey{})
 		if ffRequestID != nil {
-			req.Header.Set(ffapi.FFRequestIDHeader, ffRequestID.(string))
+			req.Header.Set(ffapi.RequestIDHeader(), ffRequestID.(string))
 		}
 
 		if ffrestyConfig.OnBeforeRequest != nil {

pkg/ffresty/ffresty_test.go

@@ -558,7 +558,7 @@ func TestPassthroughHeaders(t *testing.T) {
 
 	httpmock.RegisterResponder("GET", "http://localhost:12345/test",
 		func(req *http.Request) (*http.Response, error) {
-			assert.Equal(t, "customReqID", req.Header.Get(ffapi.FFRequestIDHeader))
+			assert.Equal(t, "customReqID", req.Header.Get(ffapi.RequestIDHeader()))
 			assert.Equal(t, "headervalue", req.Header.Get("someheader"))
 			assert.Equal(t, "custom value", req.Header.Get("X-Custom-Header"))
 			assert.Equal(t, "Basic dXNlcjpwYXNz", req.Header.Get("Authorization"))