Merge pull request #730 from uber/lu.grpc

Support grpc client that has multiple services defined in proto

codegen/module_system.go

@@ -899,13 +899,24 @@ func (g *gRPCClientGenerator) Generate(
 
 	reversedMethods := reverseExposedMethods(clientSpec)
 
+	serviceNames := map[string]struct{}{}
+	for key := range reversedMethods {
+		serviceName := strings.Split(key, "::")[0]
+		serviceNames[serviceName] = struct{}{}
+	}
+
+	services := []*ServiceSpec{}
+	for name := range serviceNames {
+		services = append(services, &ServiceSpec{Name: name})
+	}
+
 	// @rpatali: Update all struct to use more general field IDLFile instead of thriftFile.
 	clientMeta := &ClientMeta{
 		ProtoServices:    clientSpec.ModuleSpec.ProtoServices,
 		Instance:         instance,
 		ExportName:       clientSpec.ExportName,
 		ExportType:       clientSpec.ExportType,
-		Services:         nil,
+		Services:         services,
 		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
 		ClientID:         clientSpec.ClientID,
 		ExposedMethods:   reversedMethods,

codegen/templates/grpc_client.tmpl

@@ -1,5 +1,6 @@
 {{- /* template to render gateway gRPC client code */ -}}
 {{- $instance := .Instance }}
+{{- $services := .Services }}
 package {{$instance.PackageInfo.PackageName}}
 
 import (
@@ -24,22 +25,27 @@ import (
 type Client interface {
 {{range $i, $svc := .ProtoServices -}}
 	{{range $j, $method := $svc.RPC}}
-		{{title $method.Name}} (
+	{{$serviceMethod := printf "%s::%s" $svc.Name .Name -}}
+	{{$methodName := (title (index $exposedMethods $serviceMethod)) -}}
+	{{- if $methodName -}}
+		{{$methodName}} (
 		ctx context.Context,
 		request *gen.{{$method.Request.Name}},
 		opts ...yarpc.CallOption,
 		) (*gen.{{$method.Response.Name}}, error)
 	{{ end -}}
+	{{ end -}}
+{{ end -}}
 }
 
 // {{$clientName}} is the gRPC client for downstream service.
 type {{$clientName}} struct {
-	client gen.{{pascal $svc.Name}}YARPCClient
+	{{range $i, $s := $services -}}
+	{{camel $s.Name}}Client gen.{{pascal $s.Name}}YARPCClient
+	{{ end -}}
 	opts   *zanzibar.GRPCClientOpts
 }
 
-{{- end}}
-
 // NewClient returns a new gRPC client for service {{$clientID}}
 func {{$exportName}}(deps *module.Dependencies) Client {
 	oc := deps.Default.GRPCClientDispatcher.MustOutboundConfig("{{$clientID}}")
@@ -55,30 +61,33 @@ func {{$exportName}}(deps *module.Dependencies) Client {
 	methodNames := map[string]string{
 		{{range $i, $svc := .ProtoServices -}}
 			{{range $j, $method := $svc.RPC -}}
-				"{{printf "%s::%s" $svc.Name $method.Name}}": "{{$method.Name}}",
+			{{$serviceMethod := printf "%s::%s" $svc.Name .Name -}}
+			{{$methodName := (title (index $exposedMethods $serviceMethod)) -}}
+			{{- if $methodName -}}
+				"{{$serviceMethod}}": "{{$methodName}}",
+			{{ end -}}
 			{{- end -}}
 		{{- end}}
 	}
-	{{range $i, $svc := .ProtoServices -}}
 	return &{{$clientName}}{
-		client: gen.New{{pascal $svc.Name}}YARPCClient(oc),
+		{{range $i, $s := $services -}}
+		{{camel $s.Name}}Client: gen.New{{pascal $s.Name}}YARPCClient(oc),
+		{{ end -}}
 		opts: zanzibar.NewGRPCClientOpts(
 		deps.Default.Logger,
 		deps.Default.ContextMetrics,
 		deps.Default.ContextExtractor,
 		methodNames,
 		"{{$clientID}}",
-		"{{$svc.Name}}",
 		routingKey,
 		requestUUIDHeaderKey,
-		configureCicruitBreaker(deps, timeoutInMS),
+		configureCircuitBreaker(deps, timeoutInMS),
 		timeoutInMS,
 		),
 	}
-	{{- end}}
 }
 
-func configureCicruitBreaker(deps *module.Dependencies, timeoutVal int) bool {
+func configureCircuitBreaker(deps *module.Dependencies, timeoutVal int) bool {
 	// circuitBreakerDisabled sets whether circuit-breaker should be disabled
 	circuitBreakerDisabled := false
 	if deps.Default.Config.ContainsKey("clients.{{$clientID}}.circuitBreakerDisabled") {
@@ -123,9 +132,11 @@ func configureCicruitBreaker(deps *module.Dependencies, timeoutVal int) bool {
 
 {{range $i, $svc := .ProtoServices -}}
 {{range $j, $method := $svc.RPC -}}
-{{if $method.Name -}}
-// {{$method.Name}} is a client RPC call for method {{printf "%s::%s" $svc.Name $method.Name}}.
-func (e *{{$clientName}}) {{$method.Name}}(
+{{$serviceMethod := printf "%s::%s" $svc.Name .Name -}}
+{{$methodName := (title (index $exposedMethods $serviceMethod)) -}}
+{{if $methodName -}}
+// {{$methodName}} is a client RPC call for method {{printf "%s::%s" $svc.Name $method.Name}}.
+func (e *{{$clientName}}) {{$methodName}}(
 	ctx context.Context,
 	request *gen.{{$method.Request.Name}},
 	opts ...yarpc.CallOption,
@@ -147,7 +158,7 @@ func (e *{{$clientName}}) {{$method.Name}}(
 	ctx, cancel := context.WithTimeout(ctx, e.opts.Timeout)
 	defer cancel()
 
-	runFunc := e.client.{{$method.Name}}
+	runFunc := e.{{camel $svc.Name}}Client.{{$method.Name}}
 	callHelper.Start()
 	if e.opts.CircuitBreakerDisabled {
 		result, err = runFunc(ctx, request, opts...)

examples/example-gateway/endpoints/bounce/bounce.go

@@ -48,7 +48,7 @@ func (w bounceWorkflow) Handle(
 	reqHeaders zanzibar.Header,
 	req *bounce.Bounce_Bounce_Args,
 ) (string, zanzibar.Header, error) {
-	res, err := w.echo.Echo(ctx, &echo.Request{Message: req.Msg})
+	res, err := w.echo.EchoEcho(ctx, &echo.Request{Message: req.Msg})
 	if err != nil {
 		return "", nil, err
 	}

examples/example-gateway/endpoints/bounce/bounce_test.go

@@ -25,7 +25,7 @@ func TestEcho(t *testing.T) {
 	ctx := context.Background()
 	var result bounce.Bounce_Bounce_Result
 
-	ms.MockClients().Echo.EXPECT().Echo(gomock.Any(), &echo.Request{Message: message}).
+	ms.MockClients().Echo.EXPECT().EchoEcho(gomock.Any(), &echo.Request{Message: message}).
 		Return(&echo.Response{Message: message}, nil)
 
 	success, resHeaders, err := ms.MakeTChannelRequest(