Merge pull request #5821 from hashicorp/dani/b-5770

trhooks: Add TaskStopHook interface to services

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)
 	}

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)

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[6].Op)
+	require.Equal(t, "remove", consulOps[7].Op)
 }
 
 // testWaitForTaskToStart waits for the task to be running or fails the test