From ec18a80cc803cd22f28c6d3b22506ab888514ce5 Mon Sep 17 00:00:00 2001
From: Danielle Lancashire <dani@hashicorp.com>
Date: Wed, 12 Jun 2019 14:59:26 +0200
Subject: [PATCH] trhooks: Add TaskStopHook interface to services

We currently only run cleanup Service Hooks when a task is either
Killed, or Exited. However, due to the implementation of a task runner,
tasks are only Exited if they every correctly started running, which is
not true when you recieve an error early in the task start flow, such as
not being able to pull secrets from Vault.

This updates the service hook to also call consul deregistration
routines during a task `Stop` lifecycle event, to ensure that any
registered checks and services are cleared in such cases.
---
 client/allocrunner/alloc_runner_unix_test.go     |  3 ++-
 client/allocrunner/taskrunner/service_hook.go    | 12 ++++++++++++
 .../allocrunner/taskrunner/task_runner_test.go   |  6 +++++-
 client/storage/interface.go                      |  8 ++++++++
 client/storage/plugin_manager.go                 | 16 ++++++++++++++++
 5 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 client/storage/interface.go
 create mode 100644 client/storage/plugin_manager.go

diff --git a/client/allocrunner/alloc_runner_unix_test.go b/client/allocrunner/alloc_runner_unix_test.go
index ecb387d34e86..e0608efc75cd 100644
--- a/client/allocrunner/alloc_runner_unix_test.go
+++ b/client/allocrunner/alloc_runner_unix_test.go
@@ -116,8 +116,9 @@ func TestAllocRunner_Restore_RunningTerminal(t *testing.T) {
 	// Assert consul was cleaned up:
 	//   2 removals (canary+noncanary) during prekill
 	//   2 removals (canary+noncanary) during exited
+	//   2 removals (canary+noncanary) during stop
 	consulOps := conf2.Consul.(*consul.MockConsulServiceClient).GetOps()
-	require.Len(t, consulOps, 4)
+	require.Len(t, consulOps, 6)
 	for _, op := range consulOps {
 		require.Equal(t, "remove", op.Op)
 	}
diff --git a/client/allocrunner/taskrunner/service_hook.go b/client/allocrunner/taskrunner/service_hook.go
index c5def6b29a3c..a89110417107 100644
--- a/client/allocrunner/taskrunner/service_hook.go
+++ b/client/allocrunner/taskrunner/service_hook.go
@@ -16,6 +16,11 @@ import (
 	"github.com/hashicorp/nomad/plugins/drivers"
 )
 
+var _ interfaces.TaskPoststartHook = &serviceHook{}
+var _ interfaces.TaskPreKillHook = &serviceHook{}
+var _ interfaces.TaskExitedHook = &serviceHook{}
+var _ interfaces.TaskStopHook = &serviceHook{}
+
 type serviceHookConfig struct {
 	alloc  *structs.Allocation
 	task   *structs.Task
@@ -179,6 +184,13 @@ func (h *serviceHook) deregister() {
 
 }
 
+func (h *serviceHook) Stop(ctx context.Context, req *interfaces.TaskStopRequest, resp *interfaces.TaskStopResponse) error {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	h.deregister()
+	return nil
+}
+
 func (h *serviceHook) getTaskServices() *agentconsul.TaskServices {
 	// Interpolate with the task's environment
 	interpolatedServices := interpolateServices(h.taskEnv, h.services)
diff --git a/client/allocrunner/taskrunner/task_runner_test.go b/client/allocrunner/taskrunner/task_runner_test.go
index 7d2bd5275e60..234aa70bd303 100644
--- a/client/allocrunner/taskrunner/task_runner_test.go
+++ b/client/allocrunner/taskrunner/task_runner_test.go
@@ -1991,7 +1991,7 @@ func TestTaskRunner_UnregisterConsul_Retries(t *testing.T) {
 
 	consul := conf.Consul.(*consulapi.MockConsulServiceClient)
 	consulOps := consul.GetOps()
-	require.Len(t, consulOps, 6)
+	require.Len(t, consulOps, 8)
 
 	// Initial add
 	require.Equal(t, "add", consulOps[0].Op)
@@ -2006,6 +2006,10 @@ func TestTaskRunner_UnregisterConsul_Retries(t *testing.T) {
 	// Removing canary and non-canary entries on retry
 	require.Equal(t, "remove", consulOps[4].Op)
 	require.Equal(t, "remove", consulOps[5].Op)
+
+	// Removing canary and non-canary entries on stop
+	require.Equal(t, "remove", consulOps[1].Op)
+	require.Equal(t, "remove", consulOps[2].Op)
 }
 
 // testWaitForTaskToStart waits for the task to be running or fails the test
diff --git a/client/storage/interface.go b/client/storage/interface.go
new file mode 100644
index 000000000000..4a701ed4e55f
--- /dev/null
+++ b/client/storage/interface.go
@@ -0,0 +1,8 @@
+package storage
+
+type Volume struct {
+}
+
+type VolumeDriver interface {
+	RequiresHostAttachment(v *Volume) bool
+}
diff --git a/client/storage/plugin_manager.go b/client/storage/plugin_manager.go
new file mode 100644
index 000000000000..def550e7329f
--- /dev/null
+++ b/client/storage/plugin_manager.go
@@ -0,0 +1,16 @@
+package storage
+
+import hclog "github.com/hashicorp/go-hclog"
+
+type ManagerConfig struct {
+	SocketDir string
+}
+
+type PluginManager struct {
+	cfg    *ManagerConfig
+	logger hclog.Logger
+}
+
+func (p *PluginManager) Run() error {
+	return nil
+}