MultipleDriverClient could be useful during language-specific parsing…

…s on the large scale

Scenario: we need to parse a ton of go and python files inside the k8s environment, to save time we need to perform this parses on the large scale of go- and python-driver containers/instances etc.

Solution:
- run two separate Deployments of go- and python-driver container pods
- provide Horizontal Autoscalers for both of Deployments
- provide Services with LoadBalancer type
- during client initialization provide Services endpoints configuration, this will create two language-oriented connections
that will be responsible for sending parse request only to a dedicated Service, that will load-balance this requestbetween underlying language driver pods

Examples of endpoint formats:
- localhost:9432 - casual example there's only one driver or bblfshd server
- python=localhost:9432,go=localhost:9432 - coma-separated mapping in format language=address
- %s-driver.bblfsh.svc.example.com - DNS template based on the language

Signed-off-by: lwsanty <[email protected]>

.travis.yml

@@ -8,6 +8,7 @@ install:
     if [[ $TRAVIS_OS_NAME = linux ]]; then
       docker run --privileged -d -p 9432:9432 --name bblfshd bblfsh/bblfshd
       docker exec bblfshd bblfshctl driver install bblfsh/python-driver
+      docker exec bblfshd bblfshctl driver install bblfsh/go-driver
     fi
   - go mod download
 
@@ -21,7 +22,7 @@ script:
 
 jobs:
   include:
-    - {go: 1.11.x,  os: linux, sudo: required, dist: trusty, services: [docker]}
-    - {go: 1.12.x, os: linux, sudo: required, dist: trusty, services: [docker]}
-    - {go: 1.11.x,  os: osx, osx_image: xcode9.3}
-    - {go: 1.12.x, os: osx, osx_image: xcode9.3}
+    - {go: 1.12.x,  os: linux, sudo: required, dist: bionic, services: [docker]}
+    - {go: 1.13.x, os: linux, sudo: required, dist: bionic, services: [docker]}
+    - {go: 1.12.x,  os: osx, osx_image: xcode11.2}
+    - {go: 1.13.x, os: osx, osx_image: xcode11.2}

client.go

