Add heartbeat / automatic server cleanup code

This replaces the previous method of attempting to clean up servers when
an unexpected exit occurs in the client (e.g. SIGINT or panic) by a
heartbeat protocol.  If the server does not hear from the connecting
client within a certain interval of time (500ms in this commit), it will
de-activate itself.  This prevents dangling Docker Machine server
processes from accumulating.

Signed-off-by: Nathan LeClaire <nathan.leclaire@gmail.com>

cmd/machine.go

@@ -3,7 +3,6 @@ package main
 import (
 	"fmt"
 	"os"
-	"os/signal"
 	"path"
 	"strconv"
 
@@ -141,31 +140,10 @@ func main() {
 		},
 	}
 
-	go commands.DeferClosePluginServers()
-
-	// Cleanup to run in case the user sends an interrupt (CTRL+C) to the
-	// Machine program.  Ensure that we do not leave dangling OS processes.
-	signalCh := make(chan os.Signal, 1)
-	signal.Notify(signalCh, os.Interrupt)
-
-	// Atypical exit condition -- write to the cleanup done channel after
-	// ensuring that we have closed all exec-ed plugin servers.
-	go func() {
-		for range signalCh {
-			log.Info("\nReceieved an interrupt, performing cleanup work...")
-			close(commands.RpcClientDriversCh)
-			<-commands.RpcDriversClosedCh
-			os.Exit(1)
-		}
-	}()
-
 	// TODO: Close plugin servers in case of client panic.
 	if err := app.Run(os.Args); err != nil {
 		log.Error(err)
 	}
-
-	close(commands.RpcClientDriversCh)
-	<-commands.RpcDriversClosedCh
 }
 
 func cmdNotFound(c *cli.Context, command string) {

commands/commands.go

@@ -21,9 +21,6 @@ var (
 	ErrUnknownShell       = errors.New("Error: Unknown shell")
 	ErrNoMachineSpecified = errors.New("Error: Expected to get one or more machine names as arguments.")
 	ErrExpectedOneMachine = errors.New("Error: Expected one machine name as an argument.")
-
-	RpcClientDriversCh = make(chan *rpcdriver.RpcClientDriver)
-	RpcDriversClosedCh = make(chan bool)
 )
 
 func newPluginDriver(driverName string, rawContent []byte) (*rpcdriver.RpcClientDriver, error) {
@@ -32,47 +29,14 @@ func newPluginDriver(driverName string, rawContent []byte) (*rpcdriver.RpcClient
 		return nil, err
 	}
 
-	RpcClientDriversCh <- d
-
 	return d, nil
 }
 
-func DeferClosePluginServers() {
-	rpcClientDrivers := []*rpcdriver.RpcClientDriver{}
-
-	for d := range RpcClientDriversCh {
-		rpcClientDrivers = append(rpcClientDrivers, d)
-	}
-
-	doneCh := make(chan bool)
-
-	for _, d := range rpcClientDrivers {
-		d := d
-		go func() {
-			if err := d.Close(); err != nil {
-				log.Debugf("Error closing connection to plugin server: %s", err)
-			}
-
-			doneCh <- true
-		}()
-	}
-
-	for range rpcClientDrivers {
-		<-doneCh
-	}
-
-	RpcDriversClosedCh <- true
-}
-
 func fatal(args ...interface{}) {
-	close(RpcClientDriversCh)
-	<-RpcDriversClosedCh
 	log.Fatal(args...)
 }
 
 func fatalf(fmtString string, args ...interface{}) {
-	close(RpcClientDriversCh)
-	<-RpcDriversClosedCh
 	log.Fatalf(fmtString, args...)
 }
 

commands/create.go

@@ -305,6 +305,10 @@ func cmdCreateOuter(c *cli.Context) {
 		}
 	}
 
+	if err := driver.Close(); err != nil {
+		fatal(err)
+	}
+
 	if err := c.App.Run(os.Args); err != nil {
 		fatal(err)
 	}

libmachine/drivers/plugin/localbinary/plugin.go

