From 7970b224b07f222e14bd4d88bfd4dbd4fd4821ca Mon Sep 17 00:00:00 2001
From: Michael Schurter <mschurter@hashicorp.com>
Date: Tue, 5 Mar 2019 15:12:02 -0800
Subject: [PATCH] client: emit event and call exited hooks during cleanup

Builds upon earlier commit that cleans up restored handles of terminal
allocs by also emitting terminated events and calling exited hooks when
appropriate.
---
 client/allocrunner/alloc_runner_unix_test.go  | 18 +++++++++++++++++
 .../allocrunner/interfaces/task_lifecycle.go  |  9 ++++++---
 client/allocrunner/taskrunner/task_runner.go  | 20 ++++++++++++++++---
 3 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/client/allocrunner/alloc_runner_unix_test.go b/client/allocrunner/alloc_runner_unix_test.go
index 397ba0fcf..ecb387d34 100644
--- a/client/allocrunner/alloc_runner_unix_test.go
+++ b/client/allocrunner/alloc_runner_unix_test.go
@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/hashicorp/nomad/client/consul"
 	"github.com/hashicorp/nomad/client/state"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
@@ -111,4 +112,21 @@ func TestAllocRunner_Restore_RunningTerminal(t *testing.T) {
 
 	// Assert logmon was cleaned up
 	require.Error(t, logmonProc.Signal(syscall.Signal(0)))
+
+	// Assert consul was cleaned up:
+	//   2 removals (canary+noncanary) during prekill
+	//   2 removals (canary+noncanary) during exited
+	consulOps := conf2.Consul.(*consul.MockConsulServiceClient).GetOps()
+	require.Len(t, consulOps, 4)
+	for _, op := range consulOps {
+		require.Equal(t, "remove", op.Op)
+	}
+
+	// Assert terminated task event was emitted
+	events := ar2.AllocState().TaskStates[task.Name].Events
+	require.Len(t, events, 4)
+	require.Equal(t, events[0].Type, structs.TaskReceived)
+	require.Equal(t, events[1].Type, structs.TaskSetup)
+	require.Equal(t, events[2].Type, structs.TaskStarted)
+	require.Equal(t, events[3].Type, structs.TaskTerminated)
 }
diff --git a/client/allocrunner/interfaces/task_lifecycle.go b/client/allocrunner/interfaces/task_lifecycle.go
index ff2faf88b..747c26040 100644
--- a/client/allocrunner/interfaces/task_lifecycle.go
+++ b/client/allocrunner/interfaces/task_lifecycle.go
@@ -87,7 +87,7 @@ type TaskPrestartHook interface {
 	TaskHook
 
 	// Prestart is called before the task is started including after every
-	// restart.
+	// restart. Prestart is not called if the allocation is terminal.
 	Prestart(context.Context, *TaskPrestartRequest, *TaskPrestartResponse) error
 }
 
@@ -114,7 +114,8 @@ type TaskPoststartResponse struct{}
 type TaskPoststartHook interface {
 	TaskHook
 
-	// Poststart is called after the task has started.
+	// Poststart is called after the task has started. Poststart is not
+	// called if the allocation is terminal.
 	Poststart(context.Context, *TaskPoststartRequest, *TaskPoststartResponse) error
 }
 
@@ -125,7 +126,8 @@ type TaskPreKillHook interface {
 	TaskHook
 
 	// PreKilling is called right before a task is going to be killed or
-	// restarted. They are called concurrently with TaskRunner.Run.
+	// restarted. They are called concurrently with TaskRunner.Run and may
+	// be called without Prestart being called.
 	PreKilling(context.Context, *TaskPreKillRequest, *TaskPreKillResponse) error
 }
 
@@ -136,6 +138,7 @@ type TaskExitedHook interface {
 	TaskHook
 
 	// Exited is called after a task exits and may or may not be restarted.
+	// Prestart may or may not have been called.
 	Exited(context.Context, *TaskExitedRequest, *TaskExitedResponse) error
 }
 
diff --git a/client/allocrunner/taskrunner/task_runner.go b/client/allocrunner/taskrunner/task_runner.go
index c4cd35143..e68e2a256 100644
--- a/client/allocrunner/taskrunner/task_runner.go
+++ b/client/allocrunner/taskrunner/task_runner.go
@@ -493,7 +493,15 @@ MAIN:
 	// that should be terminal, so if the handle still exists we should
 	// kill it here.
 	if tr.getDriverHandle() != nil {
-		tr.handleKill()
+		if result = tr.handleKill(); result != nil {
+			tr.emitExitResultEvent(result)
+		}
+
+		tr.clearDriverHandle()
+
+		if err := tr.exited(); err != nil {
+			tr.logger.Error("exited hooks failed while cleaning up terminal task", "error", err)
+		}
 	}
 
 	// Mark the task as dead
@@ -540,6 +548,14 @@ func (tr *TaskRunner) handleTaskExitResult(result *drivers.ExitResult) (retryWai
 		return true
 	}
 
+	// Emit Terminated event
+	tr.emitExitResultEvent(result)
+
+	return false
+}
+
+// emitExitResultEvent emits a TaskTerminated event for an ExitResult.
+func (tr *TaskRunner) emitExitResultEvent(result *drivers.ExitResult) {
 	event := structs.NewTaskEvent(structs.TaskTerminated).
 		SetExitCode(result.ExitCode).
 		SetSignal(result.Signal).
@@ -551,8 +567,6 @@ func (tr *TaskRunner) handleTaskExitResult(result *drivers.ExitResult) (retryWai
 	if result.OOMKilled && !tr.clientConfig.DisableTaggedMetrics {
 		metrics.IncrCounterWithLabels([]string{"client", "allocs", "oom_killed"}, 1, tr.baseLabels)
 	}
-
-	return false
 }
 
 // handleUpdates runs update hooks when triggerUpdateCh is ticked and exits