plugin: Add API for plugin authors (#185)

Plugin authors don't need to know how the plugin is talking to ThriftRW. With
this API, they can simply implement a `main` function that calls `plugin.Main`
with their plugin definition, and we handle setting up the communication with
ThriftRw.

plugin/doc.go

@@ -0,0 +1,40 @@
+// Copyright (c) 2016 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+// Package plugin provides the API for writing ThriftRW plugins.
+//
+// Plugins are standalone programs with names in the format
+// thriftrw-plugin-$name where $name is the name of the plugin.
+//
+// 	// thriftrw-plugin-myfancyplugin/main.go
+// 	package main
+//
+// 	import "github.com/thriftrw/thriftrw-go/plugin"
+//
+// 	func main() {
+// 		plugin.Main(&plugin.Plugin{
+// 			Name: "myfancyplugin",
+// 			// ...
+// 		})
+// 	}
+//
+// Note that the name in the executable MUST match the name in the Plugin
+// struct.
+package plugin

plugin/mock_service_generator_test.go

@@ -0,0 +1,41 @@
+// Automatically generated by MockGen. DO NOT EDIT!
+// Source: github.com/thriftrw/thriftrw-go/plugin/api (interfaces: ServiceGenerator)
+
+package plugin
+
+import (
+	gomock "github.com/golang/mock/gomock"
+	api "github.com/thriftrw/thriftrw-go/plugin/api"
+)
+
+// Mock of ServiceGenerator interface
+type MockServiceGenerator struct {
+	ctrl     *gomock.Controller
+	recorder *_MockServiceGeneratorRecorder
+}
+
+// Recorder for MockServiceGenerator (not exported)
+type _MockServiceGeneratorRecorder struct {
+	mock *MockServiceGenerator
+}
+
+func NewMockServiceGenerator(ctrl *gomock.Controller) *MockServiceGenerator {
+	mock := &MockServiceGenerator{ctrl: ctrl}
+	mock.recorder = &_MockServiceGeneratorRecorder{mock}
+	return mock
+}
+
+func (_m *MockServiceGenerator) EXPECT() *_MockServiceGeneratorRecorder {
+	return _m.recorder
+}
+
+func (_m *MockServiceGenerator) Generate(_param0 *api.GenerateServiceRequest) (*api.GenerateServiceResponse, error) {
+	ret := _m.ctrl.Call(_m, "Generate", _param0)
+	ret0, _ := ret[0].(*api.GenerateServiceResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+func (_mr *_MockServiceGeneratorRecorder) Generate(arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "Generate", arg0)
+}

plugin/plugin.go

@@ -0,0 +1,105 @@
+// Copyright (c) 2016 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package plugin
+
+import (
+	"io"
+	"log"
+	"os"
+
+	"github.com/thriftrw/thriftrw-go/internal/envelope"
+	"github.com/thriftrw/thriftrw-go/internal/frame"
+	"github.com/thriftrw/thriftrw-go/internal/multiplex"
+	"github.com/thriftrw/thriftrw-go/plugin/api"
+	"github.com/thriftrw/thriftrw-go/plugin/api/service/plugin"
+	"github.com/thriftrw/thriftrw-go/plugin/api/service/servicegenerator"
+	"github.com/thriftrw/thriftrw-go/protocol"
+)
+
+const _fastPathFrameSize = 10 * 1024 * 1024 // 10 MB
+
+var (
+	_proto           = protocol.Binary
+	_out   io.Writer = os.Stdout
+	_in    io.Reader = os.Stdin
+)
+
+// Plugin defines a ThriftRW plugin.
+type Plugin struct {
+	// Name of the plugin. The name of the executable providing this plugin MUST
+	// BE thriftrw-plugin-$name.
+	Name string
+
+	// If non-nil, this indicates that the plugin will generate code for Thrift
+	// services.
+	ServiceGenerator api.ServiceGenerator
+}
+
+// Main serves the given plugin. It is the entry point to the plugin system.
+// User-defined plugins should call Main with their main function.
+func Main(p *Plugin) {
+	// The plugin communicates with the ThriftRW process over stdout and stdin
+	// of this process. Requests and responses are Thrift envelopes with a
+	// 4-byte big-endian encoded length prefix. Envelope names contain method
+	// names prefixed with the service name and a ":".
+
+	mainHandler := multiplex.NewHandler()
+
+	features := []api.Feature{}
+
+	if p.ServiceGenerator != nil {
+		features = append(features, api.FeatureServiceGenerator)
+		mainHandler.Put("ServiceGenerator", servicegenerator.NewHandler(p.ServiceGenerator))
+	}
+
+	// TODO(abg): Check for other features and register handlers here.
+
+	server := frame.NewServer(_in, _out)
+	mainHandler.Put("Plugin", plugin.NewHandler(pluginHandler{
+		server:   server,
+		plugin:   p,
+		features: features,
+	}))
+
+	if err := server.Serve(envelope.NewServer(_proto, mainHandler)); err != nil {
+		log.Fatalf("plugin server failed with error: %v", err)
+	}
+}
+
+// pluginHandler implements the Plugin service.
+type pluginHandler struct {
+	server   *frame.Server
+	plugin   *Plugin
+	features []api.Feature
+}
+
+func (h pluginHandler) Handshake(request *api.HandshakeRequest) (*api.HandshakeResponse, error) {
+	return &api.HandshakeResponse{
+		Name:       h.plugin.Name,
+		ApiVersion: api.Version,
+		Features:   h.features,
+	}, nil
+}
+
+func (h pluginHandler) Goodbye() error {
+	h.server.Stop()
+	return nil
+}

plugin/plugin_test.go

@@ -0,0 +1,138 @@
+package plugin
+
+import (
+	"io"
+	"testing"
+
+	"github.com/thriftrw/thriftrw-go/internal/envelope"
+	"github.com/thriftrw/thriftrw-go/internal/frame"
+	"github.com/thriftrw/thriftrw-go/internal/multiplex"
+	"github.com/thriftrw/thriftrw-go/plugin/api"
+	"github.com/thriftrw/thriftrw-go/plugin/api/service/plugin"
+	"github.com/thriftrw/thriftrw-go/plugin/api/service/servicegenerator"
+	"github.com/thriftrw/thriftrw-go/ptr"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+//go:generate mockgen -destination mock_service_generator_test.go -package plugin github.com/thriftrw/thriftrw-go/plugin/api ServiceGenerator
+
+// fakeStreams is a helper for tests to control the output and input used by
+// plugin.Main while testing.
+//
+// It returns a Writer to write to the plugin's stdin, a Reader to read from it,
+// and a function that should be called after the test is finished to restore
+// the old values.
+//
+// 	in, out, done := fakeStreams()
+// 	defer done()
+func fakeStreams() (stdin io.Writer, stdout io.Reader, done func()) {
+	stdinReader, stdinWriter := io.Pipe()
+	stdoutReader, stdoutWriter := io.Pipe()
+
+	oldIn := _in
+	oldOut := _out
+	_in = stdinReader
+	_out = stdoutWriter
+
+	return stdinWriter, stdoutReader, func() {
+		_in = oldIn
+		_out = oldOut
+	}
+}
+
+func fakeEnvelopeClient() (envelope.Client, func()) {
+	in, out, done := fakeStreams()
+	return envelope.NewClient(_proto, frame.NewClient(in, out)), done
+}
+
+func TestEmptyPlugin(t *testing.T) {
+	transport, done := fakeEnvelopeClient()
+	defer done()
+
+	go Main(&Plugin{Name: "hello"})
+
+	client := plugin.NewClient(multiplex.NewClient("Plugin", transport))
+
+	response, err := client.Handshake(&api.HandshakeRequest{})
+	require.NoError(t, err)
+	assert.Equal(t, api.Version, response.ApiVersion)
+	assert.Equal(t, "hello", response.Name)
+	assert.Empty(t, response.Features)
+
+	assert.NoError(t, client.Goodbye())
+}
+
+func TestServiceGenerator(t *testing.T) {
+	transport, done := fakeEnvelopeClient()
+	defer done()
+
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	serviceGenerator := NewMockServiceGenerator(mockCtrl)
+
+	go Main(&Plugin{
+		Name:             "hello",
+		ServiceGenerator: serviceGenerator,
+	})
+
+	pluginClient := plugin.NewClient(multiplex.NewClient("Plugin", transport))
+	defer pluginClient.Goodbye()
+
+	handshake, err := pluginClient.Handshake(&api.HandshakeRequest{})
+	require.NoError(t, err)
+	assert.Equal(t, api.Version, handshake.ApiVersion)
+	assert.Equal(t, "hello", handshake.Name)
+	assert.Contains(t, handshake.Features, api.FeatureServiceGenerator)
+
+	sgClient := servicegenerator.NewClient(multiplex.NewClient("ServiceGenerator", transport))
+	req := &api.GenerateServiceRequest{
+		RootServices: []api.ServiceID{1},
+		Services: map[api.ServiceID]*api.Service{
+			1: {
+				Name:      "MyService",
+				Package:   "github.com/thriftrw/thriftrw-go/plugin/fake/myservice",
+				Directory: "fake/myservice",
+				Functions: []*api.Function{},
+				ParentID:  (*api.ServiceID)(ptr.Int32(2)),
+				ModuleID:  1,
+			},
+			2: {
+				Name:      "BaseService",
+				Package:   "github.com/thriftrw/thriftrw-go/plugin/fake/baseservice",
+				Directory: "fake/baseservice",
+				Functions: []*api.Function{
+					{
+						Name:       "Healthy",
+						ThriftName: "healthy",
+						Arguments:  []*api.Argument{},
+					},
+				},
+				ModuleID: 1,
+			},
+		},
+		Modules: map[api.ModuleID]*api.Module{
+			1: {
+				Package:   "github.com/thriftrw/thriftrw-go/plugin/fake",
+				Directory: "fake",
+			},
+		},
+	}
+
+	res := &api.GenerateServiceResponse{
+		Files: map[string][]byte{
+			"fake/myservice/foo.go":   {1, 2, 3},
+			"fake/baseservice/bar.go": {4, 5, 6},
+			"fake/baz.go":             {7, 8, 9},
+		},
+	}
+
+	serviceGenerator.EXPECT().Generate(req).Return(res, nil)
+	gotRes, err := sgClient.Generate(req)
+	if assert.NoError(t, err) {
+		assert.Equal(t, res, gotRes)
+	}
+}