Skip to content

Commit bb9882e

Browse files
aarongableaarongable
authored
Add an optional implementation of streams using generics (#7057)
1 parent a87e923 commit bb9882e

File tree

21 files changed

+462
-906
lines changed

21 files changed

+462
-906
lines changed

‎balancer/grpclb/grpc_lb_v1/load_balancer_grpc.pb.go

+2-2
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎cmd/protoc-gen-go-grpc/grpc.go

+84-19
Original file line numberDiff line numberDiff line change
@@ -174,8 +174,13 @@ func generateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.
174174

175175
g.P("// This is a compile-time assertion to ensure that this generated file")
176176
g.P("// is compatible with the grpc package it is being compiled against.")
177-
g.P("// Requires gRPC-Go v1.62.0 or later.")
178-
g.P("const _ = ", grpcPackage.Ident("SupportPackageIsVersion8")) // When changing, update version number above.
177+
if *useGenericStreams {
178+
g.P("// Requires gRPC-Go v1.64.0 or later.")
179+
g.P("const _ = ", grpcPackage.Ident("SupportPackageIsVersion9"))
180+
} else {
181+
g.P("// Requires gRPC-Go v1.62.0 or later.")
182+
g.P("const _ = ", grpcPackage.Ident("SupportPackageIsVersion8")) // When changing, update version number above.
183+
}
179184
g.P()
180185
for _, service := range file.Services {
181186
genService(gen, file, g, service)
@@ -299,12 +304,27 @@ func clientSignature(g *protogen.GeneratedFile, method *protogen.Method) string
299304
if !method.Desc.IsStreamingClient() && !method.Desc.IsStreamingServer() {
300305
s += "*" + g.QualifiedGoIdent(method.Output.GoIdent)
301306
} else {
302-
s += method.Parent.GoName + "_" + method.GoName + "Client"
307+
if *useGenericStreams {
308+
s += clientStreamInterface(g, method)
309+
} else {
310+
s += method.Parent.GoName + "_" + method.GoName + "Client"
311+
}
303312
}
304313
s += ", error)"
305314
return s
306315
}
307316

317+
func clientStreamInterface(g *protogen.GeneratedFile, method *protogen.Method) string {
318+
typeParam := g.QualifiedGoIdent(method.Input.GoIdent) + ", " + g.QualifiedGoIdent(method.Output.GoIdent)
319+
if method.Desc.IsStreamingClient() && method.Desc.IsStreamingServer() {
320+
return g.QualifiedGoIdent(grpcPackage.Ident("BidiStreamingClient")) + "[" + typeParam + "]"
321+
} else if method.Desc.IsStreamingClient() {
322+
return g.QualifiedGoIdent(grpcPackage.Ident("ClientStreamingClient")) + "[" + typeParam + "]"
323+
} else { // i.e. if method.Desc.IsStreamingServer()
324+
return g.QualifiedGoIdent(grpcPackage.Ident("ServerStreamingClient")) + "[" + g.QualifiedGoIdent(method.Output.GoIdent) + "]"
325+
}
326+
}
327+
308328
func genClientMethod(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, method *protogen.Method, index int) {
309329
service := method.Parent
310330
fmSymbol := helper.formatFullMethodSymbol(service, method)
@@ -323,11 +343,17 @@ func genClientMethod(gen *protogen.Plugin, file *protogen.File, g *protogen.Gene
323343
g.P()
324344
return
325345
}
326-
streamType := unexport(service.GoName) + method.GoName + "Client"
346+
347+
streamImpl := unexport(service.GoName) + method.GoName + "Client"
348+
if *useGenericStreams {
349+
typeParam := g.QualifiedGoIdent(method.Input.GoIdent) + ", " + g.QualifiedGoIdent(method.Output.GoIdent)
350+
streamImpl = g.QualifiedGoIdent(grpcPackage.Ident("GenericClientStream")) + "[" + typeParam + "]"
351+
}
352+
327353
serviceDescVar := service.GoName + "_ServiceDesc"
328354
g.P("stream, err := c.cc.NewStream(ctx, &", serviceDescVar, ".Streams[", index, `], `, fmSymbol, `, cOpts...)`)
329355
g.P("if err != nil { return nil, err }")
330-
g.P("x := &", streamType, "{stream}")
356+
g.P("x := &", streamImpl, "{ClientStream: stream}")
331357
if !method.Desc.IsStreamingClient() {
332358
g.P("if err := x.ClientStream.SendMsg(in); err != nil { return nil, err }")
333359
g.P("if err := x.ClientStream.CloseSend(); err != nil { return nil, err }")
@@ -336,11 +362,20 @@ func genClientMethod(gen *protogen.Plugin, file *protogen.File, g *protogen.Gene
336362
g.P("}")
337363
g.P()
338364

365+
// Auxiliary types aliases, for backwards compatibility.
366+
if *useGenericStreams {
367+
g.P("// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.")
368+
g.P("type ", service.GoName, "_", method.GoName, "Client = ", clientStreamInterface(g, method))
369+
g.P()
370+
return
371+
}
372+
373+
// Stream auxiliary types and methods, if we're not taking advantage of the
374+
// pre-implemented generic types and their methods.
339375
genSend := method.Desc.IsStreamingClient()
340376
genRecv := method.Desc.IsStreamingServer()
341377
genCloseAndRecv := !method.Desc.IsStreamingServer()
342378

343-
// Stream auxiliary types and methods.
344379
g.P("type ", service.GoName, "_", method.GoName, "Client interface {")
345380
if genSend {
346381
g.P("Send(*", method.Input.GoIdent, ") error")
@@ -355,27 +390,27 @@ func genClientMethod(gen *protogen.Plugin, file *protogen.File, g *protogen.Gene
355390
g.P("}")
356391
g.P()
357392

358-
g.P("type ", streamType, " struct {")
393+
g.P("type ", streamImpl, " struct {")
359394
g.P(grpcPackage.Ident("ClientStream"))
360395
g.P("}")
361396
g.P()
362397

363398
if genSend {
364-
g.P("func (x *", streamType, ") Send(m *", method.Input.GoIdent, ") error {")
399+
g.P("func (x *", streamImpl, ") Send(m *", method.Input.GoIdent, ") error {")
365400
g.P("return x.ClientStream.SendMsg(m)")
366401
g.P("}")
367402
g.P()
368403
}
369404
if genRecv {
370-
g.P("func (x *", streamType, ") Recv() (*", method.Output.GoIdent, ", error) {")
405+
g.P("func (x *", streamImpl, ") Recv() (*", method.Output.GoIdent, ", error) {")
371406
g.P("m := new(", method.Output.GoIdent, ")")
372407
g.P("if err := x.ClientStream.RecvMsg(m); err != nil { return nil, err }")
373408
g.P("return m, nil")
374409
g.P("}")
375410
g.P()
376411
}
377412
if genCloseAndRecv {
378-
g.P("func (x *", streamType, ") CloseAndRecv() (*", method.Output.GoIdent, ", error) {")
413+
g.P("func (x *", streamImpl, ") CloseAndRecv() (*", method.Output.GoIdent, ", error) {")
379414
g.P("if err := x.ClientStream.CloseSend(); err != nil { return nil, err }")
380415
g.P("m := new(", method.Output.GoIdent, ")")
381416
g.P("if err := x.ClientStream.RecvMsg(m); err != nil { return nil, err }")
@@ -396,7 +431,11 @@ func serverSignature(g *protogen.GeneratedFile, method *protogen.Method) string
396431
reqArgs = append(reqArgs, "*"+g.QualifiedGoIdent(method.Input.GoIdent))
397432
}
398433
if method.Desc.IsStreamingClient() || method.Desc.IsStreamingServer() {
399-
reqArgs = append(reqArgs, method.Parent.GoName+"_"+method.GoName+"Server")
434+
if *useGenericStreams {
435+
reqArgs = append(reqArgs, serverStreamInterface(g, method))
436+
} else {
437+
reqArgs = append(reqArgs, method.Parent.GoName+"_"+method.GoName+"Server")
438+
}
400439
}
401440
return method.GoName + "(" + strings.Join(reqArgs, ", ") + ") " + ret
402441
}
@@ -442,6 +481,17 @@ func genServiceDesc(file *protogen.File, g *protogen.GeneratedFile, serviceDescV
442481
g.P()
443482
}
444483

484+
func serverStreamInterface(g *protogen.GeneratedFile, method *protogen.Method) string {
485+
typeParam := g.QualifiedGoIdent(method.Input.GoIdent) + ", " + g.QualifiedGoIdent(method.Output.GoIdent)
486+
if method.Desc.IsStreamingClient() && method.Desc.IsStreamingServer() {
487+
return g.QualifiedGoIdent(grpcPackage.Ident("BidiStreamingServer")) + "[" + typeParam + "]"
488+
} else if method.Desc.IsStreamingClient() {
489+
return g.QualifiedGoIdent(grpcPackage.Ident("ClientStreamingServer")) + "[" + typeParam + "]"
490+
} else { // i.e. if method.Desc.IsStreamingServer()
491+
return g.QualifiedGoIdent(grpcPackage.Ident("ServerStreamingServer")) + "[" + g.QualifiedGoIdent(method.Output.GoIdent) + "]"
492+
}
493+
}
494+
445495
func genServerMethod(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, method *protogen.Method, hnameFuncNameFormatter func(string) string) string {
446496
service := method.Parent
447497
hname := fmt.Sprintf("_%s_%s_Handler", service.GoName, method.GoName)
@@ -464,23 +514,38 @@ func genServerMethod(gen *protogen.Plugin, file *protogen.File, g *protogen.Gene
464514
g.P()
465515
return hname
466516
}
467-
streamType := unexport(service.GoName) + method.GoName + "Server"
517+
518+
streamImpl := unexport(service.GoName) + method.GoName + "Server"
519+
if *useGenericStreams {
520+
typeParam := g.QualifiedGoIdent(method.Input.GoIdent) + ", " + g.QualifiedGoIdent(method.Output.GoIdent)
521+
streamImpl = g.QualifiedGoIdent(grpcPackage.Ident("GenericServerStream")) + "[" + typeParam + "]"
522+
}
523+
468524
g.P("func ", hnameFuncNameFormatter(hname), "(srv interface{}, stream ", grpcPackage.Ident("ServerStream"), ") error {")
469525
if !method.Desc.IsStreamingClient() {
470526
g.P("m := new(", method.Input.GoIdent, ")")
471527
g.P("if err := stream.RecvMsg(m); err != nil { return err }")
472-
g.P("return srv.(", service.GoName, "Server).", method.GoName, "(m, &", streamType, "{stream})")
528+
g.P("return srv.(", service.GoName, "Server).", method.GoName, "(m, &", streamImpl, "{ServerStream: stream})")
473529
} else {
474-
g.P("return srv.(", service.GoName, "Server).", method.GoName, "(&", streamType, "{stream})")
530+
g.P("return srv.(", service.GoName, "Server).", method.GoName, "(&", streamImpl, "{ServerStream: stream})")
475531
}
476532
g.P("}")
477533
g.P()
478534

535+
// Auxiliary types aliases, for backwards compatibility.
536+
if *useGenericStreams {
537+
g.P("// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.")
538+
g.P("type ", service.GoName, "_", method.GoName, "Server = ", serverStreamInterface(g, method))
539+
g.P()
540+
return hname
541+
}
542+
543+
// Stream auxiliary types and methods, if we're not taking advantage of the
544+
// pre-implemented generic types and their methods.
479545
genSend := method.Desc.IsStreamingServer()
480546
genSendAndClose := !method.Desc.IsStreamingServer()
481547
genRecv := method.Desc.IsStreamingClient()
482548

483-
// Stream auxiliary types and methods.
484549
g.P("type ", service.GoName, "_", method.GoName, "Server interface {")
485550
if genSend {
486551
g.P("Send(*", method.Output.GoIdent, ") error")
@@ -495,25 +560,25 @@ func genServerMethod(gen *protogen.Plugin, file *protogen.File, g *protogen.Gene
495560
g.P("}")
496561
g.P()
497562

498-
g.P("type ", streamType, " struct {")
563+
g.P("type ", streamImpl, " struct {")
499564
g.P(grpcPackage.Ident("ServerStream"))
500565
g.P("}")
501566
g.P()
502567

503568
if genSend {
504-
g.P("func (x *", streamType, ") Send(m *", method.Output.GoIdent, ") error {")
569+
g.P("func (x *", streamImpl, ") Send(m *", method.Output.GoIdent, ") error {")
505570
g.P("return x.ServerStream.SendMsg(m)")
506571
g.P("}")
507572
g.P()
508573
}
509574
if genSendAndClose {
510-
g.P("func (x *", streamType, ") SendAndClose(m *", method.Output.GoIdent, ") error {")
575+
g.P("func (x *", streamImpl, ") SendAndClose(m *", method.Output.GoIdent, ") error {")
511576
g.P("return x.ServerStream.SendMsg(m)")
512577
g.P("}")
513578
g.P()
514579
}
515580
if genRecv {
516-
g.P("func (x *", streamType, ") Recv() (*", method.Input.GoIdent, ", error) {")
581+
g.P("func (x *", streamImpl, ") Recv() (*", method.Input.GoIdent, ", error) {")
517582
g.P("m := new(", method.Input.GoIdent, ")")
518583
g.P("if err := x.ServerStream.RecvMsg(m); err != nil { return nil, err }")
519584
g.P("return m, nil")

‎cmd/protoc-gen-go-grpc/main.go

+2
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import (
4444
const version = "1.3.0"
4545

4646
var requireUnimplemented *bool
47+
var useGenericStreams *bool
4748

4849
func main() {
4950
showVersion := flag.Bool("version", false, "print the version and exit")
@@ -55,6 +56,7 @@ func main() {
5556

5657
var flags flag.FlagSet
5758
requireUnimplemented = flags.Bool("require_unimplemented_servers", true, "set to false to match legacy behavior")
59+
useGenericStreams = flags.Bool("use_generic_streams_experimental", false, "set to true to use generic types for streaming client and server objects; this flag is EXPERIMENTAL and may be changed or removed in a future release")
5860

5961
protogen.Options{
6062
ParamFunc: flags.Set,

‎cmd/protoc-gen-go-grpc/protoc-gen-go-grpc_test.sh

+1-1
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ popd
3030

3131
protoc \
3232
--go-grpc_out="${TEMPDIR}" \
33-
--go-grpc_opt=paths=source_relative \
33+
--go-grpc_opt=paths=source_relative,use_generic_streams_experimental=true \
3434
"examples/route_guide/routeguide/route_guide.proto"
3535

3636
GOLDENFILE="examples/route_guide/routeguide/route_guide_grpc.pb.go"

‎credentials/alts/internal/proto/grpc_gcp/handshaker_grpc.pb.go

+12-50
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)