grpc.go 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464
  1. // Go support for Protocol Buffers - Google's data interchange format
  2. //
  3. // Copyright 2015 The Go Authors. All rights reserved.
  4. // https://github.com/golang/protobuf
  5. //
  6. // Redistribution and use in source and binary forms, with or without
  7. // modification, are permitted provided that the following conditions are
  8. // met:
  9. //
  10. // * Redistributions of source code must retain the above copyright
  11. // notice, this list of conditions and the following disclaimer.
  12. // * Redistributions in binary form must reproduce the above
  13. // copyright notice, this list of conditions and the following disclaimer
  14. // in the documentation and/or other materials provided with the
  15. // distribution.
  16. // * Neither the name of Google Inc. nor the names of its
  17. // contributors may be used to endorse or promote products derived from
  18. // this software without specific prior written permission.
  19. //
  20. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  21. // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  22. // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  23. // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  24. // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  25. // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  26. // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  27. // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  28. // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  29. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  30. // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  31. // Package grpc outputs gRPC service descriptions in Go code.
  32. // It runs as a plugin for the Go protocol buffer compiler plugin.
  33. // It is linked in to protoc-gen-go.
  34. package grpc
  35. import (
  36. "fmt"
  37. "path"
  38. "strconv"
  39. "strings"
  40. pb "github.com/gogo/protobuf/protoc-gen-gogo/descriptor"
  41. "github.com/gogo/protobuf/protoc-gen-gogo/generator"
  42. )
  43. // generatedCodeVersion indicates a version of the generated code.
  44. // It is incremented whenever an incompatibility between the generated code and
  45. // the grpc package is introduced; the generated code references
  46. // a constant, grpc.SupportPackageIsVersionN (where N is generatedCodeVersion).
  47. const generatedCodeVersion = 4
  48. // Paths for packages used by code generated in this file,
  49. // relative to the import_prefix of the generator.Generator.
  50. const (
  51. contextPkgPath = "golang.org/x/net/context"
  52. grpcPkgPath = "google.golang.org/grpc"
  53. )
  54. func init() {
  55. generator.RegisterPlugin(new(grpc))
  56. }
  57. // grpc is an implementation of the Go protocol buffer compiler's
  58. // plugin architecture. It generates bindings for gRPC support.
  59. type grpc struct {
  60. gen *generator.Generator
  61. }
  62. // Name returns the name of this plugin, "grpc".
  63. func (g *grpc) Name() string {
  64. return "grpc"
  65. }
  66. // The names for packages imported in the generated code.
  67. // They may vary from the final path component of the import path
  68. // if the name is used by other packages.
  69. var (
  70. contextPkg string
  71. grpcPkg string
  72. )
  73. // Init initializes the plugin.
  74. func (g *grpc) Init(gen *generator.Generator) {
  75. g.gen = gen
  76. contextPkg = generator.RegisterUniquePackageName("context", nil)
  77. grpcPkg = generator.RegisterUniquePackageName("grpc", nil)
  78. }
  79. // Given a type name defined in a .proto, return its object.
  80. // Also record that we're using it, to guarantee the associated import.
  81. func (g *grpc) objectNamed(name string) generator.Object {
  82. g.gen.RecordTypeUse(name)
  83. return g.gen.ObjectNamed(name)
  84. }
  85. // Given a type name defined in a .proto, return its name as we will print it.
  86. func (g *grpc) typeName(str string) string {
  87. return g.gen.TypeName(g.objectNamed(str))
  88. }
  89. // P forwards to g.gen.P.
  90. func (g *grpc) P(args ...interface{}) { g.gen.P(args...) }
  91. // Generate generates code for the services in the given file.
  92. func (g *grpc) Generate(file *generator.FileDescriptor) {
  93. if len(file.FileDescriptorProto.Service) == 0 {
  94. return
  95. }
  96. g.P("// Reference imports to suppress errors if they are not otherwise used.")
  97. g.P("var _ ", contextPkg, ".Context")
  98. g.P("var _ ", grpcPkg, ".ClientConn")
  99. g.P()
  100. // Assert version compatibility.
  101. g.P("// This is a compile-time assertion to ensure that this generated file")
  102. g.P("// is compatible with the grpc package it is being compiled against.")
  103. g.P("const _ = ", grpcPkg, ".SupportPackageIsVersion", generatedCodeVersion)
  104. g.P()
  105. for i, service := range file.FileDescriptorProto.Service {
  106. g.generateService(file, service, i)
  107. }
  108. }
  109. // GenerateImports generates the import declaration for this file.
  110. func (g *grpc) GenerateImports(file *generator.FileDescriptor) {
  111. if len(file.FileDescriptorProto.Service) == 0 {
  112. return
  113. }
  114. g.P("import (")
  115. g.P(contextPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, contextPkgPath)))
  116. g.P(grpcPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, grpcPkgPath)))
  117. g.P(")")
  118. g.P()
  119. }
  120. // reservedClientName records whether a client name is reserved on the client side.
  121. var reservedClientName = map[string]bool{
  122. // TODO: do we need any in gRPC?
  123. }
  124. func unexport(s string) string { return strings.ToLower(s[:1]) + s[1:] }
  125. // generateService generates all the code for the named service.
  126. func (g *grpc) generateService(file *generator.FileDescriptor, service *pb.ServiceDescriptorProto, index int) {
  127. path := fmt.Sprintf("6,%d", index) // 6 means service.
  128. origServName := service.GetName()
  129. fullServName := origServName
  130. if pkg := file.GetPackage(); pkg != "" {
  131. fullServName = pkg + "." + fullServName
  132. }
  133. servName := generator.CamelCase(origServName)
  134. g.P()
  135. g.P("// Client API for ", servName, " service")
  136. g.P()
  137. // Client interface.
  138. g.P("type ", servName, "Client interface {")
  139. for i, method := range service.Method {
  140. g.gen.PrintComments(fmt.Sprintf("%s,2,%d", path, i)) // 2 means method in a service.
  141. g.P(g.generateClientSignature(servName, method))
  142. }
  143. g.P("}")
  144. g.P()
  145. // Client structure.
  146. g.P("type ", unexport(servName), "Client struct {")
  147. g.P("cc *", grpcPkg, ".ClientConn")
  148. g.P("}")
  149. g.P()
  150. // NewClient factory.
  151. g.P("func New", servName, "Client (cc *", grpcPkg, ".ClientConn) ", servName, "Client {")
  152. g.P("return &", unexport(servName), "Client{cc}")
  153. g.P("}")
  154. g.P()
  155. var methodIndex, streamIndex int
  156. serviceDescVar := "_" + servName + "_serviceDesc"
  157. // Client method implementations.
  158. for _, method := range service.Method {
  159. var descExpr string
  160. if !method.GetServerStreaming() && !method.GetClientStreaming() {
  161. // Unary RPC method
  162. descExpr = fmt.Sprintf("&%s.Methods[%d]", serviceDescVar, methodIndex)
  163. methodIndex++
  164. } else {
  165. // Streaming RPC method
  166. descExpr = fmt.Sprintf("&%s.Streams[%d]", serviceDescVar, streamIndex)
  167. streamIndex++
  168. }
  169. g.generateClientMethod(servName, fullServName, serviceDescVar, method, descExpr)
  170. }
  171. g.P("// Server API for ", servName, " service")
  172. g.P()
  173. // Server interface.
  174. serverType := servName + "Server"
  175. g.P("type ", serverType, " interface {")
  176. for i, method := range service.Method {
  177. g.gen.PrintComments(fmt.Sprintf("%s,2,%d", path, i)) // 2 means method in a service.
  178. g.P(g.generateServerSignature(servName, method))
  179. }
  180. g.P("}")
  181. g.P()
  182. // Server registration.
  183. g.P("func Register", servName, "Server(s *", grpcPkg, ".Server, srv ", serverType, ") {")
  184. g.P("s.RegisterService(&", serviceDescVar, `, srv)`)
  185. g.P("}")
  186. g.P()
  187. // Server handler implementations.
  188. var handlerNames []string
  189. for _, method := range service.Method {
  190. hname := g.generateServerMethod(servName, fullServName, method)
  191. handlerNames = append(handlerNames, hname)
  192. }
  193. // Service descriptor.
  194. g.P("var ", serviceDescVar, " = ", grpcPkg, ".ServiceDesc {")
  195. g.P("ServiceName: ", strconv.Quote(fullServName), ",")
  196. g.P("HandlerType: (*", serverType, ")(nil),")
  197. g.P("Methods: []", grpcPkg, ".MethodDesc{")
  198. for i, method := range service.Method {
  199. if method.GetServerStreaming() || method.GetClientStreaming() {
  200. continue
  201. }
  202. g.P("{")
  203. g.P("MethodName: ", strconv.Quote(method.GetName()), ",")
  204. g.P("Handler: ", handlerNames[i], ",")
  205. g.P("},")
  206. }
  207. g.P("},")
  208. g.P("Streams: []", grpcPkg, ".StreamDesc{")
  209. for i, method := range service.Method {
  210. if !method.GetServerStreaming() && !method.GetClientStreaming() {
  211. continue
  212. }
  213. g.P("{")
  214. g.P("StreamName: ", strconv.Quote(method.GetName()), ",")
  215. g.P("Handler: ", handlerNames[i], ",")
  216. if method.GetServerStreaming() {
  217. g.P("ServerStreams: true,")
  218. }
  219. if method.GetClientStreaming() {
  220. g.P("ClientStreams: true,")
  221. }
  222. g.P("},")
  223. }
  224. g.P("},")
  225. g.P("Metadata: \"", file.GetName(), "\",")
  226. g.P("}")
  227. g.P()
  228. }
  229. // generateClientSignature returns the client-side signature for a method.
  230. func (g *grpc) generateClientSignature(servName string, method *pb.MethodDescriptorProto) string {
  231. origMethName := method.GetName()
  232. methName := generator.CamelCase(origMethName)
  233. if reservedClientName[methName] {
  234. methName += "_"
  235. }
  236. reqArg := ", in *" + g.typeName(method.GetInputType())
  237. if method.GetClientStreaming() {
  238. reqArg = ""
  239. }
  240. respName := "*" + g.typeName(method.GetOutputType())
  241. if method.GetServerStreaming() || method.GetClientStreaming() {
  242. respName = servName + "_" + generator.CamelCase(origMethName) + "Client"
  243. }
  244. return fmt.Sprintf("%s(ctx %s.Context%s, opts ...%s.CallOption) (%s, error)", methName, contextPkg, reqArg, grpcPkg, respName)
  245. }
  246. func (g *grpc) generateClientMethod(servName, fullServName, serviceDescVar string, method *pb.MethodDescriptorProto, descExpr string) {
  247. sname := fmt.Sprintf("/%s/%s", fullServName, method.GetName())
  248. methName := generator.CamelCase(method.GetName())
  249. inType := g.typeName(method.GetInputType())
  250. outType := g.typeName(method.GetOutputType())
  251. g.P("func (c *", unexport(servName), "Client) ", g.generateClientSignature(servName, method), "{")
  252. if !method.GetServerStreaming() && !method.GetClientStreaming() {
  253. g.P("out := new(", outType, ")")
  254. // TODO: Pass descExpr to Invoke.
  255. g.P("err := ", grpcPkg, `.Invoke(ctx, "`, sname, `", in, out, c.cc, opts...)`)
  256. g.P("if err != nil { return nil, err }")
  257. g.P("return out, nil")
  258. g.P("}")
  259. g.P()
  260. return
  261. }
  262. streamType := unexport(servName) + methName + "Client"
  263. g.P("stream, err := ", grpcPkg, ".NewClientStream(ctx, ", descExpr, `, c.cc, "`, sname, `", opts...)`)
  264. g.P("if err != nil { return nil, err }")
  265. g.P("x := &", streamType, "{stream}")
  266. if !method.GetClientStreaming() {
  267. g.P("if err := x.ClientStream.SendMsg(in); err != nil { return nil, err }")
  268. g.P("if err := x.ClientStream.CloseSend(); err != nil { return nil, err }")
  269. }
  270. g.P("return x, nil")
  271. g.P("}")
  272. g.P()
  273. genSend := method.GetClientStreaming()
  274. genRecv := method.GetServerStreaming()
  275. genCloseAndRecv := !method.GetServerStreaming()
  276. // Stream auxiliary types and methods.
  277. g.P("type ", servName, "_", methName, "Client interface {")
  278. if genSend {
  279. g.P("Send(*", inType, ") error")
  280. }
  281. if genRecv {
  282. g.P("Recv() (*", outType, ", error)")
  283. }
  284. if genCloseAndRecv {
  285. g.P("CloseAndRecv() (*", outType, ", error)")
  286. }
  287. g.P(grpcPkg, ".ClientStream")
  288. g.P("}")
  289. g.P()
  290. g.P("type ", streamType, " struct {")
  291. g.P(grpcPkg, ".ClientStream")
  292. g.P("}")
  293. g.P()
  294. if genSend {
  295. g.P("func (x *", streamType, ") Send(m *", inType, ") error {")
  296. g.P("return x.ClientStream.SendMsg(m)")
  297. g.P("}")
  298. g.P()
  299. }
  300. if genRecv {
  301. g.P("func (x *", streamType, ") Recv() (*", outType, ", error) {")
  302. g.P("m := new(", outType, ")")
  303. g.P("if err := x.ClientStream.RecvMsg(m); err != nil { return nil, err }")
  304. g.P("return m, nil")
  305. g.P("}")
  306. g.P()
  307. }
  308. if genCloseAndRecv {
  309. g.P("func (x *", streamType, ") CloseAndRecv() (*", outType, ", error) {")
  310. g.P("if err := x.ClientStream.CloseSend(); err != nil { return nil, err }")
  311. g.P("m := new(", outType, ")")
  312. g.P("if err := x.ClientStream.RecvMsg(m); err != nil { return nil, err }")
  313. g.P("return m, nil")
  314. g.P("}")
  315. g.P()
  316. }
  317. }
  318. // generateServerSignature returns the server-side signature for a method.
  319. func (g *grpc) generateServerSignature(servName string, method *pb.MethodDescriptorProto) string {
  320. origMethName := method.GetName()
  321. methName := generator.CamelCase(origMethName)
  322. if reservedClientName[methName] {
  323. methName += "_"
  324. }
  325. var reqArgs []string
  326. ret := "error"
  327. if !method.GetServerStreaming() && !method.GetClientStreaming() {
  328. reqArgs = append(reqArgs, contextPkg+".Context")
  329. ret = "(*" + g.typeName(method.GetOutputType()) + ", error)"
  330. }
  331. if !method.GetClientStreaming() {
  332. reqArgs = append(reqArgs, "*"+g.typeName(method.GetInputType()))
  333. }
  334. if method.GetServerStreaming() || method.GetClientStreaming() {
  335. reqArgs = append(reqArgs, servName+"_"+generator.CamelCase(origMethName)+"Server")
  336. }
  337. return methName + "(" + strings.Join(reqArgs, ", ") + ") " + ret
  338. }
  339. func (g *grpc) generateServerMethod(servName, fullServName string, method *pb.MethodDescriptorProto) string {
  340. methName := generator.CamelCase(method.GetName())
  341. hname := fmt.Sprintf("_%s_%s_Handler", servName, methName)
  342. inType := g.typeName(method.GetInputType())
  343. outType := g.typeName(method.GetOutputType())
  344. if !method.GetServerStreaming() && !method.GetClientStreaming() {
  345. g.P("func ", hname, "(srv interface{}, ctx ", contextPkg, ".Context, dec func(interface{}) error, interceptor ", grpcPkg, ".UnaryServerInterceptor) (interface{}, error) {")
  346. g.P("in := new(", inType, ")")
  347. g.P("if err := dec(in); err != nil { return nil, err }")
  348. g.P("if interceptor == nil { return srv.(", servName, "Server).", methName, "(ctx, in) }")
  349. g.P("info := &", grpcPkg, ".UnaryServerInfo{")
  350. g.P("Server: srv,")
  351. g.P("FullMethod: ", strconv.Quote(fmt.Sprintf("/%s/%s", fullServName, methName)), ",")
  352. g.P("}")
  353. g.P("handler := func(ctx ", contextPkg, ".Context, req interface{}) (interface{}, error) {")
  354. g.P("return srv.(", servName, "Server).", methName, "(ctx, req.(*", inType, "))")
  355. g.P("}")
  356. g.P("return interceptor(ctx, in, info, handler)")
  357. g.P("}")
  358. g.P()
  359. return hname
  360. }
  361. streamType := unexport(servName) + methName + "Server"
  362. g.P("func ", hname, "(srv interface{}, stream ", grpcPkg, ".ServerStream) error {")
  363. if !method.GetClientStreaming() {
  364. g.P("m := new(", inType, ")")
  365. g.P("if err := stream.RecvMsg(m); err != nil { return err }")
  366. g.P("return srv.(", servName, "Server).", methName, "(m, &", streamType, "{stream})")
  367. } else {
  368. g.P("return srv.(", servName, "Server).", methName, "(&", streamType, "{stream})")
  369. }
  370. g.P("}")
  371. g.P()
  372. genSend := method.GetServerStreaming()
  373. genSendAndClose := !method.GetServerStreaming()
  374. genRecv := method.GetClientStreaming()
  375. // Stream auxiliary types and methods.
  376. g.P("type ", servName, "_", methName, "Server interface {")
  377. if genSend {
  378. g.P("Send(*", outType, ") error")
  379. }
  380. if genSendAndClose {
  381. g.P("SendAndClose(*", outType, ") error")
  382. }
  383. if genRecv {
  384. g.P("Recv() (*", inType, ", error)")
  385. }
  386. g.P(grpcPkg, ".ServerStream")
  387. g.P("}")
  388. g.P()
  389. g.P("type ", streamType, " struct {")
  390. g.P(grpcPkg, ".ServerStream")
  391. g.P("}")
  392. g.P()
  393. if genSend {
  394. g.P("func (x *", streamType, ") Send(m *", outType, ") error {")
  395. g.P("return x.ServerStream.SendMsg(m)")
  396. g.P("}")
  397. g.P()
  398. }
  399. if genSendAndClose {
  400. g.P("func (x *", streamType, ") SendAndClose(m *", outType, ") error {")
  401. g.P("return x.ServerStream.SendMsg(m)")
  402. g.P("}")
  403. g.P()
  404. }
  405. if genRecv {
  406. g.P("func (x *", streamType, ") Recv() (*", inType, ", error) {")
  407. g.P("m := new(", inType, ")")
  408. g.P("if err := x.ServerStream.RecvMsg(m); err != nil { return nil, err }")
  409. g.P("return m, nil")
  410. g.P("}")
  411. g.P()
  412. }
  413. return hname
  414. }