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 }