Merge 5f89983 into 8f952a9

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Unreleased
 ### Changed
 - Most built-in logs like `Finished an outgoing client HTTP request` now use the context logger. 
+- **Context Extractor**: Added [`ContextExtractor`](https://godoc.org/github.com/uber/zanzibar/runtime#ContextExtractor) interface. It can be used to define "extractors" or functions to pull out dynamic fields like trace ID, request headers, etc. out of the context to be used in log fields and metric tags. These can be used to pull out fields that are application-specific without adding code to zanzibar.
+
+  **BREAKING CHANGE**
+- **GetContextScopeExtractors**:Added [`GetContextScopeExtractors`](https://godoc.org/github.com/uber/zanzibar/runtime#GetContextScopeExtractors) interface. For migration, define var AppOptions *zanzibar.Options at PackageRoot and implement interface accordingly.
 
 ### Fixed
 - HTTP `DELETE` methods on clients can now send a JSON payload. Previously it was silently discarded. 

codegen/module.go

@@ -828,6 +828,7 @@ func readPackageInfo(
 		// uniqueness of the provided package aliases. Note that the default
 		// package is "PackageName".
 		PackageAlias:          defaultAlias + "static",
+		PackageRoot:           packageRoot,
 		GeneratedPackageAlias: defaultAlias + "generated",
 		ModulePackageAlias:    defaultAlias + "module",
 		PackagePath: path.Join(
@@ -1231,6 +1232,8 @@ type PackageInfo struct {
 	PackageName string
 	// PackageAlias is the unique import alias for non-generated packages
 	PackageAlias string
+	// PackageRoot is the unique import root for non-generated packages
+	PackageRoot string
 	// GeneratedPackageAlias is the unique import alias for generated packages
 	GeneratedPackageAlias string
 	// ModulePackageAlias is the unique import alias for the module system's,

codegen/templates/main.tmpl

@@ -17,6 +17,7 @@ import (
 	"github.com/uber/zanzibar/config"
 	"github.com/uber/zanzibar/runtime"
 
+	app "{{$instance.PackageInfo.PackageRoot}}"
 	service "{{$instance.PackageInfo.GeneratedPackagePath}}"
 	module "{{$instance.PackageInfo.ModulePackagePath}}"
 )
@@ -42,8 +43,8 @@ func getConfig() *zanzibar.StaticConfig {
 
 func createGateway() (*zanzibar.Gateway, error) {
 	config := getConfig()
-	
-	gateway, _, err := service.CreateGateway(config, nil)
+
+	gateway, _, err := service.CreateGateway(config, app.AppOptions)
 	if err != nil {
 		return nil, err
 	}

runtime/context.go

@@ -30,11 +30,16 @@ import (
 
 type contextFieldKey string
 
+// ContextScopeTagsExtractor defines func where extracts tags from context
+type ContextScopeTagsExtractor func(context.Context) map[string]string
+
 const (
-	endpointKey        = contextFieldKey("endpoint")
-	requestUUIDKey     = contextFieldKey("requestUUID")
-	routingDelegateKey = contextFieldKey("rd")
-	requestLogFields   = contextFieldKey("requestLogFields")
+	endpointKey           = contextFieldKey("endpoint")
+	requestUUIDKey        = contextFieldKey("requestUUID")
+	routingDelegateKey    = contextFieldKey("rd")
+	endpointRequestHeader = contextFieldKey("endpointRequestHeader")
+	requestLogFields      = contextFieldKey("requestLogFields")
+	requestScopeFields    = contextFieldKey("requestScopeFields")
 )
 
 const (
@@ -49,6 +54,20 @@ const (
 	logFieldHandlerID           = "handlerID"
 )
 
+const (
+	scopeFieldRequestMethod = "method"
+	scopeFieldEndpointID    = "endpointID"
+	scopeFieldHandlerID     = "handlerID"
+)
+
+const (
+	// EndpointScope defines the name of endpoint scope
+	EndpointScope = "endpoint"
+
+	// ClientScope defines the name of client scope
+	ClientScope = "client"
+)
+
 // WithEndpointField adds the endpoint information in the
 // request context.
 func WithEndpointField(ctx context.Context, endpoint string) context.Context {
@@ -64,6 +83,30 @@ func GetRequestEndpointFromCtx(ctx context.Context) string {
 	return ""
 }
 
+// WithEndpointRequestHeadersField adds the endpoint request header information in the
+// request context.
+func WithEndpointRequestHeadersField(ctx context.Context, requestHeaders map[string]string) context.Context {
+	headers := GetEndpointRequestHeadersFromCtx(ctx)
+	for k, v := range requestHeaders {
+		headers[k] = v
+	}
+
+	return context.WithValue(ctx, endpointRequestHeader, headers)
+}
+
+// GetEndpointRequestHeadersFromCtx returns the endpoint request headers, if it exists on context
+func GetEndpointRequestHeadersFromCtx(ctx context.Context) map[string]string {
+	requestHeaders := make(map[string]string)
+	if val := ctx.Value(endpointRequestHeader); val != nil {
+		headers, _ := val.(map[string]string)
+		for k, v := range headers {
+			requestHeaders[k] = v
+		}
+	}
+
+	return requestHeaders
+}
+
 // withRequestFields annotates zanzibar request context to context.Context. In
 // future, we can use a request context struct to add more context in terms of
 // request handler, etc if need be.
@@ -105,6 +148,29 @@ func WithLogFields(ctx context.Context, newFields ...zap.Field) context.Context
 	return context.WithValue(ctx, requestLogFields, accumulateLogFields(ctx, newFields))
 }
 
+// WithScopeFields returns a new context with the given scope fields attached to context.Context
+func WithScopeFields(ctx context.Context, newFields map[string]string) context.Context {
+	fields := GetScopeFieldsFromCtx(ctx)
+	for k, v := range newFields {
+		fields[k] = v
+	}
+
+	return context.WithValue(ctx, requestScopeFields, fields)
+}
+
+// GetScopeFieldsFromCtx returns the tag info extracted from context.
+func GetScopeFieldsFromCtx(ctx context.Context) map[string]string {
+	fields := make(map[string]string)
+	if val := ctx.Value(requestScopeFields); val != nil {
+		headers, _ := val.(map[string]string)
+		for k, v := range headers {
+			fields[k] = v
+		}
+	}
+
+	return fields
+}
+
 func accumulateLogFields(ctx context.Context, newFields []zap.Field) []zap.Field {
 	previousFieldsUntyped := ctx.Value(requestLogFields)
 	if previousFieldsUntyped == nil {
@@ -124,6 +190,41 @@ func accumulateLogFields(ctx context.Context, newFields []zap.Field) []zap.Field
 	return fields
 }
 
+// ContextExtractor is a extractor that extracts some log fields from the context
+type ContextExtractor interface {
+	ExtractScopeTags(ctx context.Context) map[string]string
+}
+
+// AddContextScopeTagsExtractor added a scope tags extractor to contextExtractor.
+func (c *ContextExtractors) AddContextScopeTagsExtractor(extractors ...ContextScopeTagsExtractor) {
+	c.contextScopeExtractors = extractors
+}
+
+// MakeContextExtractor returns a extractor that extracts log fields a context.
+func (c *ContextExtractors) MakeContextExtractor() ContextExtractor {
+	return &ContextExtractors{
+		contextScopeExtractors: c.contextScopeExtractors,
+	}
+}
+
+// ContextExtractors warps extractors for context
+type ContextExtractors struct {
+	contextScopeExtractors []ContextScopeTagsExtractor
+}
+
+// ExtractScopeTags extracts scope fields from a context into a tag.
+func (c *ContextExtractors) ExtractScopeTags(ctx context.Context) map[string]string {
+	tags := make(map[string]string)
+	for _, fn := range c.contextScopeExtractors {
+		sc := fn(ctx)
+		for k, v := range sc {
+			tags[k] = v
+		}
+	}
+
+	return tags
+}
+
 // ContextLogger is a logger that extracts some log fields from the context before passing through to underlying zap logger.
 type ContextLogger interface {
 	Debug(ctx context.Context, msg string, fields ...zap.Field)

runtime/context_test.go

@@ -53,6 +53,52 @@ func TestGetRequestEndpointFromCtx(t *testing.T) {
 	assert.Equal(t, expected, endpoint)
 }
 
+func TestWithEndpointRequestHeadersField(t *testing.T) {
+	expected := map[string]string{"region": "san_francisco", "dc": "sjc1"}
+	ctx := WithEndpointRequestHeadersField(context.TODO(), expected)
+	rh := ctx.Value(endpointRequestHeader)
+	requestHeaders, ok := rh.(map[string]string)
+
+	assert.True(t, ok)
+	assert.Equal(t, requestHeaders, expected)
+}
+
+func TestGetEndpointRequestHeadersFromCtx(t *testing.T) {
+	expected := map[string]string{"region": "san_francisco", "dc": "sjc1"}
+	headers := map[string]string{"region": "san_francisco", "dc": "sjc1"}
+	ctx := WithEndpointRequestHeadersField(context.TODO(), headers)
+	requestHeaders := GetEndpointRequestHeadersFromCtx(ctx)
+	assert.Equal(t, expected, requestHeaders)
+
+	expected = map[string]string{}
+	ctx = context.TODO()
+	requestHeaders = GetEndpointRequestHeadersFromCtx(ctx)
+	assert.Equal(t, expected, requestHeaders)
+}
+
+func TestWithScopeFields(t *testing.T) {
+	expected := map[string]string{"endpoint": "tincup", "handler": "exchange"}
+	ctx := WithScopeFields(context.TODO(), expected)
+	rs := ctx.Value(requestScopeFields)
+	scopes, ok := rs.(map[string]string)
+
+	assert.True(t, ok)
+	assert.Equal(t, scopes, expected)
+}
+
+func TestGetScopeFieldsFromCtx(t *testing.T) {
+	expected := map[string]string{"endpoint": "tincup", "handler": "exchange"}
+	scope := map[string]string{"endpoint": "tincup", "handler": "exchange"}
+	ctx := WithScopeFields(context.TODO(), scope)
+	scopes := GetScopeFieldsFromCtx(ctx)
+	assert.Equal(t, expected, scopes)
+
+	expected = map[string]string{}
+	ctx = context.TODO()
+	scopes = GetScopeFieldsFromCtx(ctx)
+	assert.Equal(t, expected, scopes)
+}
+
 func TestWithRequestFields(t *testing.T) {
 	ctx := withRequestFields(context.TODO())
 
@@ -143,3 +189,22 @@ func TestContextLoggerPanic(t *testing.T) {
 
 	contextLogger.Panic(ctx, "msg", zap.String("argField", "argValue"))
 }
+
+func TestExtractScope(t *testing.T) {
+	headers := map[string]string{"x-uber-region-id": "san_francisco"}
+	ctx := WithEndpointRequestHeadersField(context.TODO(), headers)
+	contextScopeExtractors := []ContextScopeTagsExtractor{func(ctx context.Context) map[string]string {
+		headers := GetEndpointRequestHeadersFromCtx(ctx)
+		return map[string]string{"region-id": headers["x-uber-region-id"]}
+	}}
+
+	expected := map[string]string{"region-id": "san_francisco"}
+	contextExtractors := &ContextExtractors{}
+	for _, scopeExtractor := range contextScopeExtractors {
+		contextExtractors.AddContextScopeTagsExtractor(scopeExtractor)
+	}
+
+	contextExtractor := contextExtractors.MakeContextExtractor()
+	tags := contextExtractor.ExtractScopeTags(ctx)
+	assert.Equal(t, tags, expected)
+}