@@ -77,7 +77,7 @@ type LocalBinaryExecutor struct {
 
 func NewLocalBinaryPlugin(driverName string) *LocalBinaryPlugin {
 	return &LocalBinaryPlugin{
-		stopCh: make(chan bool, 1),
+		stopCh: make(chan bool),
 		addrCh: make(chan string, 1),
 		Executor: &LocalBinaryExecutor{
 			DriverName: driverName,
@@ -132,13 +132,19 @@ func (lbe *LocalBinaryExecutor) Close() error {
 }
 
 func stream(scanner *bufio.Scanner, streamOutCh chan<- string, stopCh <-chan bool) {
-	for scanner.Scan() {
+	lines := make(chan string)
+	go func() {
+		for scanner.Scan() {
+			lines <- scanner.Text()
+		}
+	}()
+	for {
 		select {
 		case <-stopCh:
 			close(streamOutCh)
 			return
-		default:
-			streamOutCh <- strings.Trim(scanner.Text(), "\n")
+		case line := <-lines:
+			streamOutCh <- strings.Trim(line, "\n")
 			if err := scanner.Err(); err != nil {
 				log.Warnf("Scanning stream: %s", err)
 			}
@@ -148,7 +154,7 @@ func stream(scanner *bufio.Scanner, streamOutCh chan<- string, stopCh <-chan boo
 
 func (lbp *LocalBinaryPlugin) AttachStream(scanner *bufio.Scanner) (<-chan string, chan<- bool) {
 	streamOutCh := make(chan string)
-	stopCh := make(chan bool, 1)
+	stopCh := make(chan bool)
 	go stream(scanner, streamOutCh, stopCh)
 	return streamOutCh, stopCh
 }

libmachine/drivers/plugin/localbinary/plugin_test.go

@@ -106,7 +106,7 @@ func TestExecServer(t *testing.T) {
 		stopCh:      make(chan bool, 1),
 	}
 
-	finalErr := make(chan error, 1)
+	finalErr := make(chan error)
 
 	// Start the docker-machine-foo plugin server
 	go func() {

libmachine/drivers/plugin/register_driver.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/rpc"
 	"os"
+	"time"
 
 	"github.com/docker/machine/libmachine"
 	"github.com/docker/machine/libmachine/drivers"
@@ -13,6 +14,10 @@ import (
 	"github.com/docker/machine/libmachine/drivers/rpc"
 )
 
+var (
+	heartbeatTimeout = 500 * time.Millisecond
+)
+
 func RegisterDriver(d drivers.Driver) {
 	if os.Getenv(localbinary.PluginEnvKey) != localbinary.PluginEnvVal {
 		fmt.Fprintln(os.Stderr, `This is a Docker Machine plugin binary.
@@ -23,11 +28,8 @@ Please use this plugin through the main 'docker-machine' binary.`)
 
 	libmachine.SetDebug(true)
 
-	rpcd := new(rpcdriver.RpcServerDriver)
-	rpcd.ActualDriver = d
-	rpcd.CloseCh = make(chan bool)
+	rpcd := rpcdriver.NewRpcServerDriver(d)
 	rpc.Register(rpcd)
-
 	rpc.HandleHTTP()
 
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
@@ -41,5 +43,14 @@ Please use this plugin through the main 'docker-machine' binary.`)
 
 	go http.Serve(listener, nil)
 
-	<-rpcd.CloseCh
+	for {
+		select {
+		case <-rpcd.CloseCh:
+			os.Exit(0)
+		case <-rpcd.HeartbeatCh:
+			continue
+		case <-time.After(heartbeatTimeout):
+			os.Exit(1)
+		}
+	}
 }

libmachine/drivers/rpc/client_driver.go

@@ -3,6 +3,7 @@ package rpcdriver
 import (
 	"fmt"
 	"net/rpc"
+	"time"
 
 	"github.com/docker/machine/libmachine/drivers"
 	"github.com/docker/machine/libmachine/drivers/plugin/localbinary"
@@ -11,8 +12,14 @@ import (
 	"github.com/docker/machine/libmachine/state"
 )
 
+var (
+	heartbeatInterval = 200 * time.Millisecond
+)
+
 type RpcClientDriver struct {
-	Client *InternalClient
+	plugin          localbinary.DriverPlugin
+	heartbeatDoneCh chan bool
+	Client          *InternalClient
 }
 
 type RpcCall struct {
@@ -22,36 +29,27 @@ type RpcCall struct {
 }
 
 type InternalClient struct {
-	RpcClient *rpc.Client
-	Calls     chan RpcCall
-	CallErrs  chan error
+	MachineName string
+	RpcClient   *rpc.Client
 }
 
 func (ic *InternalClient) Call(serviceMethod string, args interface{}, reply interface{}) error {
-	ic.Calls <- RpcCall{
-		ServiceMethod: serviceMethod,
-		Args:          args,
-		Reply:         reply,
+	if serviceMethod != "RpcServerDriver.Heartbeat" {
+		log.Debugf("(%s) Calling %+v", ic.MachineName, serviceMethod)
 	}
-
-	return <-ic.CallErrs
+	return ic.RpcClient.Call(serviceMethod, args, reply)
 }
 
 func NewInternalClient(rpcclient *rpc.Client) *InternalClient {
 	return &InternalClient{
 		RpcClient: rpcclient,
-		Calls:     make(chan RpcCall),
-		CallErrs:  make(chan error),
 	}
 }
-
 func NewRpcClientDriver(rawDriverData []byte, driverName string) (*RpcClientDriver, error) {
 	mcnName := ""
 
 	p := localbinary.NewLocalBinaryPlugin(driverName)
 
-	c := &RpcClientDriver{}
-
 	go func() {
 		if err := p.Serve(); err != nil {
 			// If we can't safely load the server, best to just
@@ -70,31 +68,24 @@ func NewRpcClientDriver(rawDriverData []byte, driverName string) (*RpcClientDriv
 		return nil, err
 	}
 
-	c.Client = NewInternalClient(rpcclient)
+	c := &RpcClientDriver{
+		Client:          NewInternalClient(rpcclient),
+		heartbeatDoneCh: make(chan bool),
+	}
 
-	go func() {
+	go func(heartbeatDoneCh <-chan bool) {
 		for {
-			call := <-c.Client.Calls
-
-			log.Debugf("(%s) Got msg %+v", mcnName, call)
-
-			if call.ServiceMethod == "RpcServerDriver.Close" {
-				p.Close()
-			}
-
-			c.Client.CallErrs <- c.Client.RpcClient.Call(call.ServiceMethod, call.Args, call.Reply)
-
-			if call.ServiceMethod == "RpcServerDriver.Close" {
-				// If we're messaging the server to close,
-				// we're not accepting any more RPC calls at
-				// all, so return from this function
-				// (subsequent "requests" to make a call by
-				// sending on the Calls channel will simply
-				// block and never go through)
+			select {
+			case <-heartbeatDoneCh:
 				return
+			default:
+				if err := c.Client.Call("RpcServerDriver.Heartbeat", struct{}{}, nil); err != nil {
+					log.Warnf("Error attempting heartbeat call to plugin server: %s", err)
+				}
+				time.Sleep(heartbeatInterval)
 			}
 		}
-	}()
+	}(c.heartbeatDoneCh)
 
 	var version int
 	if err := c.Client.Call("RpcServerDriver.GetVersion", struct{}{}, &version); err != nil {
@@ -108,6 +99,8 @@ func NewRpcClientDriver(rawDriverData []byte, driverName string) (*RpcClientDriv
 
 	mcnName = c.GetMachineName()
 	p.MachineName = mcnName
+	c.Client.MachineName = mcnName
+	c.plugin = p
 
 	return c, nil
 }
@@ -121,6 +114,15 @@ func (c *RpcClientDriver) UnmarshalJSON(data []byte) error {
 }
 
 func (c *RpcClientDriver) Close() error {
+	c.heartbeatDoneCh <- true
+	close(c.heartbeatDoneCh)
+
+	log.Debug("Making call to close connection to plugin binary")
+
+	if err := c.plugin.Close(); err != nil {
+		return err
+	}
+
 	log.Debug("Making call to close driver server")
 
 	if err := c.Client.Call("RpcServerDriver.Close", struct{}{}, nil); err != nil {

libmachine/drivers/rpc/server_driver.go

@@ -67,6 +67,15 @@ func (r RpcFlags) Bool(key string) bool {
 type RpcServerDriver struct {
 	ActualDriver drivers.Driver
 	CloseCh      chan bool
+	HeartbeatCh  chan bool
+}
+
+func NewRpcServerDriver(d drivers.Driver) *RpcServerDriver {
+	return &RpcServerDriver{
+		ActualDriver: d,
+		CloseCh:      make(chan bool),
+		HeartbeatCh:  make(chan bool),
+	}
 }
 
 func (r *RpcServerDriver) Close(_, _ *struct{}) error {
@@ -181,3 +190,8 @@ func (r *RpcServerDriver) Start(_ *struct{}, _ *struct{}) error {
 func (r *RpcServerDriver) Stop(_ *struct{}, _ *struct{}) error {
 	return r.ActualDriver.Stop()
 }
+
+func (r *RpcServerDriver) Heartbeat(_ *struct{}, _ *struct{}) error {
+	r.HeartbeatCh <- true
+	return nil
+}