From c2c984ea508acd12c5e910c80ef2a16a14f8201c Mon Sep 17 00:00:00 2001
From: Nick Ethier <nethier@hashicorp.com>
Date: Fri, 15 Mar 2019 23:50:17 -0400
Subject: [PATCH] executor: block shutdown on process exiting

---
 drivers/shared/executor/executor.go       |  8 ++++++
 drivers/shared/executor/executor_linux.go | 15 +++++++++--
 drivers/shared/executor/executor_test.go  | 31 +++++++++++++++++++++++
 3 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/drivers/shared/executor/executor.go b/drivers/shared/executor/executor.go
index 156905eef..070e101d5 100644
--- a/drivers/shared/executor/executor.go
+++ b/drivers/shared/executor/executor.go
@@ -469,6 +469,14 @@ func (e *UniversalExecutor) Shutdown(signal string, grace time.Duration) error {
 		proc.Kill()
 	}
 
+	// Wait for process to exit
+	select {
+	case <-e.processExited:
+	case <-time.After(time.Second * 15):
+		e.logger.Warn("process did not exit after 15 seconds")
+		merr.Errors = append(merr.Errors, fmt.Errorf("process did not exit after 15 seconds"))
+	}
+
 	// Prefer killing the process via the resource container.
 	if !(e.commandCfg.ResourceLimits || e.commandCfg.BasicProcessCgroup) {
 		if err := e.cleanupChildProcesses(proc); err != nil && err.Error() != finishedErr {
diff --git a/drivers/shared/executor/executor_linux.go b/drivers/shared/executor/executor_linux.go
index 565561d6e..3ac4618aa 100644
--- a/drivers/shared/executor/executor_linux.go
+++ b/drivers/shared/executor/executor_linux.go
@@ -339,10 +339,21 @@ func (l *LibcontainerExecutor) Shutdown(signal string, grace time.Duration) erro
 		case <-time.After(grace):
 			// Force kill all container processes after grace period,
 			// hence `true` argument.
-			return l.container.Signal(os.Kill, true)
+			if err := l.container.Signal(os.Kill, true); err != nil {
+				return err
+			}
 		}
 	} else {
-		return l.container.Signal(os.Kill, true)
+		if err := l.container.Signal(os.Kill, true); err != nil {
+			return err
+		}
+	}
+
+	select {
+	case <-l.userProcExited:
+		return nil
+	case <-time.After(time.Second * 15):
+		return fmt.Errorf("process failed to exit after 15 seconds")
 	}
 }
 
diff --git a/drivers/shared/executor/executor_test.go b/drivers/shared/executor/executor_test.go
index aad0647bd..1f8000bd8 100644
--- a/drivers/shared/executor/executor_test.go
+++ b/drivers/shared/executor/executor_test.go
@@ -20,6 +20,7 @@ import (
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/plugins/drivers"
 	tu "github.com/hashicorp/nomad/testutil"
+	ps "github.com/mitchellh/go-ps"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -234,6 +235,36 @@ func TestExecutor_Start_Kill(pt *testing.T) {
 	}
 }
 
+func TestExecutor_Shutdown_Exit(t *testing.T) {
+	require := require.New(t)
+	t.Parallel()
+	execCmd, allocDir := testExecutorCommand(t)
+	execCmd.Cmd = "/bin/sleep"
+	execCmd.Args = []string{"100"}
+	cfg := &ExecutorConfig{
+		LogFile: "/dev/null",
+	}
+	executor, pluginClient, err := CreateExecutor(testlog.HCLogger(t), nil, cfg)
+	require.NoError(err)
+
+	proc, err := executor.Launch(execCmd)
+	require.NoError(err)
+	require.NotZero(proc.Pid)
+
+	executor.Shutdown("", 0)
+	pluginClient.Kill()
+	tu.WaitForResult(func() (bool, error) {
+		p, err := ps.FindProcess(proc.Pid)
+		if err != nil {
+			return false, err
+		}
+		return p == nil, fmt.Errorf("process found: %d", proc.Pid)
+	}, func(err error) {
+		require.NoError(err)
+	})
+	require.NoError(allocDir.Destroy())
+}
+
 func TestUniversalExecutor_MakeExecutable(t *testing.T) {
 	t.Parallel()
 	// Create a temp file