From 3af2d05f6be8f75e69dad3cb794463be77d82188 Mon Sep 17 00:00:00 2001
From: Drew Bailey <2614075+drewbailey@users.noreply.github.com>
Date: Wed, 8 Apr 2020 10:24:19 -0400
Subject: [PATCH] Run task shutdown_delay regardless of service registration

task shutdown_delay will currently only run if there are registered
services for the task. This implementation detail isn't explicity stated
anywhere and is defined outside of the service stanza.

This change moves shutdown_delay to be evaluated after prekill hooks are
run, outside of any task runner hooks.

just use time.sleep
---
 client/allocrunner/taskrunner/service_hook.go | 14 --------------
 client/allocrunner/taskrunner/task_runner.go  |  8 ++++++++
 2 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/client/allocrunner/taskrunner/service_hook.go b/client/allocrunner/taskrunner/service_hook.go
index eaf5cdb89..e329af93f 100644
--- a/client/allocrunner/taskrunner/service_hook.go
+++ b/client/allocrunner/taskrunner/service_hook.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"sync"
-	"time"
 
 	log "github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
@@ -40,7 +39,6 @@ type serviceHook struct {
 	logger    log.Logger
 
 	// The following fields may be updated
-	delay      time.Duration
 	driverExec tinterfaces.ScriptExecutor
 	driverNet  *drivers.DriverNetwork
 	canary     bool
@@ -64,7 +62,6 @@ func newServiceHook(c serviceHookConfig) *serviceHook {
 		taskName:  c.task.Name,
 		services:  c.task.Services,
 		restarter: c.restarter,
-		delay:     c.task.ShutdownDelay,
 	}
 
 	if res := c.alloc.AllocatedResources.Tasks[c.task.Name]; res != nil {
@@ -140,7 +137,6 @@ func (h *serviceHook) updateHookFields(req *interfaces.TaskUpdateRequest) error
 	}
 
 	// Update service hook fields
-	h.delay = task.ShutdownDelay
 	h.taskEnv = req.TaskEnv
 	h.services = task.Services
 	h.networks = networks
@@ -156,16 +152,6 @@ func (h *serviceHook) PreKilling(ctx context.Context, req *interfaces.TaskPreKil
 	// Deregister before killing task
 	h.deregister()
 
-	// If there's no shutdown delay, exit early
-	if h.delay == 0 {
-		return nil
-	}
-
-	h.logger.Debug("waiting before killing task", "shutdown_delay", h.delay)
-	select {
-	case <-ctx.Done():
-	case <-time.After(h.delay):
-	}
 	return nil
 }
 
diff --git a/client/allocrunner/taskrunner/task_runner.go b/client/allocrunner/taskrunner/task_runner.go
index 43712f140..592d70bed 100644
--- a/client/allocrunner/taskrunner/task_runner.go
+++ b/client/allocrunner/taskrunner/task_runner.go
@@ -844,6 +844,14 @@ func (tr *TaskRunner) handleKill() *drivers.ExitResult {
 	// Run the pre killing hooks
 	tr.preKill()
 
+	// Wait for task ShutdownDelay after running prekill hooks
+	// This allows for things like service de-registration to run
+	// before waiting to kill task
+	if delay := tr.Task().ShutdownDelay; delay != 0 {
+		tr.logger.Debug("waiting before killing task", "shutdown_delay", delay)
+		time.Sleep(delay)
+	}
+
 	// Tell the restart tracker that the task has been killed so it doesn't
 	// attempt to restart it.
 	tr.restartTracker.SetKilled()