go.uber.org/yarpc@v1.72.1/encoding/protobuf/protoc-gen-yarpc-go/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. 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 "go.uber.org/yarpc/internal/protoplugin" 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 jsonpb.AnyResolver, options ...protobuf.ClientOption) {{$service.GetName}}YARPCClient { 92 return &_{{$service.GetName}}YARPCCaller{protobuf.NewStreamClient( 93 protobuf.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 ...protobuf.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 jsonpb.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 protobuf.BuildProcedures( 150 protobuf.BuildProceduresParams{ 151 ServiceName: "{{trimPrefixPeriod $service.FQSN}}", 152 UnaryHandlerParams: []protobuf.BuildProceduresUnaryHandlerParams{ 153 {{range $method := unaryMethods $service}}{ 154 MethodName: "{{$method.GetName}}", 155 Handler: protobuf.NewUnaryHandler( 156 protobuf.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: []protobuf.BuildProceduresOnewayHandlerParams{ 166 {{range $method := onewayMethods $service}}{ 167 MethodName: "{{$method.GetName}}", 168 Handler: protobuf.NewOnewayHandler( 169 protobuf.OnewayHandlerParams{ 170 Handle: handler.{{$method.GetName}}, 171 NewRequest: new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, 172 }, 173 ), 174 }, 175 {{end}} 176 }, 177 StreamHandlerParams: []protobuf.BuildProceduresStreamHandlerParams{ 178 {{range $method := clientServerStreamingMethods $service}}{ 179 MethodName: "{{$method.GetName}}", 180 Handler: protobuf.NewStreamHandler( 181 protobuf.StreamHandlerParams{ 182 Handle: handler.{{$method.GetName}}, 183 }, 184 ), 185 }, 186 {{end}} 187 {{range $method := serverStreamingMethods $service}}{ 188 MethodName: "{{$method.GetName}}", 189 Handler: protobuf.NewStreamHandler( 190 protobuf.StreamHandlerParams{ 191 Handle: handler.{{$method.GetName}}, 192 }, 193 ), 194 }, 195 {{end}} 196 {{range $method := clientStreamingMethods $service}}{ 197 MethodName: "{{$method.GetName}}", 198 Handler: protobuf.NewStreamHandler( 199 protobuf.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 jsonpb.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 ...protobuf.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(protobuf.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 jsonpb.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: {{$service.GetName}}ReflectionMeta, 304 } 305 } 306 } 307 308 // {{$service.GetName}}ReflectionMeta is the reflection server metadata 309 // required for using the gRPC reflection protocol with YARPC. 310 // 311 // See https://github.com/grpc/grpc/blob/master/doc/server-reflection.md. 312 var {{$service.GetName}}ReflectionMeta = reflection.ServerMeta{ 313 ServiceName: "{{trimPrefixPeriod $service.FQSN}}", 314 FileDescriptors: {{ fileDescriptorClosureVarName .File }}, 315 } 316 317 type _{{$service.GetName}}YARPCCaller struct { 318 streamClient protobuf.StreamClient 319 } 320 321 {{range $method := unaryMethods $service}} 322 func (c *_{{$service.GetName}}YARPCCaller) {{$method.GetName}}(ctx context.Context, request *{{$method.RequestType.GoType $packagePath}}, options ...yarpc.CallOption) (*{{$method.ResponseType.GoType $packagePath}}, error) { 323 responseMessage, err := c.streamClient.Call(ctx, "{{$method.GetName}}", request, new{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, options...) 324 if responseMessage == nil { 325 return nil, err 326 } 327 response, ok := responseMessage.(*{{$method.ResponseType.GoType $packagePath}}) 328 if !ok { 329 return nil, protobuf.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, responseMessage) 330 } 331 return response, err 332 } 333 {{end}} 334 {{range $method := onewayMethods $service}} 335 func (c *_{{$service.GetName}}YARPCCaller) {{$method.GetName}}(ctx context.Context, request *{{$method.RequestType.GoType $packagePath}}, options ...yarpc.CallOption) (yarpc.Ack, error) { 336 return c.streamClient.CallOneway(ctx, "{{$method.GetName}}", request, options...) 337 } 338 {{end}} 339 {{range $method := clientStreamingMethods $service}} 340 func (c *_{{$service.GetName}}YARPCCaller) {{$method.GetName}}(ctx context.Context, options ...yarpc.CallOption) ({{$service.GetName}}Service{{$method.GetName}}YARPCClient, error) { 341 stream, err := c.streamClient.CallStream(ctx, "{{$method.GetName}}", options...) 342 if err != nil { 343 return nil, err 344 } 345 return &_{{$service.GetName}}Service{{$method.GetName}}YARPCClient{stream: stream}, nil 346 } 347 {{end}} 348 {{range $method := serverStreamingMethods $service}} 349 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) { 350 stream, err := c.streamClient.CallStream(ctx, "{{$method.GetName}}", options...) 351 if err != nil { 352 return nil, err 353 } 354 if err := stream.Send(request); err != nil { 355 return nil, err 356 } 357 return &_{{$service.GetName}}Service{{$method.GetName}}YARPCClient{stream: stream}, nil 358 } 359 {{end}} 360 {{range $method := clientServerStreamingMethods $service}} 361 func (c *_{{$service.GetName}}YARPCCaller) {{$method.GetName}}(ctx context.Context, options ...yarpc.CallOption) ({{$service.GetName}}Service{{$method.GetName}}YARPCClient, error) { 362 stream, err := c.streamClient.CallStream(ctx, "{{$method.GetName}}", options...) 363 if err != nil { 364 return nil, err 365 } 366 return &_{{$service.GetName}}Service{{$method.GetName}}YARPCClient{stream: stream}, nil 367 } 368 {{end}} 369 370 type _{{$service.GetName}}YARPCHandler struct { 371 server {{$service.GetName}}YARPCServer 372 } 373 374 {{range $method := unaryMethods $service}} 375 func (h *_{{$service.GetName}}YARPCHandler) {{$method.GetName}}(ctx context.Context, requestMessage proto.Message) (proto.Message, error) { 376 var request *{{$method.RequestType.GoType $packagePath}} 377 var ok bool 378 if requestMessage != nil { 379 request, ok = requestMessage.(*{{$method.RequestType.GoType $packagePath}}) 380 if !ok { 381 return nil, protobuf.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, requestMessage) 382 } 383 } 384 response, err := h.server.{{$method.GetName}}(ctx, request) 385 if response == nil { 386 return nil, err 387 } 388 return response, err 389 } 390 {{end}} 391 {{range $method := onewayMethods $service}} 392 func (h *_{{$service.GetName}}YARPCHandler) {{$method.GetName}}(ctx context.Context, requestMessage proto.Message) error { 393 var request *{{$method.RequestType.GoType $packagePath}} 394 var ok bool 395 if requestMessage != nil { 396 request, ok = requestMessage.(*{{$method.RequestType.GoType $packagePath}}) 397 if !ok { 398 return protobuf.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, requestMessage) 399 } 400 } 401 return h.server.{{$method.GetName}}(ctx, request) 402 } 403 {{end}} 404 {{range $method := clientStreamingMethods $service}} 405 func (h *_{{$service.GetName}}YARPCHandler) {{$method.GetName}}(serverStream *protobuf.ServerStream) error { 406 response, err := h.server.{{$method.GetName}}(&_{{$service.GetName}}Service{{$method.GetName}}YARPCServer{serverStream: serverStream}) 407 if err != nil { 408 return err 409 } 410 return serverStream.Send(response) 411 } 412 {{end}} 413 {{range $method := serverStreamingMethods $service}} 414 func (h *_{{$service.GetName}}YARPCHandler) {{$method.GetName}}(serverStream *protobuf.ServerStream) error { 415 requestMessage, err := serverStream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest) 416 if requestMessage == nil { 417 return err 418 } 419 420 request, ok := requestMessage.(*{{$method.RequestType.GoType $packagePath}}) 421 if !ok { 422 return protobuf.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, requestMessage) 423 } 424 return h.server.{{$method.GetName}}(request, &_{{$service.GetName}}Service{{$method.GetName}}YARPCServer{serverStream: serverStream}) 425 } 426 {{end}} 427 {{range $method := clientServerStreamingMethods $service}} 428 func (h *_{{$service.GetName}}YARPCHandler) {{$method.GetName}}(serverStream *protobuf.ServerStream) error { 429 return h.server.{{$method.GetName}}(&_{{$service.GetName}}Service{{$method.GetName}}YARPCServer{serverStream: serverStream}) 430 } 431 {{end}} 432 433 {{range $method := clientStreamingMethods $service}} 434 type _{{$service.GetName}}Service{{$method.GetName}}YARPCClient struct { 435 stream *protobuf.ClientStream 436 } 437 438 func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Context() context.Context { 439 return c.stream.Context() 440 } 441 442 func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Send(request *{{$method.RequestType.GoType $packagePath}}, options ...yarpc.StreamOption) error { 443 return c.stream.Send(request, options...) 444 } 445 446 func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) CloseAndRecv(options ...yarpc.StreamOption) (*{{$method.ResponseType.GoType $packagePath}}, error) { 447 if err := c.stream.Close(options...); err != nil { 448 return nil, err 449 } 450 responseMessage, err := c.stream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, options...) 451 if responseMessage == nil { 452 return nil, err 453 } 454 response, ok := responseMessage.(*{{$method.ResponseType.GoType $packagePath}}) 455 if !ok { 456 return nil, protobuf.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, responseMessage) 457 } 458 return response, err 459 } 460 {{end}} 461 462 {{range $method := serverStreamingMethods $service}} 463 type _{{$service.GetName}}Service{{$method.GetName}}YARPCClient struct { 464 stream *protobuf.ClientStream 465 } 466 467 func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Context() context.Context { 468 return c.stream.Context() 469 } 470 471 func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Recv(options ...yarpc.StreamOption) (*{{$method.ResponseType.GoType $packagePath}}, error) { 472 responseMessage, err := c.stream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, options...) 473 if responseMessage == nil { 474 return nil, err 475 } 476 response, ok := responseMessage.(*{{$method.ResponseType.GoType $packagePath}}) 477 if !ok { 478 return nil, protobuf.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, responseMessage) 479 } 480 return response, err 481 } 482 483 func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) CloseSend(options ...yarpc.StreamOption) error { 484 return c.stream.Close(options...) 485 } 486 {{end}} 487 488 {{range $method := clientServerStreamingMethods $service}} 489 type _{{$service.GetName}}Service{{$method.GetName}}YARPCClient struct { 490 stream *protobuf.ClientStream 491 } 492 493 func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Context() context.Context { 494 return c.stream.Context() 495 } 496 497 func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Send(request *{{$method.RequestType.GoType $packagePath}}, options ...yarpc.StreamOption) error { 498 return c.stream.Send(request, options...) 499 } 500 501 func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) Recv(options ...yarpc.StreamOption) (*{{$method.ResponseType.GoType $packagePath}}, error) { 502 responseMessage, err := c.stream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, options...) 503 if responseMessage == nil { 504 return nil, err 505 } 506 response, ok := responseMessage.(*{{$method.ResponseType.GoType $packagePath}}) 507 if !ok { 508 return nil, protobuf.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCResponse, responseMessage) 509 } 510 return response, err 511 } 512 513 func (c *_{{$service.GetName}}Service{{$method.GetName}}YARPCClient) CloseSend(options ...yarpc.StreamOption) error { 514 return c.stream.Close(options...) 515 } 516 {{end}} 517 518 {{range $method := clientStreamingMethods $service}} 519 type _{{$service.GetName}}Service{{$method.GetName}}YARPCServer struct { 520 serverStream *protobuf.ServerStream 521 } 522 523 func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Context() context.Context { 524 return s.serverStream.Context() 525 } 526 527 func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Recv(options ...yarpc.StreamOption) (*{{$method.RequestType.GoType $packagePath}}, error) { 528 requestMessage, err := s.serverStream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, options...) 529 if requestMessage == nil { 530 return nil, err 531 } 532 request, ok := requestMessage.(*{{$method.RequestType.GoType $packagePath}}) 533 if !ok { 534 return nil, protobuf.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, requestMessage) 535 } 536 return request, err 537 } 538 {{end}} 539 540 {{range $method := serverStreamingMethods $service}} 541 type _{{$service.GetName}}Service{{$method.GetName}}YARPCServer struct { 542 serverStream *protobuf.ServerStream 543 } 544 545 func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Context() context.Context { 546 return s.serverStream.Context() 547 } 548 549 func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Send(response *{{$method.ResponseType.GoType $packagePath}}, options ...yarpc.StreamOption) error { 550 return s.serverStream.Send(response, options...) 551 } 552 {{end}} 553 554 {{range $method := clientServerStreamingMethods $service}} 555 type _{{$service.GetName}}Service{{$method.GetName}}YARPCServer struct { 556 serverStream *protobuf.ServerStream 557 } 558 559 func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Context() context.Context { 560 return s.serverStream.Context() 561 } 562 563 func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Recv(options ...yarpc.StreamOption) (*{{$method.RequestType.GoType $packagePath}}, error) { 564 requestMessage, err := s.serverStream.Receive(new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, options...) 565 if requestMessage == nil { 566 return nil, err 567 } 568 request, ok := requestMessage.(*{{$method.RequestType.GoType $packagePath}}) 569 if !ok { 570 return nil, protobuf.CastError(empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest, requestMessage) 571 } 572 return request, err 573 } 574 575 func (s *_{{$service.GetName}}Service{{$method.GetName}}YARPCServer) Send(response *{{$method.ResponseType.GoType $packagePath}}, options ...yarpc.StreamOption) error { 576 return s.serverStream.Send(response, options...) 577 } 578 {{end}} 579 580 {{range $method := $service.Methods}} 581 func new{{$service.GetName}}Service{{$method.GetName}}YARPCRequest() proto.Message { 582 return &{{$method.RequestType.GoType $packagePath}}{} 583 } 584 585 func new{{$service.GetName}}Service{{$method.GetName}}YARPCResponse() proto.Message { 586 return &{{$method.ResponseType.GoType $packagePath}}{} 587 } 588 {{end}} 589 var ( 590 {{range $method := $service.Methods}} 591 empty{{$service.GetName}}Service{{$method.GetName}}YARPCRequest = &{{$method.RequestType.GoType $packagePath}}{} 592 empty{{$service.GetName}}Service{{$method.GetName}}YARPCResponse = &{{$method.ResponseType.GoType $packagePath}}{}{{end}} 593 ) 594 {{end}} 595 596 var {{ fileDescriptorClosureVarName .File }} = [][]byte{ 597 // {{ .Name }} 598 {{ encodedFileDescriptor .File }},{{range $dependency := .TransitiveDependencies }} 599 // {{ $dependency.Name }} 600 {{ encodedFileDescriptor $dependency }},{{end}} 601 } 602 603 {{if .Services}}func init() { {{range $service := .Services}} 604 yarpc.RegisterClientBuilder( 605 func(clientConfig transport.ClientConfig, structField reflect.StructField) {{$service.GetName}}YARPCClient { 606 return New{{$service.GetName}}YARPCClient(clientConfig, protobuf.ClientBuilderOptions(clientConfig, structField)...) 607 }, 608 ){{end}} 609 }{{end}} 610 ` 611 612 // Runner is the Runner used for protoc-gen-yarpc-go. 613 var Runner = protoplugin.NewRunner( 614 template.Must(template.New("tmpl").Funcs( 615 template.FuncMap{ 616 "unaryMethods": unaryMethods, 617 "onewayMethods": onewayMethods, 618 "clientStreamingMethods": clientStreamingMethods, 619 "serverStreamingMethods": serverStreamingMethods, 620 "clientServerStreamingMethods": clientServerStreamingMethods, 621 "encodedFileDescriptor": encodedFileDescriptor, 622 "fileDescriptorClosureVarName": fileDescriptorClosureVarName, 623 "trimPrefixPeriod": trimPrefixPeriod, 624 }).Parse(tmpl)), 625 checkTemplateInfo, 626 []string{ 627 "context", 628 "io/ioutil", 629 "reflect", 630 "github.com/gogo/protobuf/jsonpb", 631 "github.com/gogo/protobuf/proto", 632 "go.uber.org/fx", 633 "go.uber.org/yarpc", 634 "go.uber.org/yarpc/api/transport", 635 "go.uber.org/yarpc/api/x/restriction", 636 "go.uber.org/yarpc/encoding/protobuf", 637 "go.uber.org/yarpc/encoding/protobuf/reflection", 638 }, 639 func(file *protoplugin.File) (string, error) { 640 name := file.GetName() 641 return fmt.Sprintf("%s.pb.yarpc.go", strings.TrimSuffix(name, filepath.Ext(name))), nil 642 }, 643 func(key string, value string) error { 644 return nil 645 }, 646 ) 647 648 func checkTemplateInfo(templateInfo *protoplugin.TemplateInfo) error { 649 return nil 650 } 651 652 func unaryMethods(service *protoplugin.Service) ([]*protoplugin.Method, error) { 653 methods := make([]*protoplugin.Method, 0, len(service.Methods)) 654 for _, method := range service.Methods { 655 if !method.GetClientStreaming() && !method.GetServerStreaming() && method.ResponseType.FQMN() != ".uber.yarpc.Oneway" { 656 methods = append(methods, method) 657 } 658 } 659 return methods, nil 660 } 661 662 func onewayMethods(service *protoplugin.Service) ([]*protoplugin.Method, error) { 663 methods := make([]*protoplugin.Method, 0, len(service.Methods)) 664 for _, method := range service.Methods { 665 if !method.GetClientStreaming() && !method.GetServerStreaming() && method.ResponseType.FQMN() == ".uber.yarpc.Oneway" { 666 methods = append(methods, method) 667 } 668 } 669 return methods, nil 670 } 671 672 func clientStreamingMethods(service *protoplugin.Service) ([]*protoplugin.Method, error) { 673 methods := make([]*protoplugin.Method, 0, len(service.Methods)) 674 for _, method := range service.Methods { 675 if method.GetClientStreaming() && !method.GetServerStreaming() { 676 methods = append(methods, method) 677 } 678 } 679 return methods, nil 680 } 681 682 func serverStreamingMethods(service *protoplugin.Service) ([]*protoplugin.Method, error) { 683 methods := make([]*protoplugin.Method, 0, len(service.Methods)) 684 for _, method := range service.Methods { 685 if !method.GetClientStreaming() && method.GetServerStreaming() { 686 methods = append(methods, method) 687 } 688 } 689 return methods, nil 690 } 691 692 func clientServerStreamingMethods(service *protoplugin.Service) ([]*protoplugin.Method, error) { 693 methods := make([]*protoplugin.Method, 0, len(service.Methods)) 694 for _, method := range service.Methods { 695 if method.GetClientStreaming() && method.GetServerStreaming() { 696 methods = append(methods, method) 697 } 698 } 699 return methods, nil 700 } 701 702 // fileDescriptorClosureVarName is used to refer to a variable that contains a closure of all encoded 703 // file descriptors required to interpret a specific proto file. It is used in the yarpc codebase to 704 // attach reflection information to services. 705 func fileDescriptorClosureVarName(f *protoplugin.File) (string, error) { 706 name := f.GetName() 707 if name == "" { 708 return "", fmt.Errorf("could not create fileDescriptorClosureVarName: %s has no name", f) 709 } 710 711 // Use a sha256 of the filename instead of the filename to prevent any characters that are illegal 712 // as golang identifiers and to discourage external usage of this constant. 713 h := sha256.Sum256([]byte(name)) 714 return fmt.Sprintf("yarpcFileDescriptorClosure%s", hex.EncodeToString(h[:8])), nil 715 } 716 717 func encodedFileDescriptor(f *protoplugin.File) (string, error) { 718 fdBytes, err := f.SerializedFileDescriptor() 719 if err != nil { 720 return "", err 721 } 722 723 // Create string that contains a golang byte slice literal containing the 724 // serialized file descriptor: 725 // 726 // []byte{ 727 // 0x00, 0x01, 0x02, ..., 0xFF, // Up to 16 bytes per line 728 // } 729 // 730 var buf bytes.Buffer 731 buf.WriteString("[]byte{\n") 732 for len(fdBytes) > 0 { 733 n := 16 734 if n > len(fdBytes) { 735 n = len(fdBytes) 736 } 737 for _, c := range fdBytes[:n] { 738 fmt.Fprintf(&buf, "0x%02x,", c) 739 } 740 buf.WriteString("\n") 741 fdBytes = fdBytes[n:] 742 } 743 buf.WriteString("}") 744 return buf.String(), nil 745 } 746 747 func trimPrefixPeriod(s string) string { 748 return strings.TrimPrefix(s, ".") 749 }