Implement a new code generator that utilizes google.golang.org/protobuf

client.go

@@ -26,7 +26,6 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/gogo/protobuf/proto"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc/codes"

cmd/protoc-gen-go-ttrpc/generator.go

@@ -0,0 +1,109 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package main
+
+import (
+	"google.golang.org/protobuf/compiler/protogen"
+	"strings"
+)
+
+// generator is a Go code generator that uses ttrpc.Server and ttrpc.Client.
+// Unlike the original gogo version, this doesn't generate serializers for message types and
+// let protoc-gen-go handle them.
+type generator struct {
+	out      *protogen.GeneratedFile
+	contextT string
+	serverT  string
+	clientT  string
+	methodT  string
+}
+
+func newGenerator(out *protogen.GeneratedFile) *generator {
+	gen := generator{out: out}
+	gen.contextT = out.QualifiedGoIdent(protogen.GoIdent{GoImportPath: "context", GoName: "Context"})
+	gen.serverT = out.QualifiedGoIdent(protogen.GoIdent{GoImportPath: "github.com/containerd/ttrpc", GoName: "Server"})
+	gen.clientT = out.QualifiedGoIdent(protogen.GoIdent{GoImportPath: "github.com/containerd/ttrpc", GoName: "Client"})
+	gen.methodT = out.QualifiedGoIdent(protogen.GoIdent{GoImportPath: "github.com/containerd/ttrpc", GoName: "Method"})
+	return &gen
+}
+
+func generate(plugin *protogen.Plugin, input *protogen.File) error {
+	file := plugin.NewGeneratedFile(input.GeneratedFilenamePrefix+"_ttrpc.pb.go", input.GoImportPath)
+	file.P("// Code generated by protoc-gen-go-ttrpc. DO NOT EDIT.")
+	file.P("// source: ", input.Desc.Path())
+	file.P("package ", input.GoPackageName)
+
+	gen := newGenerator(file)
+	for _, service := range input.Services {
+		gen.genService(service)
+	}
+	return nil
+}
+
+func (gen *generator) genService(service *protogen.Service) {
+	fullName := service.Desc.FullName()
+	p := gen.out
+
+	serviceName := service.GoName + "Service"
+	p.P("type ", serviceName, " interface{")
+	for _, method := range service.Methods {
+		p.P(method.GoName,
+			"(ctx ", gen.contextT, ",",
+			"req *", method.Input.GoIdent, ")",
+			"(*", method.Output.GoIdent, ", error)")
+	}
+	p.P("}")
+
+	// registration method
+	p.P("func Register", serviceName, "(srv *", gen.serverT, ", svc ", serviceName, "){")
+	p.P(`srv.Register("`, fullName, `", map[string]`, gen.methodT, "{")
+	for _, method := range service.Methods {
+		p.P(`"`, method.GoName, `": func(ctx `, gen.contextT, ", unmarshal func(interface{}) error)(interface{}, error){")
+		p.P("var req ", method.Input.GoIdent)
+		p.P("if err := unmarshal(&req); err != nil {")
+		p.P("return nil, err")
+		p.P("}")
+		p.P("return svc.", method.GoName, "(ctx, &req)")
+		p.P("},")
+	}
+	p.P("})")
+	p.P("}")
+
+	clientType := service.GoName + "Client"
+	clientStructType := strings.ToLower(clientType[:1]) + clientType[1:]
+	p.P("type ", clientStructType, " struct{")
+	p.P("client *", gen.clientT)
+	p.P("}")
+	p.P("func New", clientType, "(client *", gen.clientT, ")", serviceName, "{")
+	p.P("return &", clientStructType, "{")
+	p.P("client:client,")
+	p.P("}")
+	p.P("}")
+
+	for _, method := range service.Methods {
+		p.P("func (c *", clientStructType, ")", method.GoName, "(",
+			"ctx ", gen.contextT, ",",
+			"req *", method.Input.GoIdent, ")",
+			"(*", method.Output.GoIdent, ", error){")
+		p.P("var resp ", method.Output.GoIdent)
+		p.P(`if err := c.client.Call(ctx, "`, fullName, `", "`, method.Desc.Name(), `", req, &resp); err != nil {`)
+		p.P("return nil, err")
+		p.P("}")
+		p.P("return &resp, nil")
+		p.P("}")
+	}
+}

cmd/protoc-gen-go-ttrpc/main.go

@@ -0,0 +1,35 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package main
+
+import (
+	"google.golang.org/protobuf/compiler/protogen"
+)
+
+func main() {
+	protogen.Options{}.Run(func(gen *protogen.Plugin) error {
+		for _, f := range gen.Files {
+			if !f.Generate {
+				continue
+			}
+			if err := generate(gen, f); err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+}

codec.go

@@ -17,15 +17,24 @@
 package ttrpc
 
 import (
-	"github.com/gogo/protobuf/proto"
 	"github.com/pkg/errors"
+
+	vtproto "github.com/planetscale/vtprotobuf/codec/grpc"
+	gproto "github.com/golang/protobuf/proto"
+	"google.golang.org/grpc/encoding"
 )
 
 type codec struct{}
 
+var proto encoding.Codec
+
+func init() {
+	proto = &vtcodec{vt:&vtproto.Codec{}}
+}
+
 func (c codec) Marshal(msg interface{}) ([]byte, error) {
 	switch v := msg.(type) {
-	case proto.Message:
+	case gproto.Message:
 		return proto.Marshal(v)
 	default:
 		return nil, errors.Errorf("ttrpc: cannot marshal unknown type: %T", msg)
@@ -34,7 +43,7 @@ func (c codec) Marshal(msg interface{}) ([]byte, error) {
 
 func (c codec) Unmarshal(p []byte, msg interface{}) error {
 	switch v := msg.(type) {
-	case proto.Message:
+	case gproto.Message:
 		return proto.Unmarshal(p, v)
 	default:
 		return errors.Errorf("ttrpc: cannot unmarshal into unknown type: %T", msg)

example/Protobuild.toml

@@ -1,6 +1,5 @@
 version = "unstable"
-generator = "gogottrpc"
-plugins = ["ttrpc"]
+plugins = ["go-grpc", "go-ttrpc"]
 
 # Control protoc include paths. Below are usually some good defaults, but feel
 # free to try it without them if it works for your project.
@@ -9,11 +8,6 @@ plugins = ["ttrpc"]
   # treat the root of the project as an include, but this may not be necessary.
   # before = ["./protobuf"]
 
-  # Paths that should be treated as include roots in relation to the vendor
-  # directory. These will be calculated with the vendor directory nearest the
-  # target package.
-  packages = ["github.com/gogo/protobuf"]
-
   # Paths that will be added untouched to the end of the includes. We use
   # `/usr/local/include` to pickup the common install location of protobuf.
   # This is the default.

example/cmd/main.go

@@ -26,7 +26,7 @@ import (
 
 	ttrpc "github.com/containerd/ttrpc"
 	"github.com/containerd/ttrpc/example"
-	"github.com/gogo/protobuf/types"
+	"google.golang.org/protobuf/types/known/emptypb"
 )
 
 const socket = "example-ttrpc-server"
@@ -130,6 +130,6 @@ func (s *exampleServer) Method1(ctx context.Context, r *example.Method1Request)
 	}, nil
 }
 
-func (s *exampleServer) Method2(ctx context.Context, r *example.Method1Request) (*types.Empty, error) {
-	return &types.Empty{}, nil
+func (s *exampleServer) Method2(ctx context.Context, r *example.Method1Request) (*emptypb.Empty, error) {
+	return &emptypb.Empty{}, nil
 }