Fix multiple issues with a driver pool

cmd/bblfshd/main.go

@@ -5,12 +5,15 @@ import (
 	"flag"
 	"fmt"
 	"net"
+	"net/http"
 	"os"
 	"os/signal"
 	"sync"
 	"syscall"
 	"time"
 
+	_ "net/http/pprof"
+
 	"github.com/bblfsh/bblfshd/daemon"
 	"github.com/bblfsh/bblfshd/runtime"
 
@@ -46,6 +49,10 @@ var (
 		format *string
 		fields *string
 	}
+	pprof struct {
+		enabled *bool
+		address *string
+	}
 	cmd *flag.FlagSet
 
 	usrListener net.Listener
@@ -66,6 +73,8 @@ func init() {
 	log.level = cmd.String("log-level", "info", "log level: panic, fatal, error, warning, info, debug.")
 	log.format = cmd.String("log-format", "text", "format of the logs: text or json.")
 	log.fields = cmd.String("log-fields", "", "extra fields to add to every log line in json format.")
+	pprof.enabled = cmd.Bool("profiler", false, "run profiler http endpoint (pprof).")
+	pprof.address = cmd.String("profiler-address", ":6060", "profiler address to listen on.")
 	cmd.Parse(os.Args[1:])
 
 	buildLogger()
@@ -101,6 +110,15 @@ func installRecommended(d *daemon.Daemon) error {
 func main() {
 	logrus.Infof("bblfshd version: %s (build: %s)", version, build)
 
+	if *pprof.enabled {
+		logrus.Infof("running pprof on %s", *pprof.address)
+		go func() {
+			if err := http.ListenAndServe(*pprof.address, nil); err != nil {
+				logrus.Errorf("cannot start pprof: %v", err)
+			}
+		}()
+	}
+
 	c, err := jaegercfg.FromEnv()
 	if err != nil {
 		logrus.Errorf("error configuring tracer: %s", err)

daemon/common_test.go

@@ -1,6 +1,7 @@
 package daemon
 
 import (
+	"context"
 	"encoding/hex"
 	"fmt"
 	"os"
@@ -44,7 +45,7 @@ func (d *mockDriver) ServiceV2() protocol2.DriverClient {
 	return nil
 }
 
-func (d *mockDriver) Start() error {
+func (d *mockDriver) Start(ctx context.Context) error {
 	return nil
 }
 

daemon/daemon.go

@@ -29,7 +29,7 @@ type Daemon struct {
 	runtime   *runtime.Runtime
 	driverEnv []string
 
-	mu   sync.Mutex
+	mu   sync.RWMutex
 	pool map[string]*DriverPool
 }
 
@@ -131,10 +131,17 @@ func (d *Daemon) RemoveDriver(language string) error {
 }
 
 func (d *Daemon) DriverPool(ctx context.Context, language string) (*DriverPool, error) {
+	d.mu.RLock()
+	dp, ok := d.pool[language]
+	d.mu.RUnlock()
+	if ok {
+		return dp, nil
+	}
+
 	d.mu.Lock()
 	defer d.mu.Unlock()
-
-	if dp, ok := d.pool[language]; ok {
+	dp, ok = d.pool[language]
+	if ok {
 		return dp, nil
 	}
 
@@ -179,7 +186,7 @@ func (d *Daemon) newDriverPool(rctx context.Context, language string, image runt
 			return nil, err
 		}
 
-		if err := driver.Start(); err != nil {
+		if err := driver.Start(rctx); err != nil {
 			return nil, err
 		}
 
@@ -191,8 +198,12 @@ func (d *Daemon) newDriverPool(rctx context.Context, language string, image runt
 		"language": language,
 	})
 
+	if err := dp.Start(); err != nil {
+		return nil, err
+	}
+
 	d.pool[language] = dp
-	return dp, dp.Start()
+	return dp, nil
 }
 
 func (d *Daemon) removePool(language string) error {
@@ -211,8 +222,8 @@ func (d *Daemon) removePool(language string) error {
 
 // Current returns the current list of driver pools.
 func (d *Daemon) Current() map[string]*DriverPool {
-	d.mu.Lock()
-	defer d.mu.Unlock()
+	d.mu.RLock()
+	defer d.mu.RUnlock()
 
 	out := make(map[string]*DriverPool, len(d.pool))
 	for k, pool := range d.pool {
@@ -224,6 +235,8 @@ func (d *Daemon) Current() map[string]*DriverPool {
 
 // Stop stops all the pools and containers.
 func (d *Daemon) Stop() error {
+	d.mu.Lock()
+	defer d.mu.Unlock()
 	var err error
 	for _, dp := range d.pool {
 		if cerr := dp.Stop(); cerr != nil && err != nil {

daemon/driver.go

@@ -24,7 +24,7 @@ import (
 
 type Driver interface {
 	ID() string
-	Start() error
+	Start(ctx context.Context) error
 	Stop() error
 	Status() (protocol.Status, error)
 	State() (*protocol.DriverInstanceState, error)
@@ -47,10 +47,9 @@ type DriverInstance struct {
 }
 
 const (
-	DriverBinary      = "/opt/driver/bin/driver"
-	GRPCSocket        = "rpc.sock"
-	TmpPathPattern    = "/tmp/%s"
-	ConnectionTimeout = 5 * time.Second
+	DriverBinary   = "/opt/driver/bin/driver"
+	GRPCSocket     = "rpc.sock"
+	TmpPathPattern = "/tmp/%s"
 )
 
 type Options struct {
@@ -116,12 +115,13 @@ func (i *DriverInstance) ID() string {
 }
 
 // Start starts a container and connects to it.
-func (i *DriverInstance) Start() error {
+func (i *DriverInstance) Start(ctx context.Context) error {
 	if err := i.Container.Start(); err != nil {
 		return err
 	}
 
-	if err := i.dial(); err != nil {
+	if err := i.dial(ctx); err != nil {
+		_ = i.Container.Stop()
 		return err
 	}
 
@@ -132,19 +132,24 @@ func (i *DriverInstance) Start() error {
 	return nil
 }
 
-func (i *DriverInstance) dial() error {
+func (i *DriverInstance) dial(ctx context.Context) error {
 	addr := filepath.Join(i.tmp, GRPCSocket)
 
 	opts := []grpc.DialOption{
 		grpc.WithDialer(func(addr string, t time.Duration) (net.Conn, error) {
 			return net.DialTimeout("unix", addr, t)
 		}),
+		// always wait for the connection to become active
 		grpc.WithBlock(),
-		grpc.WithTimeout(ConnectionTimeout),
+		// we want to know sooner rather than later
+		// TODO(dennwc): sometimes the initialization of the container takes >5 sec
+		//               meaning that the time between Container.Start and the actual
+		//               execution of a Go server (not the native driver) takes this long
+		grpc.WithBackoffMaxDelay(time.Second),
 		grpc.WithInsecure(),
 	}
 	opts = append(opts, protocol2.DialOptions()...)
-	conn, err := grpc.Dial(addr, opts...)
+	conn, err := grpc.DialContext(ctx, addr, opts...)
 
 	i.conn = conn
 	i.srv1 = protocol1.NewProtocolServiceClient(conn)

daemon/driver_test.go

@@ -1,6 +1,7 @@
 package daemon
 
 import (
+	"context"
 	"io/ioutil"
 	"os"
 	"testing"
@@ -32,7 +33,10 @@ func TestNewDriver(t *testing.T) {
 
 	require.NoError(err)
 
-	err = i.Start()
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	err = i.Start(ctx)
 	require.NoError(err)
 
 	time.Sleep(50 * time.Millisecond)
@@ -60,7 +64,10 @@ func TestDriverInstance_State(t *testing.T) {
 	require.Len(state.Processes, 0)
 	require.True(state.Created.IsZero())
 
-	err = i.Start()
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	err = i.Start(ctx)
 	require.NoError(err)
 	defer func() {
 		err = i.Stop()

daemon/pool_test.go

@@ -92,16 +92,18 @@ func TestDriverPoolExecute_Recovery(t *testing.T) {
 	err := dp.Start()
 	require.NoError(err)
 
+	ctx := context.Background()
+
 	for i := 0; i < 100; i++ {
-		err := dp.Execute(func(_ context.Context, d Driver) error {
+		err := dp.ExecuteCtx(ctx, func(_ context.Context, d Driver) error {
 			require.NotNil(d)
 
 			if i%10 == 0 {
 				d.(*mockDriver).MockStatus = protocol.Stopped
 			}
 
 			return nil
-		}, 0)
+		})
 
 		require.Nil(err)
 		require.Len(dp.Current(), 1)