Merge pull request #756 from uber/OVERRIDE-ENV-FOR-SHADOW-REQUEST

Override env for shadow request

examples/example-gateway/build.yaml

@@ -35,3 +35,4 @@ moduleSearchPaths:
 defaultDependencies:
   endpoint:
     - middlewares/default/*
+shadowRequestHeader: x-shadow-request

examples/example-gateway/config/production.yaml

@@ -90,3 +90,4 @@ grpc.clientServiceNameMapping:
   echo: echo
 router.whitelistedPaths:
   - /path/whitelisted
+service.shadow.env.override.enable: true

examples/example-gateway/config/test.yaml

@@ -89,3 +89,4 @@ clients.baz.alternates:
       port: 8114
 grpc.clientServiceNameMapping:
   echo: echo
+service.shadow.env.override.enable: true

runtime/constants.go

@@ -51,6 +51,10 @@ const (
 	clientHTTPUnmarshalError = "client.http-unmarshal-error"
 	// clientTchannelReadError is the metric for tracking errors in reading tchannel response
 	clientTchannelUnmarshalError = "client.tchannel-unmarshal-error"
+
+	// shadow headers and environment
+	shadowEnvironment = "shadow"
+	environmentKey    = "env"
 )
 
 var knownMetrics = []string{

runtime/router.go

@@ -89,6 +89,7 @@ type RouterEndpoint struct {
 	contextLogger    ContextLogger
 	scope            tally.Scope
 	tracer           opentracing.Tracer
+	config           *StaticConfig
 }
 
 // NewRouterEndpoint creates an endpoint that can be registered to HTTPRouter
@@ -108,6 +109,7 @@ func NewRouterEndpoint(
 		scope:            deps.Scope,
 		tracer:           deps.Tracer,
 		JSONWrapper:      deps.JSONWrapper,
+		config:           deps.Config,
 	}
 }
 

runtime/server_http_request.go

@@ -88,6 +88,7 @@ func NewServerHTTPRequest(
 	}
 	if endpoint.contextExtractor != nil {
 		headers := map[string]string{}
+
 		for k, v := range r.Header {
 			// TODO: this 0th element logic is probably not correct
 			headers[k] = v[0]
@@ -99,6 +100,18 @@ func NewServerHTTPRequest(
 
 		logFields = append(logFields, endpoint.contextExtractor.ExtractLogFields(ctx)...)
 	}
+
+	// Overriding the environment for shadow requests
+	if endpoint.config != nil {
+		if endpoint.config.ContainsKey("service.shadow.env.override.enable") &&
+			endpoint.config.MustGetBoolean("service.shadow.env.override.enable") &&
+			endpoint.config.ContainsKey("shadowRequestHeader") &&
+			r.Header.Get(endpoint.config.MustGetString("shadowRequestHeader")) != "" {
+			scopeTags[environmentKey] = shadowEnvironment
+			logFields = append(logFields, zap.String(environmentKey, shadowEnvironment))
+		}
+	}
+
 	ctx = WithScopeTags(ctx, scopeTags)
 	ctx = WithLogFields(ctx, logFields...)
 

runtime/server_http_request_test.go

@@ -2238,7 +2238,7 @@ func TestSpanCreated(t *testing.T) {
 	assert.Equal(t, "200 OK", resp.Status)
 }
 
-func TestIncomingHTTPRequestServerLog(t *testing.T) {
+func testIncomingHTTPRequestServerLog(t *testing.T, isShadowRequest bool, environment string) {
 	gateway, err := benchGateway.CreateGateway(
 		defaultTestConfig,
 		defaultTestOptions,
@@ -2256,6 +2256,7 @@ func TestIncomingHTTPRequestServerLog(t *testing.T) {
 		Logger:        bgateway.ActualGateway.Logger,
 		ContextLogger: bgateway.ActualGateway.ContextLogger,
 		Tracer:        bgateway.ActualGateway.Tracer,
+		Config:        bgateway.ActualGateway.Config,
 	}
 	err = bgateway.ActualGateway.HTTPRouter.Handle(
 		"GET", "/foo", http.HandlerFunc(zanzibar.NewRouterEndpoint(
@@ -2273,13 +2274,20 @@ func TestIncomingHTTPRequestServerLog(t *testing.T) {
 	)
 	assert.NoError(t, err)
 
-	_, err = gateway.MakeRequest("GET", "/foo?bar=bar", nil, nil)
+	var headers map[string]string
+	if isShadowRequest {
+		headers = map[string]string{
+			"X-Shadow-Request": "true",
+		}
+	}
+	_, err = gateway.MakeRequest("GET", "/foo?bar=bar", headers, nil)
 	assert.NoError(t, err)
 
 	allLogs := bgateway.AllLogs()
 	assert.Equal(t, 1, len(allLogs["Finished an incoming server HTTP request with 200 status code"]))
 
 	tags := allLogs["Finished an incoming server HTTP request with 200 status code"][0]
+
 	dynamicHeaders := []string{
 		"requestUUID",
 		"remoteAddr",
@@ -2290,14 +2298,19 @@ func TestIncomingHTTPRequestServerLog(t *testing.T) {
 		"pid",
 		"timestamp-finished",
 	}
+
+	if isShadowRequest {
+		dynamicHeaders = append(dynamicHeaders, "X-Shadow-Request")
+	}
+
 	for _, dynamicValue := range dynamicHeaders {
 		assert.Contains(t, tags, dynamicValue)
 		delete(tags, dynamicValue)
 	}
 
 	expectedValues := map[string]interface{}{
 		"msg":             "Finished an incoming server HTTP request with 200 status code",
-		"env":             "test",
+		"env":             environment,
 		"level":           "debug",
 		"zone":            "unknown",
 		"service":         "example-gateway",
@@ -2319,3 +2332,28 @@ func TestIncomingHTTPRequestServerLog(t *testing.T) {
 		assert.Equal(t, expectedValue, tags[expectedKey], "unexpected header %q", expectedKey)
 	}
 }
+
+func TestIncomingHTTPRequestServerLogForDiffRequestTypes(t *testing.T) {
+	tests := []struct {
+		name                string
+		isShadowRequest     bool
+		expectedEnvironment string
+	}{
+		{
+			"Test incoming http request server log for normal request",
+			false,
+			"test",
+		},
+		{
+			"Test incoming http request server log for shadow request",
+			true,
+			"shadow",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			testIncomingHTTPRequestServerLog(t, tt.isShadowRequest, tt.expectedEnvironment)
+		})
+	}
+}

test/lib/util/util.go

@@ -52,6 +52,10 @@ func DefaultConfigFiles(serviceName string) []string {
 			"..", "..", "..", "examples", "example-gateway", "config",
 			"example-gateway", "test.yaml",
 		),
+		filepath.Join(
+			getDirName(),
+			"..", "..", "..", "examples", "example-gateway", "build.yaml",
+		),
 	}
 }