Merge 2fc0a74 into 9f97351

codegen/module_system.go

@@ -50,6 +50,7 @@ type EndpointMeta struct {
 	ResRequiredHeadersKeys []string
 	TraceKey               string
 	DeputyReqHeader        string
+	ShadowReqHeader        string
 }
 
 // EndpointCollectionMeta saves information used to generate an initializer
@@ -548,6 +549,7 @@ func (g *TChannelClientGenerator) Generate(
 		SidecarRouter:    clientSpec.SidecarRouter,
 		StagingReqHeader: g.packageHelper.StagingReqHeader(),
 		DeputyReqHeader:  g.packageHelper.DeputyReqHeader(),
+		ShadowReqHeader:  g.packageHelper.ShadowReqHeader(),
 	}
 
 	client, err := g.templates.ExecTemplate(
@@ -916,6 +918,13 @@ func (g *EndpointGenerator) generateEndpointFile(
 		TransformTo: shk,
 		Field:       &compile.FieldSpec{Required: false},
 	}
+
+	shk = textproto.CanonicalMIMEHeaderKey(g.packageHelper.ShadowReqHeader())
+	reqHeaders[shk] = &TypedHeader{
+		Name:        shk,
+		TransformTo: shk,
+		Field:       &compile.FieldSpec{Required: false},
+	}
 	// TODO: http client needs to support multiple thrift services
 	meta := &EndpointMeta{
 		Instance:               instance,
@@ -935,6 +944,7 @@ func (g *EndpointGenerator) generateEndpointFile(
 		WorkflowPkg:            workflowPkg,
 		TraceKey:               g.packageHelper.traceKey,
 		DeputyReqHeader:        g.packageHelper.DeputyReqHeader(),
+		ShadowReqHeader:        g.packageHelper.ShadowReqHeader(),
 	}
 
 	var endpoint []byte
@@ -1276,6 +1286,7 @@ type ClientMeta struct {
 	Fixture          *Fixture
 	StagingReqHeader string
 	DeputyReqHeader  string
+	ShadowReqHeader  string
 }
 
 func findMethod(

codegen/package.go

@@ -55,6 +55,8 @@ type PackageHelper struct {
 	stagingReqHeader string
 	// Use deputy client when this header is set
 	deputyReqHeader string
+	// This is a shadow request - call Apprentice instead of a downstream client
+	shadowReqHeader string
 	// traceKey is the key for unique trace id that identifies request / response pair
 	traceKey string
 }
@@ -86,6 +88,8 @@ type PackageHelperOptions struct {
 	StagingReqHeader string
 	// header key to redirect client requests to local environment, defaults to "x-deputy-forwarded"
 	DeputyReqHeader string
+	// header key to redirect client request to apprentice, defaults to "X-Uber-Should-Request-Downstream"
+	ShadowReqHeader string
 	// header key to uniquely identifies request/response pair, defaults to "x-trace-id"
 	TraceKey string
 }
@@ -146,6 +150,13 @@ func (p *PackageHelperOptions) deputyReqHeader() string {
 	return "x-deputy-forwarded"
 }
 
+func (p *PackageHelperOptions) shadowReqHeader() string {
+	if p.ShadowReqHeader != "" {
+		return p.ShadowReqHeader
+	}
+	return "X-Shadow-Start-Trace-Id"
+}
+
 func (p *PackageHelperOptions) traceKey() string {
 	if p.TraceKey != "" {
 		return p.TraceKey
@@ -192,6 +203,7 @@ func NewPackageHelper(
 		annotationPrefix:   options.annotationPrefix(),
 		stagingReqHeader:   options.stagingReqHeader(),
 		deputyReqHeader:    options.deputyReqHeader(),
+		shadowReqHeader:    options.shadowReqHeader(),
 		traceKey:           options.traceKey(),
 	}
 	return p, nil
@@ -319,3 +331,9 @@ func (p PackageHelper) StagingReqHeader() string {
 func (p PackageHelper) DeputyReqHeader() string {
 	return p.deputyReqHeader
 }
+
+// ShadowReqHeader returns the header name that will be checked to determine
+// if a request should go to apprentice or downstream client
+func (p PackageHelper) ShadowReqHeader() string {
+	return p.shadowReqHeader
+}

codegen/runner/runner.go

@@ -89,6 +89,10 @@ func main() {
 	if config.ContainsKey("deputyReqHeader") {
 		deputyReqHeader = config.MustGetString("deputyReqHeader")
 	}
+	shadowReqHeader := "X-Shadow-Start-Trace-Id"
+	if config.ContainsKey("shadowReqHeader") {
+		shadowReqHeader = config.MustGetString("shadowReqHeader")
+	}
 	options := &codegen.PackageHelperOptions{
 		RelThriftRootDir:       config.MustGetString("thriftRootDir"),
 		RelTargetGenDir:        config.MustGetString("targetGenDir"),
@@ -98,6 +102,7 @@ func main() {
 		CopyrightHeader:        string(copyright),
 		StagingReqHeader:       stagingReqHeader,
 		DeputyReqHeader:        deputyReqHeader,
+		ShadowReqHeader:        shadowReqHeader,
 		TraceKey:               config.MustGetString("traceKey"),
 	}
 

codegen/templates/http_client.tmpl

@@ -47,6 +47,9 @@ type Client interface {
 type {{$clientName}} struct {
 	clientID string
 	httpClient   *zanzibar.HTTPClient
+	{{if ne $clientID "apprentice" -}}
+	apprenticeClient *zanzibar.TChannelClient
+	{{end -}}
 
 	{{if $sidecarRouter -}}
 	calleeHeader string
@@ -76,6 +79,44 @@ func {{$exportName}}(deps *module.Dependencies) Client {
 		deps.Default.Config.MustGetStruct("clients.{{$clientID}}.defaultHeaders", &defaultHeaders)
 	}
 
+	{{if ne $clientID "apprentice" -}}
+	apprenticeMethodNames := map[string]string{
+		"Apprentice::getRequest": "getRequest",
+		"Apprentice::getResponse": "getResponse",
+	}
+
+	apprenticeTimeout := time.Millisecond * time.Duration(50)
+	if deps.Default.Config.ContainsKey("clients.apprentice.timeout") {
+		apprenticeTimeout = time.Millisecond * time.Duration(
+			deps.Default.Config.MustGetInt("clients.apprentice.timeout"),
+		)
+	}
+
+	apprenticeTimeoutPerAttempt := time.Millisecond * time.Duration(50)
+	if deps.Default.Config.ContainsKey("clients.apprentice.timeoutPerAttempt") {
+		apprenticeTimeoutPerAttempt = time.Millisecond * time.Duration(
+			deps.Default.Config.MustGetInt("clients.apprentice.timeoutPerAttempt"),
+		)
+	}
+
+	apprenticeRoutingKey := ""
+
+	apprenticeClient := zanzibar.NewTChannelClient(
+		deps.Default.Channel,
+		deps.Default.Logger,
+		deps.Default.Scope,
+		&zanzibar.TChannelClientOption{
+			ServiceName:       "apprentice",
+			ClientID:          "apprentice",
+			MethodNames:       apprenticeMethodNames,
+			Timeout:           apprenticeTimeout,
+			TimeoutPerAttempt: apprenticeTimeoutPerAttempt,
+			RoutingKey:        &apprenticeRoutingKey,
+			AltSubchannelName: "apprentice",
+		},
+	)
+	{{end}}
+
 	return &{{$clientName}}{
 		clientID: "{{$clientID}}",
 		{{if $sidecarRouter -}}
@@ -96,6 +137,9 @@ func {{$exportName}}(deps *module.Dependencies) Client {
 			defaultHeaders,
 			timeout,
 		),
+		{{if ne $clientID "apprentice" -}}
+		apprenticeClient: apprenticeClient,
+		{{end -}}
 	}
 }
 

codegen/templates/tchannel_client.tmpl

@@ -117,14 +117,58 @@ func {{$exportName}}(deps *module.Dependencies) Client {
 		},
 	)
 
+	{{if ne $clientID "apprentice" -}}
+	apprenticeMethodNames := map[string]string{
+		"Apprentice::getRequest": "getRequest",
+		"Apprentice::getResponse": "getResponse",
+	}
+
+	apprenticeTimeout := time.Millisecond * time.Duration(50)
+	if deps.Default.Config.ContainsKey("clients.apprentice.timeout") {
+		apprenticeTimeout = time.Millisecond * time.Duration(
+			deps.Default.Config.MustGetInt("clients.apprentice.timeout"),
+		)
+	}
+
+	apprenticeTimeoutPerAttempt := time.Millisecond * time.Duration(50)
+	if deps.Default.Config.ContainsKey("clients.apprentice.timeoutPerAttempt") {
+		apprenticeTimeoutPerAttempt = time.Millisecond * time.Duration(
+			deps.Default.Config.MustGetInt("clients.apprentice.timeoutPerAttempt"),
+		)
+	}
+
+	apprenticeRoutingKey := ""
+
+	apprenticeClient := zanzibar.NewTChannelClient(
+		deps.Default.Channel,
+		deps.Default.Logger,
+		deps.Default.Scope,
+		&zanzibar.TChannelClientOption{
+			ServiceName:       "apprentice",
+			ClientID:          "apprentice",
+			MethodNames:       apprenticeMethodNames,
+			Timeout:           apprenticeTimeout,
+			TimeoutPerAttempt: apprenticeTimeoutPerAttempt,
+			RoutingKey:        &apprenticeRoutingKey,
+			AltSubchannelName: "apprentice",
+		},
+	)
+	{{end}}
+
 	return &{{$clientName}}{
 		client: client,
+		{{if ne $clientID "apprentice" -}}
+		apprenticeClient: apprenticeClient,
+		{{end -}}
 	}
 }
 
 // {{$clientName}} is the TChannel client for downstream service.
 type {{$clientName}} struct {
 	client *zanzibar.TChannelClient
+	{{if ne $clientID "apprentice" -}}
+	apprenticeClient *zanzibar.TChannelClient
+	{{end -}}
 }
 
 {{range $svc := .Services}}