go.uber.org/yarpc@v1.72.1/encoding/protobuf/protoc-gen-yarpc-go-v2/internal/lib/lib.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  // Package lib contains the library code for protoc-gen-yarpc-go-v2.
    22  //
    23  // It is split into a separate package so it can be called by the testing package.
    24  package lib
    25  
    26  import (
    27  	"bytes"
    28  	"crypto/sha256"
    29  	"encoding/hex"
    30  	"fmt"
    31  	"path/filepath"
    32  	"strings"
    33  	"text/template"
    34  
    35  	protoplugin "go.uber.org/yarpc/internal/protoplugin-v2"
    36  )
    37  
    38  const tmpl = `{{$packagePath := .GoPackage.Path}}{{$packageName := .GoPackage.Name}}
    39  // Code generated by protoc-gen-yarpc-go. DO NOT EDIT.
    40  // source: {{.GetName}}
    41  
    42  package {{$packageName}}
    43  {{if .Services}}
    44  import (
    45  	{{range $i := .Imports}}{{if $i.Standard}}{{$i | printf "%s\n"}}{{end}}{{end}}
    46  
    47  	{{range $i := .Imports}}{{if not $i.Standard}}{{$i | printf "%s\n"}}{{end}}{{end}}
    48  ){{end}}
    49  
    50  {{if ne (len .Services) 0}}var _ = ioutil.NopCloser{{end}}
    51  
    52  {{range $service := .Services}}
    53  // {{$service.GetName}}YARPCClient is the YARPC client-side interface for the {{$service.GetName}} service.
    54  type {{$service.GetName}}YARPCClient interface {
    55  	{{range $method := unaryMethods $service}}{{$method.GetName}}(context.Context, *{{$method.RequestType.GoType $packagePath}}, ...yarpc.CallOption) (*{{$method.ResponseType.GoType $packagePath}}, error)
    56  	{{end}}{{range $method := onewayMethods $service}}{{$method.GetName}}(context.Context, *{{$method.RequestType.GoType $packagePath}}, ...yarpc.CallOption) (yarpc.Ack, error)
    57  	{{end}}{{range $method := clientStreamingMethods $service}}{{$method.GetName}}(context.Context, ...yarpc.CallOption) ({{$service.GetName}}Service{{$method.GetName}}YARPCClient, error)
    58  	{{end}}{{range $method := serverStreamingMethods $service}}{{$method.GetName}}(context.Context, *{{$method.RequestType.GoType $packagePath}}, ...yarpc.CallOption) ({{$service.GetName}}Service{{$method.GetName}}YARPCClient, error)
    59  	{{end}}{{range $method := clientServerStreamingMethods $service}}{{$method.GetName}}(context.Context, ...yarpc.CallOption) ({{$service.GetName}}Service{{$method.GetName}}YARPCClient, error)
    60  	{{end}}
    61  }
    62  
    63  {{range $method := clientStreamingMethods $service}}
    64  // {{$service.GetName}}Service{{$method.GetName}}YARPCClient sends {{$method.RequestType.GoType $packagePath}}s and receives the single {{$method.ResponseType.GoType $packagePath}} when sending is done.
    65  type {{$service.GetName}}Service{{$method.GetName}}YARPCClient interface {
    66  	Context() context.Context
    67  	Send(*{{$method.RequestType.GoType $packagePath}}, ...yarpc.StreamOption) error
    68  	CloseAndRecv(...yarpc.StreamOption) (*{{$method.ResponseType.GoType $packagePath}}, error)
    69  }
    70  {{end}}
    71  
    72  {{range $method := serverStreamingMethods $service}}
    73  // {{$service.GetName}}Service{{$method.GetName}}YARPCClient receives {{$method.ResponseType.GoType $packagePath}}s, returning io.EOF when the stream is complete.
    74  type {{$service.GetName}}Service{{$method.GetName}}YARPCClient interface {
    75  	Context() context.Context
    76  	Recv(...yarpc.StreamOption) (*{{$method.ResponseType.GoType $packagePath}}, error)
    77  	CloseSend(...yarpc.StreamOption) error
    78  }
    79  {{end}}
    80  
    81  {{range $method := clientServerStreamingMethods $service}}
    82  // {{$service.GetName}}Service{{$method.GetName}}YARPCClient sends {{$method.RequestType.GoType $packagePath}}s and receives {{$method.ResponseType.GoType $packagePath}}s, returning io.EOF when the stream is complete.
    83  type {{$service.GetName}}Service{{$method.GetName}}YARPCClient interface {
    84  	Context() context.Context
    85  	Send(*{{$method.RequestType.GoType $packagePath}}, ...yarpc.StreamOption) error
    86  	Recv(...yarpc.StreamOption) (*{{$method.ResponseType.GoType $packagePath}}, error)
    87  	CloseSend(...yarpc.StreamOption) error
    88  }
    89  {{end}}
    90  
    91  func new{{$service.GetName}}YARPCClient(clientConfig transport.ClientConfig, anyResolver v2.AnyResolver, options ...v2.ClientOption) {{$service.GetName}}YARPCClient {
    92  	return &_{{$service.GetName}}YARPCCaller{v2.NewStreamClient(
    93  		v2.ClientParams{
    94  			ServiceName: "{{trimPrefixPeriod $service.FQSN}}",
    95  			ClientConfig: clientConfig,
    96  			AnyResolver: anyResolver,
    97  			Options: options,
    98  		},
    99  	)}
   100  }
   101  
   102  // New{{$service.GetName}}YARPCClient builds a new YARPC client for the {{$service.GetName}} service.
   103  func New{{$service.GetName}}YARPCClient(clientConfig transport.ClientConfig, options ...v2.ClientOption) {{$service.GetName}}YARPCClient {
   104  	return new{{$service.GetName}}YARPCClient(clientConfig, nil, options...)
   105  }
   106  
   107  // {{$service.GetName}}YARPCServer is the YARPC server-side interface for the {{$service.GetName}} service.
   108  type {{$service.GetName}}YARPCServer interface {
   109  	{{range $method := unaryMethods $service}}{{$method.GetName}}(context.Context, *{{$method.RequestType.GoType $packagePath}}) (*{{$method.ResponseType.GoType $packagePath}}, error)
   110  	{{end}}{{range $method := onewayMethods $service}}{{$method.GetName}}(context.Context, *{{$method.RequestType.GoType $packagePath}}) error
   111  	{{end}}{{range $method := clientStreamingMethods $service}}{{$method.GetName}}({{$service.GetName}}Service{{$method.GetName}}YARPCServer) (*{{$method.ResponseType.GoType $packagePath}}, error)
   112  	{{end}}{{range $method := serverStreamingMethods $service}}{{$method.GetName}}(*{{$method.RequestType.GoType $packagePath}}, {{$service.GetName}}Service{{$method.GetName}}YARPCServer) error
   113  	{{end}}{{range $method := clientServerStreamingMethods $service}}{{$method.GetName}}({{$service.GetName}}Service{{$method.GetName}}YARPCServer) error
   114  	{{end}}
   115  }
   116  
   117  {{range $method := clientStreamingMethods $service}}
   118  // {{$service.GetName}}Service{{$method.GetName}}YARPCServer receives {{$method.RequestType.GoType $packagePath}}s.
   119  type {{$service.GetName}}Service{{$method.GetName}}YARPCServer interface {
   120  	Context() context.Context
   121  	Recv(...yarpc.StreamOption) (*{{$method.RequestType.GoType $packagePath}}, error)
   122  }
   123  {{end}}
   124  
   125  {{range $method := serverStreamingMethods $service}}
   126  // {{$service.GetName}}Service{{$method.GetName}}YARPCServer sends {{$method.ResponseType.GoType $packagePath}}s.
   127  type {{$service.GetName}}Service{{$method.GetName}}YARPCServer interface {
   128  	Context() context.Context
   129  	Send(*{{$method.ResponseType.GoType $packagePath}}, ...yarpc.StreamOption) error
   130  }
   131  {{end}}
   132  
   133  {{range $method := clientServerStreamingMethods $service}}
   134  // {{$service.GetName}}Service{{$method.GetName}}YARPCServer receives {{$method.RequestType.GoType $packagePath}}s and sends {{$method.ResponseType.GoType $packagePath}}.
   135  type {{$service.GetName}}Service{{$method.GetName}}YARPCServer interface {
   136  	Context() context.Context
   137  	Recv(...yarpc.StreamOption) (*{{$method.RequestType.GoType $packagePath}}, error)
   138  	Send(*{{$method.ResponseType.GoType $packagePath}}, ...yarpc.StreamOption) error
   139  }
   140  {{end}}
   141  
   142  type build{{$service.GetName}}YARPCProceduresParams struct {
   143  	Server      {{$service.GetName}}YARPCServer
   144  	AnyResolver v2.AnyResolver
   145  }
   146  
   147  func build{{$service.GetName}}YARPCProcedures(params build{{$service.GetName}}YARPCProceduresParams) []transport.Procedure {
   148  	handler := &_{{$service.GetName}}YARPCHandler{params.Server}
   149  	return v2.BuildProcedures(
   150  		v2.BuildProceduresParams{
   151  			ServiceName: "{{trimPrefixPeriod $service.FQSN}}",
   152  			UnaryHandlerParams: []v2.BuildProceduresUnaryHandlerParams{
   153  			{{range $method := unaryMethods $service}}{
   154  					MethodName: "{{$method.GetName}}",
   155  					Handler: v2.NewUnaryHandler(
   156  						v2.UnaryHandlerParams{
   157  							Handle: handler.{{$method.GetName}},
   158  							NewRequest: new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest,
   159  							AnyResolver: params.AnyResolver,
   160  						},
   161  					),
   162  				},
   163  			{{end}}
   164  			},
   165  			OnewayHandlerParams: []v2.BuildProceduresOnewayHandlerParams{
   166  			{{range $method := onewayMethods $service}}{
   167  					MethodName: "{{$method.GetName}}",
   168  					Handler: v2.NewOnewayHandler(
   169  						v2.OnewayHandlerParams{
   170  							Handle: handler.{{$method.GetName}},
   171  							NewRequest: new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest,
   172  						},
   173  					),
   174  				},
   175  			{{end}}
   176  			},
   177  			StreamHandlerParams: []v2.BuildProceduresStreamHandlerParams{
   178  			{{range $method := clientServerStreamingMethods $service}}{
   179  					MethodName: "{{$method.GetName}}",
   180  					Handler: v2.NewStreamHandler(
   181  						v2.StreamHandlerParams{
   182  							Handle: handler.{{$method.GetName}},
   183  						},
   184  					),
   185  				},
   186  			{{end}}
   187  			{{range $method := serverStreamingMethods $service}}{
   188  					MethodName: "{{$method.GetName}}",
   189  					Handler: v2.NewStreamHandler(
   190  						v2.StreamHandlerParams{
   191  							Handle: handler.{{$method.GetName}},
   192  						},
   193  					),
   194  				},
   195  			{{end}}
   196  			{{range $method := clientStreamingMethods $service}}{
   197  					MethodName: "{{$method.GetName}}",
   198  					Handler: v2.NewStreamHandler(
   199  						v2.StreamHandlerParams{
   200  							Handle: handler.{{$method.GetName}},
   201  						},
   202  					),
   203  				},
   204  			{{end}}
   205  			},
   206  		},
   207  	)
   208  }
   209  
   210  // Build{{$service.GetName}}YARPCProcedures prepares an implementation of the {{$service.GetName}} service for YARPC registration.
   211  func Build{{$service.GetName}}YARPCProcedures(server {{$service.GetName}}YARPCServer) []transport.Procedure {
   212  	return build{{$service.GetName}}YARPCProcedures(build{{$service.GetName}}YARPCProceduresParams{Server:server})
   213  }
   214  
   215  // Fx{{$service.GetName}}YARPCClientParams defines the input
   216  // for NewFx{{$service.GetName}}YARPCClient. It provides the
   217  // paramaters to get a {{$service.GetName}}YARPCClient in an
   218  // Fx application.
   219  type Fx{{$service.GetName}}YARPCClientParams struct {
   220  	fx.In
   221  
   222  	Provider yarpc.ClientConfig
   223  	AnyResolver v2.AnyResolver ` + "`" + `name:"yarpcfx" optional:"true"` + "`" + `
   224  	Restriction restriction.Checker ` + "`" + `optional:"true"` + "`" + `
   225  }
   226  
   227  // Fx{{$service.GetName}}YARPCClientResult defines the output
   228  // of NewFx{{$service.GetName}}YARPCClient. It provides a
   229  // {{$service.GetName}}YARPCClient to an Fx application.
   230  type Fx{{$service.GetName}}YARPCClientResult struct {
   231  	fx.Out
   232  
   233  	Client {{$service.GetName}}YARPCClient
   234  
   235  	// We are using an fx.Out struct here instead of just returning a client
   236  	// so that we can add more values or add named versions of the client in
   237  	// the future without breaking any existing code.
   238  }
   239  
   240  // NewFx{{$service.GetName}}YARPCClient provides a {{$service.GetName}}YARPCClient
   241  // to an Fx application using the given name for routing.
   242  //
   243  //  fx.Provide(
   244  //    {{$packageName}}.NewFx{{$service.GetName}}YARPCClient("service-name"),
   245  //    ...
   246  //  )
   247  func NewFx{{$service.GetName}}YARPCClient(name string, options ...v2.ClientOption) interface{} {
   248  	return func(params Fx{{$service.GetName}}YARPCClientParams) Fx{{$service.GetName}}YARPCClientResult {
   249  		cc := params.Provider.ClientConfig(name)
   250  
   251  		if params.Restriction != nil{
   252  			if namer, ok := cc.GetUnaryOutbound().(transport.Namer); ok {
   253  				if err := params.Restriction.Check(v2.Encoding, namer.TransportName()); err != nil {
   254  					panic(err.Error())
   255  				}
   256  			}
   257  		}
   258  
   259  		return Fx{{$service.GetName}}YARPCClientResult{
   260  			Client: new{{$service.GetName}}YARPCClient(cc, params.AnyResolver, options...),
   261  		}
   262  	}
   263  }
   264  
   265  // Fx{{$service.GetName}}YARPCProceduresParams defines the input
   266  // for NewFx{{$service.GetName}}YARPCProcedures. It provides the
   267  // paramaters to get {{$service.GetName}}YARPCServer procedures in an
   268  // Fx application.
   269  type Fx{{$service.GetName}}YARPCProceduresParams struct {
   270  	fx.In
   271  
   272  	Server {{$service.GetName}}YARPCServer
   273  	AnyResolver v2.AnyResolver ` + "`" + `name:"yarpcfx" optional:"true"` + "`" + `
   274  }
   275  
   276  // Fx{{$service.GetName}}YARPCProceduresResult defines the output
   277  // of NewFx{{$service.GetName}}YARPCProcedures. It provides
   278  // {{$service.GetName}}YARPCServer procedures to an Fx application.
   279  //
   280  // The procedures are provided to the "yarpcfx" value group.
   281  // Dig 1.2 or newer must be used for this feature to work.
   282  type Fx{{$service.GetName}}YARPCProceduresResult struct {
   283  	fx.Out
   284  
   285  	Procedures []transport.Procedure ` + "`" + `group:"yarpcfx"` + "`" + `
   286  	ReflectionMeta reflection.ServerMeta ` + "`" + `group:"yarpcfx"` + "`" + `
   287  }
   288  
   289  // NewFx{{$service.GetName}}YARPCProcedures provides {{$service.GetName}}YARPCServer procedures to an Fx application.
   290  // It expects a {{$service.GetName}}YARPCServer to be present in the container.
   291  //
   292  //  fx.Provide(
   293  //    {{$packageName}}.NewFx{{$service.GetName}}YARPCProcedures(),
   294  //    ...
   295  //  )
   296  func NewFx{{$service.GetName}}YARPCProcedures() interface{} {
   297  	return func(params Fx{{$service.GetName}}YARPCProceduresParams) Fx{{$service.GetName}}YARPCProceduresResult {
   298  		return Fx{{$service.GetName}}YARPCProceduresResult{
   299  			Procedures: build{{$service.GetName}}YARPCProcedures(build{{$service.GetName}}YARPCProceduresParams{
   300  				Server:      params.Server,
   301  				AnyResolver: params.AnyResolver,
   302  			}),
   303  			ReflectionMeta: reflection.ServerMeta{
   304  				ServiceName: "{{trimPrefixPeriod $service.FQSN}}",
   305  				FileDescriptors: {{ fileDescriptorClosureVarName .File }},
   306  			},
   307  		}
   308  	}
   309  }
   310  
   311  type _{{$service.GetName}}YARPCCaller struct {
   312  	streamClient v2.StreamClient
   313  }
   314  
   315  {{range $method := unaryMethods $service}}
   316  func (c *_{{$service.GetName}}YARPCCaller) {{$method.GetName}}(ctx context.Context, request *{{$method.RequestType.GoType $packagePath}}, options ...yarpc.CallOption) (*{{$method.ResponseType.GoType $packagePath}}, error) {
   317  	responseMessage, err := c.streamClient.Call(ctx, "{{$method.GetName}}", request, new{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, options...)
   318  	if responseMessage == nil {
   319  		return nil, err
   320  	}
   321  	response, ok := responseMessage.(*{{$method.ResponseType.GoType $packagePath}})
   322  	if !ok {
   323  		return nil, v2.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, responseMessage)
   324  	}
   325  	return response, err
   326  }
   327  {{end}}
   328  {{range $method := onewayMethods $service}}
   329  func (c *_{{$service.GetName}}YARPCCaller) {{$method.GetName}}(ctx context.Context, request *{{$method.RequestType.GoType $packagePath}}, options ...yarpc.CallOption) (yarpc.Ack, error) {
   330  	return c.streamClient.CallOneway(ctx, "{{$method.GetName}}", request, options...)
   331  }
   332  {{end}}
   333  {{range $method := clientStreamingMethods $service}}
   334  func (c *_{{$service.GetName}}YARPCCaller) {{$method.GetName}}(ctx context.Context, options ...yarpc.CallOption) ({{$service.GetName}}Service{{$method.GetName}}YARPCClient, error) {
   335  	stream, err := c.streamClient.CallStream(ctx, "{{$method.GetName}}", options...)
   336  	if err != nil {
   337  		return nil, err
   338  	}
   339  	return &_{{$service.GetName}}Service{{$method.GetName}}YARPCClient{stream: stream}, nil
   340  }
   341  {{end}}
   342  {{range $method := serverStreamingMethods $service}}
   343  func (c *_{{$service.GetName}}YARPCCaller) {{$method.GetName}}(ctx context.Context, request *{{$method.RequestType.GoType $packagePath}}, options ...yarpc.CallOption) ({{$service.GetName}}Service{{$method.GetName}}YARPCClient, error) {
   344  	stream, err := c.streamClient.CallStream(ctx, "{{$method.GetName}}", options...)
   345  	if err != nil {
   346  		return nil, err
   347  	}
   348  	if err := stream.Send(request); err != nil {
   349  		return nil, err
   350  	}
   351  	return &_{{$service.GetName}}Service{{$method.GetName}}YARPCClient{stream: stream}, nil
   352  }
   353  {{end}}
   354  {{range $method := clientServerStreamingMethods $service}}
   355  func (c *_{{$service.GetName}}YARPCCaller) {{$method.GetName}}(ctx context.Context, options ...yarpc.CallOption) ({{$service.GetName}}Service{{$method.GetName}}YARPCClient, error) {
   356  	stream, err := c.streamClient.CallStream(ctx, "{{$method.GetName}}", options...)
   357  	if err != nil {
   358  		return nil, err
   359  	}
   360  	return &_{{$service.GetName}}Service{{$method.GetName}}YARPCClient{stream: stream}, nil
   361  }
   362  {{end}}
   363  
   364  type _{{$service.GetName}}YARPCHandler struct {
   365  	server {{$service.GetName}}YARPCServer
   366  }
   367  
   368  {{range $method := unaryMethods $service}}
   369  func (h *_{{$service.GetName}}YARPCHandler) {{$method.GetName}}(ctx context.Context, requestMessage proto.Message) (proto.Message, error) {
   370  	var request *{{$method.RequestType.GoType $packagePath}}
   371  	var ok bool
   372  	if requestMessage != nil {
   373  		request, ok = requestMessage.(*{{$method.RequestType.GoType $packagePath}})
   374  		if !ok {
   375  			return nil, v2.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, requestMessage)
   376  		}
   377  	}
   378  	response, err := h.server.{{$method.GetName}}(ctx, request)
   379  	if response == nil {
   380  		return nil, err
   381  	}
   382  	return response, err
   383  }
   384  {{end}}
   385  {{range $method := onewayMethods $service}}
   386  func (h *_{{$service.GetName}}YARPCHandler) {{$method.GetName}}(ctx context.Context, requestMessage proto.Message) error {
   387  	var request *{{$method.RequestType.GoType $packagePath}}
   388  	var ok bool
   389  	if requestMessage != nil {
   390  		request, ok = requestMessage.(*{{$method.RequestType.GoType $packagePath}})
   391  		if !ok {
   392  			return v2.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, requestMessage)
   393  		}
   394  	}
   395  	return h.server.{{$method.GetName}}(ctx, request)
   396  }
   397  {{end}}
   398  {{range $method := clientStreamingMethods $service}}
   399  func (h *_{{$service.GetName}}YARPCHandler) {{$method.GetName}}(serverStream *v2.ServerStream) error {
   400  	response, err := h.server.{{$method.GetName}}(&_{{$service.GetName}}Service{{$method.GetName}}YARPCServer{serverStream: serverStream})
   401  	if err != nil {
   402  		return err
   403  	}
   404  	return serverStream.Send(response)
   405  }
   406  {{end}}
   407  {{range $method := serverStreamingMethods $service}}
   408  func (h *_{{$service.GetName}}YARPCHandler) {{$method.GetName}}(serverStream *v2.ServerStream) error {
   409  	requestMessage, err := serverStream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest)
   410  	if requestMessage == nil {
   411          return err
   412      }
   413  
   414  	request, ok := requestMessage.(*{{$method.RequestType.GoType $packagePath}})
   415  	if !ok {
   416  		return v2.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, requestMessage)
   417  	}
   418  	return h.server.{{$method.GetName}}(request, &_{{$service.GetName}}Service{{$method.GetName}}YARPCServer{serverStream: serverStream})
   419  }
   420  {{end}}
   421  {{range $method := clientServerStreamingMethods $service}}
   422  func (h *_{{$service.GetName}}YARPCHandler) {{$method.GetName}}(serverStream *v2.ServerStream) error {
   423  	return h.server.{{$method.GetName}}(&_{{$service.GetName}}Service{{$method.GetName}}YARPCServer{serverStream: serverStream})
   424  }
   425  {{end}}
   426  
   427  {{range $method := clientStreamingMethods $service}}
   428  type _{{$service.GetName}}Service{{$method.GetName}}YARPCClient struct {
   429  	stream *v2.ClientStream
   430  }
   431  
   432  func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Context() context.Context {
   433  	return c.stream.Context()
   434  }
   435  
   436  func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Send(request *{{$method.RequestType.GoType $packagePath}}, options ...yarpc.StreamOption) error {
   437  	return c.stream.Send(request, options...)
   438  }
   439  
   440  func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) CloseAndRecv(options ...yarpc.StreamOption) (*{{$method.ResponseType.GoType $packagePath}}, error) {
   441  	if err := c.stream.Close(options...); err != nil {
   442  		return nil, err
   443  	}
   444  	responseMessage, err := c.stream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, options...)
   445  	if responseMessage == nil {
   446          return nil, err
   447      }
   448  	response, ok := responseMessage.(*{{$method.ResponseType.GoType $packagePath}})
   449  	if !ok {
   450  		return nil, v2.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, responseMessage)
   451  	}
   452  	return response, err
   453  }
   454  {{end}}
   455  
   456  {{range $method := serverStreamingMethods $service}}
   457  type _{{$service.GetName}}Service{{$method.GetName}}YARPCClient struct {
   458  	stream *v2.ClientStream
   459  }
   460  
   461  func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Context() context.Context {
   462  	return c.stream.Context()
   463  }
   464  
   465  func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Recv(options ...yarpc.StreamOption) (*{{$method.ResponseType.GoType $packagePath}}, error) {
   466  	responseMessage, err := c.stream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, options...)
   467  	if responseMessage == nil {
   468          return nil, err
   469      }
   470  	response, ok := responseMessage.(*{{$method.ResponseType.GoType $packagePath}})
   471  	if !ok {
   472  		return nil, v2.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, responseMessage)
   473  	}
   474  	return response, err
   475  }
   476  
   477  func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) CloseSend(options ...yarpc.StreamOption) error {
   478  	return c.stream.Close(options...)
   479  }
   480  {{end}}
   481  
   482  {{range $method := clientServerStreamingMethods $service}}
   483  type _{{$service.GetName}}Service{{$method.GetName}}YARPCClient struct {
   484  	stream *v2.ClientStream
   485  }
   486  
   487  func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Context() context.Context {
   488  	return c.stream.Context()
   489  }
   490  
   491  func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Send(request *{{$method.RequestType.GoType $packagePath}}, options ...yarpc.StreamOption) error {
   492  	return c.stream.Send(request, options...)
   493  }
   494  
   495  func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Recv(options ...yarpc.StreamOption) (*{{$method.ResponseType.GoType $packagePath}}, error) {
   496  	responseMessage, err := c.stream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, options...)
   497  	if responseMessage == nil {
   498          return nil, err
   499      }
   500  	response, ok := responseMessage.(*{{$method.ResponseType.GoType $packagePath}})
   501  	if !ok {
   502  		return nil, v2.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, responseMessage)
   503  	}
   504  	return response, err
   505  }
   506  
   507  func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) CloseSend(options ...yarpc.StreamOption) error {
   508  	return c.stream.Close(options...)
   509  }
   510  {{end}}
   511  
   512  {{range $method := clientStreamingMethods $service}}
   513  type _{{$service.GetName}}Service{{$method.GetName}}YARPCServer struct {
   514  	serverStream *v2.ServerStream
   515  }
   516  
   517  func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Context() context.Context {
   518  	return s.serverStream.Context()
   519  }
   520  
   521  func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Recv(options ...yarpc.StreamOption) (*{{$method.RequestType.GoType $packagePath}}, error) {
   522  	requestMessage, err := s.serverStream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, options...)
   523  	if requestMessage == nil {
   524          return nil, err
   525      }
   526  	request, ok := requestMessage.(*{{$method.RequestType.GoType $packagePath}})
   527  	if !ok {
   528  		return nil, v2.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, requestMessage)
   529  	}
   530  	return request, err
   531  }
   532  {{end}}
   533  
   534  {{range $method := serverStreamingMethods $service}}
   535  type _{{$service.GetName}}Service{{$method.GetName}}YARPCServer struct {
   536  	serverStream *v2.ServerStream
   537  }
   538  
   539  func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Context() context.Context {
   540  	return s.serverStream.Context()
   541  }
   542  
   543  func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Send(response *{{$method.ResponseType.GoType $packagePath}}, options ...yarpc.StreamOption) error {
   544  	return s.serverStream.Send(response, options...)
   545  }
   546  {{end}}
   547  
   548  {{range $method := clientServerStreamingMethods $service}}
   549  type _{{$service.GetName}}Service{{$method.GetName}}YARPCServer struct {
   550  	serverStream *v2.ServerStream
   551  }
   552  
   553  func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Context() context.Context {
   554  	return s.serverStream.Context()
   555  }
   556  
   557  func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Recv(options ...yarpc.StreamOption) (*{{$method.RequestType.GoType $packagePath}}, error) {
   558  	requestMessage, err := s.serverStream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, options...)
   559  	if requestMessage == nil {
   560          return nil, err
   561      }
   562  	request, ok := requestMessage.(*{{$method.RequestType.GoType $packagePath}})
   563  	if !ok {
   564  		return nil, v2.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, requestMessage)
   565  	}
   566  	return request, err
   567  }
   568  
   569  func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Send(response *{{$method.ResponseType.GoType $packagePath}}, options ...yarpc.StreamOption) error {
   570  	return s.serverStream.Send(response, options...)
   571  }
   572  {{end}}
   573  
   574  {{range $method := $service.Methods}}
   575  func new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest() proto.Message {
   576  	return &{{$method.RequestType.GoType $packagePath}}{}
   577  }
   578  
   579  func new{{$service.GetName}}Service{{$method.GetName}}YARPCResponse() proto.Message {
   580  	return &{{$method.ResponseType.GoType $packagePath}}{}
   581  }
   582  {{end}}
   583  var (
   584  {{range $method := $service.Methods}}
   585  	empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest = &{{$method.RequestType.GoType $packagePath}}{}
   586  	empty{{$service.GetName}}Service{{$method.GetName}}YARPCResponse = &{{$method.ResponseType.GoType $packagePath}}{}{{end}}
   587  )
   588  {{end}}
   589  
   590  var {{ fileDescriptorClosureVarName .File }} = [][]byte{
   591  	// {{ .Name }}
   592  	{{ encodedFileDescriptor .File }},{{range $dependency := .TransitiveDependencies }}
   593  	// {{ $dependency.Name }}
   594  	{{ encodedFileDescriptor $dependency }},{{end}}
   595  }
   596  
   597  {{if .Services}}func init() { {{range $service := .Services}}
   598  	yarpc.RegisterClientBuilder(
   599  		func(clientConfig transport.ClientConfig, structField reflect.StructField) {{$service.GetName}}YARPCClient {
   600  			return New{{$service.GetName}}YARPCClient(clientConfig, v2.ClientBuilderOptions(clientConfig, structField)...)
   601  		},
   602  	){{end}}
   603  }{{end}}
   604  `
   605  
   606  // Runner is the Runner used for protoc-gen-yarpc-go-v2.
   607  var Runner = protoplugin.NewRunner(
   608  	template.Must(template.New("tmpl").Funcs(
   609  		template.FuncMap{
   610  			"unaryMethods":                 unaryMethods,
   611  			"onewayMethods":                onewayMethods,
   612  			"clientStreamingMethods":       clientStreamingMethods,
   613  			"serverStreamingMethods":       serverStreamingMethods,
   614  			"clientServerStreamingMethods": clientServerStreamingMethods,
   615  			"encodedFileDescriptor":        encodedFileDescriptor,
   616  			"fileDescriptorClosureVarName": fileDescriptorClosureVarName,
   617  			"trimPrefixPeriod":             trimPrefixPeriod,
   618  		}).Parse(tmpl)),
   619  	checkTemplateInfo,
   620  	[]string{
   621  		"context",
   622  		"io/ioutil",
   623  		"reflect",
   624  		"go.uber.org/yarpc/encoding/protobuf/v2",
   625  		"google.golang.org/protobuf/proto",
   626  		"go.uber.org/fx",
   627  		"go.uber.org/yarpc",
   628  		"go.uber.org/yarpc/api/transport",
   629  		"go.uber.org/yarpc/api/x/restriction",
   630  		"go.uber.org/yarpc/encoding/protobuf/v2",
   631  		"go.uber.org/yarpc/encoding/protobuf/reflection",
   632  	},
   633  	func(file *protoplugin.File) (string, error) {
   634  		name := file.GetName()
   635  		return fmt.Sprintf("%s.pb.yarpc.go", strings.TrimSuffix(name, filepath.Ext(name))), nil
   636  	},
   637  	func(key string, value string) error {
   638  		return nil
   639  	},
   640  )
   641  
   642  func checkTemplateInfo(templateInfo *protoplugin.TemplateInfo) error {
   643  	return nil
   644  }
   645  
   646  func unaryMethods(service *protoplugin.Service) ([]*protoplugin.Method, error) {
   647  	methods := make([]*protoplugin.Method, 0, len(service.Methods))
   648  	for _, method := range service.Methods {
   649  		if !method.GetClientStreaming() && !method.GetServerStreaming() && method.ResponseType.FQMN() != ".uber.yarpc.Oneway" {
   650  			methods = append(methods, method)
   651  		}
   652  	}
   653  	return methods, nil
   654  }
   655  
   656  func onewayMethods(service *protoplugin.Service) ([]*protoplugin.Method, error) {
   657  	methods := make([]*protoplugin.Method, 0, len(service.Methods))
   658  	for _, method := range service.Methods {
   659  		if !method.GetClientStreaming() && !method.GetServerStreaming() && method.ResponseType.FQMN() == ".uber.yarpc.Oneway" {
   660  			methods = append(methods, method)
   661  		}
   662  	}
   663  	return methods, nil
   664  }
   665  
   666  func clientStreamingMethods(service *protoplugin.Service) ([]*protoplugin.Method, error) {
   667  	methods := make([]*protoplugin.Method, 0, len(service.Methods))
   668  	for _, method := range service.Methods {
   669  		if method.GetClientStreaming() && !method.GetServerStreaming() {
   670  			methods = append(methods, method)
   671  		}
   672  	}
   673  	return methods, nil
   674  }
   675  
   676  func serverStreamingMethods(service *protoplugin.Service) ([]*protoplugin.Method, error) {
   677  	methods := make([]*protoplugin.Method, 0, len(service.Methods))
   678  	for _, method := range service.Methods {
   679  		if !method.GetClientStreaming() && method.GetServerStreaming() {
   680  			methods = append(methods, method)
   681  		}
   682  	}
   683  	return methods, nil
   684  }
   685  
   686  func clientServerStreamingMethods(service *protoplugin.Service) ([]*protoplugin.Method, error) {
   687  	methods := make([]*protoplugin.Method, 0, len(service.Methods))
   688  	for _, method := range service.Methods {
   689  		if method.GetClientStreaming() && method.GetServerStreaming() {
   690  			methods = append(methods, method)
   691  		}
   692  	}
   693  	return methods, nil
   694  }
   695  
   696  // fileDescriptorClosureVarName is used to refer to a variable that contains a closure of all encoded
   697  // file descriptors required to interpret a specific proto file. It is used in the yarpc codebase to
   698  // attach reflection information to services.
   699  func fileDescriptorClosureVarName(f *protoplugin.File) (string, error) {
   700  	name := f.GetName()
   701  	if name == "" {
   702  		return "", fmt.Errorf("could not create fileDescriptorClosureVarName: %s has no name", f)
   703  	}
   704  
   705  	// Use a sha256 of the filename instead of the filename to prevent any characters that are illegal
   706  	// as golang identifiers and to discourage external usage of this constant.
   707  	h := sha256.Sum256([]byte(name))
   708  	return fmt.Sprintf("yarpcFileDescriptorClosure%s", hex.EncodeToString(h[:8])), nil
   709  }
   710  
   711  func encodedFileDescriptor(f *protoplugin.File) (string, error) {
   712  	fdBytes, err := f.SerializedFileDescriptor()
   713  	if err != nil {
   714  		return "", err
   715  	}
   716  
   717  	// Create string that contains a golang byte slice literal containing the
   718  	// serialized file descriptor:
   719  	//
   720  	// []byte{
   721  	//     0x00, 0x01, 0x02, ..., 0xFF,	// Up to 16 bytes per line
   722  	// }
   723  	//
   724  	var buf bytes.Buffer
   725  	buf.WriteString("[]byte{\n")
   726  	for len(fdBytes) > 0 {
   727  		n := 16
   728  		if n > len(fdBytes) {
   729  			n = len(fdBytes)
   730  		}
   731  		for _, c := range fdBytes[:n] {
   732  			fmt.Fprintf(&buf, "0x%02x,", c)
   733  		}
   734  		buf.WriteString("\n")
   735  		fdBytes = fdBytes[n:]
   736  	}
   737  	buf.WriteString("}")
   738  	return buf.String(), nil
   739  }
   740  
   741  func trimPrefixPeriod(s string) string {
   742  	return strings.TrimPrefix(s, ".")
   743  }