@@ -2,17 +2,20 @@ package bblfsh
 
 import (
 	"context"
+	"fmt"
+	"io"
+	"strings"
 	"time"
 
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/keepalive"
-	"google.golang.org/grpc/status"
-
 	"github.com/bblfsh/sdk/v3/driver"
 	"github.com/bblfsh/sdk/v3/driver/manifest"
 	protocol2 "github.com/bblfsh/sdk/v3/protocol"
 	protocol1 "gopkg.in/bblfsh/sdk.v1/protocol"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/keepalive"
+	"google.golang.org/grpc/status"
 )
 
 const (
@@ -28,9 +31,11 @@ const (
 	keepalivePingWithoutStream = true
 )
 
+type getConnFunc func(ctx context.Context, language string) (*grpc.ClientConn, error)
+
 // Client holds the public client API to interact with the bblfsh daemon.
 type Client struct {
-	*grpc.ClientConn
+	closer  io.Closer
 	driver2 protocol2.DriverClient
 	driver  driver.Driver
 }
@@ -50,11 +55,62 @@ func NewClientContext(ctx context.Context, endpoint string, options ...grpc.Dial
 	// this allows to override any default option
 	opts = append(opts, options...)
 
-	conn, err := grpc.DialContext(ctx, endpoint, opts...)
-	if err != nil {
-		return nil, err
+	switch {
+	case strings.Contains(endpoint, ","):
+		endpoints, err := parseEndpoints(endpoint)
+		if err != nil {
+			return nil, err
+		}
+		return NewClientWithConnectionsContext(func(ctx context.Context, lang string) (*grpc.ClientConn, error) {
+			e, ok := endpoints[lang]
+			if !ok {
+				return nil, &driver.ErrMissingDriver{Language: lang}
+			}
+			conn, err := grpc.DialContext(ctx, e, opts...)
+			if err != nil {
+				return nil, err
+			}
+
+			return conn, nil
+		})
+	case strings.Contains(endpoint, "%s"):
+		return NewClientWithConnectionsContext(func(ctx context.Context, lang string) (*grpc.ClientConn, error) {
+			conn, err := grpc.DialContext(ctx, fmt.Sprintf(endpoint, lang), opts...)
+			if err != nil {
+				return nil, err
+			}
+
+			return conn, nil
+		})
+	default:
+		conn, err := grpc.DialContext(ctx, endpoint, opts...)
+		if err != nil {
+			return nil, err
+		}
+		return NewClientWithConnectionContext(ctx, conn)
 	}
-	return NewClientWithConnectionContext(ctx, conn)
+}
+
+func NewClientWithConnectionsContext(getConn getConnFunc) (*Client, error) {
+	dc := newMultipleDriverClient(getConn)
+
+	return &Client{
+		driver2: dc,
+		driver:  protocol2.DriverFromClient(dc, &multipleDriverHostClient{}),
+	}, nil
+}
+
+func parseEndpoints(endpoints string) (map[string]string, error) {
+	result := make(map[string]string)
+	pairs := strings.Split(endpoints, ",")
+	for _, p := range pairs {
+		vals := strings.Split(p, "=")
+		if len(vals) != 2 {
+			return nil, fmt.Errorf("formatting is broken in section: %q", p)
+		}
+		result[vals[0]] = vals[1]
+	}
+	return result, nil
 }
 
 // NewClient is the same as NewClientContext, but assumes a default timeout for the connection.
@@ -89,21 +145,21 @@ func NewClientWithConnectionContext(ctx context.Context, conn *grpc.ClientConn)
 	if err == nil {
 		// supports v2
 		return &Client{
-			ClientConn: conn,
-			driver2:    protocol2.NewDriverClient(conn),
-			driver:     protocol2.AsDriver(conn),
+			closer:  conn,
+			driver2: protocol2.NewDriverClient(conn),
+			driver:  protocol2.AsDriver(conn),
 		}, nil
 	} else if !isServiceNotSupported(err) {
 		return nil, err
 	}
 	s1 := protocol1.NewProtocolServiceClient(conn)
 	return &Client{
-		ClientConn: conn,
-		driver2:    protocol2.NewDriverClient(conn),
+		closer:  conn,
+		driver2: protocol2.NewDriverClient(conn),
 		driver: &driverPartialV2{
 			// use only Parse from v2
 			Driver: protocol2.AsDriver(conn),
-			// use v1 for version and supported languages
+			// use v1 for version and supported langConns
 			service1: s1,
 		},
 	}, nil
@@ -128,7 +184,7 @@ func (d *driverPartialV2) Version(ctx context.Context) (driver.Version, error) {
 	}, nil
 }
 
-// Languages implements a driver.Host using v1 protocol.
+// langConns implements a driver.Host using v1 protocol.
 func (d *driverPartialV2) Languages(ctx context.Context) ([]manifest.Manifest, error) {
 	resp, err := d.service1.SupportedLanguages(ctx, &protocol1.SupportedLanguagesRequest{})
 	if err != nil {
@@ -163,7 +219,14 @@ func (c *Client) NewVersionRequest() *VersionRequest {
 	return &VersionRequest{ctx: context.Background(), client: c}
 }
 
-// NewSupportedLanguagesRequest is a parsing request to get the supported languages.
+// NewSupportedLanguagesRequest is a parsing request to get the supported langConns.
 func (c *Client) NewSupportedLanguagesRequest() *SupportedLanguagesRequest {
 	return &SupportedLanguagesRequest{ctx: context.Background(), client: c}
 }
+
+func (c *Client) Close() error {
+	if c.closer == nil {
+		return nil
+	}
+	return c.closer.Close()
+}

client_test.go

@@ -8,13 +8,13 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func newClient(t testing.TB) *Client {
+func newClient(t testing.TB, endpoint string) *Client {
 	if testing.Short() {
 		t.Skip("skipping integration test")
 	}
 	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
 	defer cancel()
-	cli, err := NewClientContext(ctx, "localhost:9432")
+	cli, err := NewClientContext(ctx, endpoint)
 	if err == context.DeadlineExceeded {
 		t.Skip("bblfshd is not running")
 	}
@@ -34,7 +34,7 @@ var clientTests = []struct {
 }
 
 func TestClient(t *testing.T) {
-	cli := newClient(t)
+	cli := newClient(t, "localhost:9432")
 	for _, c := range clientTests {
 		c := c
 		t.Run(c.name, func(t *testing.T) {
@@ -43,6 +43,15 @@ func TestClient(t *testing.T) {
 	}
 }
 
+func TestMultiConnections(t *testing.T) {
+	cli := newClient(t, "python=localhost:9432,go=localhost:9432")
+
+	// it's not a mistake that we run 2 same requests, it checks the actual map of already initialized connections
+	testNativeParseRequestCustom(t, cli, "python", "import foo")
+	testNativeParseRequestCustom(t, cli, "python", "import foo")
+	testNativeParseRequestCustom(t, cli, "go", "package main")
+}
+
 func testParseRequest(t *testing.T, cli *Client) {
 	res, err := cli.NewParseRequest().Language("python").Content("import foo").Do()
 	require.NoError(t, err)
@@ -67,7 +76,11 @@ func testParseRequestMode(t *testing.T, cli *Client) {
 }
 
 func testNativeParseRequest(t *testing.T, cli *Client) {
-	res, err := cli.NewParseRequest().Mode(Native).Language("python").Content("import foo").Do()
+	testNativeParseRequestCustom(t, cli, "python", "import foo")
+}
+
+func testNativeParseRequestCustom(t *testing.T, cli *Client, lang, content string) {
+	res, err := cli.NewParseRequest().Mode(Native).Language(lang).Content(content).Do()
 	require.NoError(t, err)
 
 	require.Equal(t, 0, len(res.Errors))

driver-client.go

@@ -0,0 +1,96 @@
+package bblfsh
+
+import (
+	"context"
+
+	protocol2 "github.com/bblfsh/sdk/v3/protocol"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+/*
+	multipleDriverClient could be useful during language-specific parsings on the large scale
+
+	Examples of endpoint formats:
+	- localhost:9432 - casual example there's only one driver or bblfshd server
+	- python=localhost:9432,go=localhost:9432 - coma-separated mapping in format language=address
+	- %s-driver.bblfsh.svc.example.com - DNS template based on the language
+*/
+
+// multipleDriverClient is a DriverClient implementation, contains connection getter and a map[language]connection
+type multipleDriverClient struct {
+	getConn        getConnFunc
+	langConnDriver langConnDriver
+}
+
+type langConnDriver map[string]*connDriver
+
+type connDriver struct {
+	conn   *grpc.ClientConn
+	driver protocol2.DriverClient
+}
+
+// multipleDriverHostClient is a DriverHostClient implementation, currently does almost nothing
+type multipleDriverHostClient struct{}
+
+// newMultipleDriverClient is a multipleDriverClient constructor
+func newMultipleDriverClient(getConn getConnFunc) *multipleDriverClient {
+	return &multipleDriverClient{
+		getConn:        getConn,
+		langConnDriver: make(langConnDriver),
+	}
+}
+
+// Parse gets connection from a given map, or creates a new connection, then inits driver client and performs Parse
+func (c *multipleDriverClient) Parse(
+	ctx context.Context,
+	in *protocol2.ParseRequest,
+	opts ...grpc.CallOption) (*protocol2.ParseResponse, error) {
+	lang := in.Language
+
+	connD, ok := c.langConnDriver[lang]
+	if !ok {
+		gConn, err := c.getConn(ctx, lang)
+		if err != nil {
+			return nil, err
+		}
+		connD = &connDriver{
+			conn: gConn,
+		}
+	}
+
+	if connD.driver == nil {
+		connD.driver = protocol2.NewDriverClient(connD.conn)
+	}
+	c.langConnDriver[lang] = connD
+
+	return connD.driver.Parse(ctx, in, opts...)
+}
+
+func (c *multipleDriverClient) Close() error {
+	var lastErr error
+	for k, v := range c.langConnDriver {
+		if err := v.conn.Close(); err != nil {
+			lastErr = err
+		}
+		delete(c.langConnDriver, k)
+	}
+	c.langConnDriver = make(langConnDriver)
+	return lastErr
+}
+
+func (hc *multipleDriverHostClient) ServerVersion(
+	ctx context.Context,
+	in *protocol2.VersionRequest,
+	opts ...grpc.CallOption) (*protocol2.VersionResponse, error) {
+	return nil, status.Error(codes.Unimplemented, "ServerVersion is not implemented")
+}
+
+func (hc *multipleDriverHostClient) SupportedLanguages(
+	ctx context.Context,
+	in *protocol2.SupportedLanguagesRequest,
+	opts ...grpc.CallOption) (*protocol2.SupportedLanguagesResponse, error) {
+	return nil, status.Error(codes.Unimplemented, "SupportedLanguages is not implemented")
+}

go.mod

@@ -3,9 +3,11 @@ module github.com/bblfsh/go-client/v4
 go 1.12
 
 require (
-	github.com/bblfsh/sdk/v3 v3.2.5
+	github.com/bblfsh/sdk/v3 v3.3.1
 	github.com/jessevdk/go-flags v1.4.0
 	github.com/stretchr/testify v1.3.0
 	google.golang.org/grpc v1.20.1
 	gopkg.in/bblfsh/sdk.v1 v1.17.0
 )
+
+replace github.com/bblfsh/sdk/v3 => github.com/bblfsh/sdk/v3 v3.3.2-0.20191114212015-37abd692204c