From 39bfe4b2eee8b5006592c949c981a3b3d6f822e7 Mon Sep 17 00:00:00 2001
From: Drew Bailey <2614075+drewbailey@users.noreply.github.com>
Date: Fri, 3 Apr 2020 15:05:46 -0400
Subject: [PATCH] Fixes bug that prevented group shutdown_delay updates

Group shutdown delay updates were not properly handled in Update hook.
This commit also ensures that plan output is displayed.
---
 client/allocrunner/alloc_runner_test.go      |  2 +-
 client/allocrunner/groupservice_hook.go      |  8 ++----
 client/allocrunner/groupservice_hook_test.go | 28 ++++++++++++++++++++
 command/agent/job_endpoint.go                |  5 +---
 nomad/structs/diff_test.go                   | 20 ++++++++++++++
 nomad/structs/structs.go                     |  8 ++----
 6 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/client/allocrunner/alloc_runner_test.go b/client/allocrunner/alloc_runner_test.go
index 3379ba3d3004..c68bf15b0e1b 100644
--- a/client/allocrunner/alloc_runner_test.go
+++ b/client/allocrunner/alloc_runner_test.go
@@ -180,7 +180,7 @@ func TestAllocRunner_TaskGroup_ShutdownDelay(t *testing.T) {
 
 	// Set a shutdown delay
 	shutdownDelay := 1 * time.Second
-	alloc.Job.TaskGroups[0].ShutdownDelay = &shutdownDelay
+	alloc.Job.TaskGroups[0].ShutdownDelay = shutdownDelay
 
 	conf, cleanup := testAllocRunnerConfig(t, alloc)
 	defer cleanup()
diff --git a/client/allocrunner/groupservice_hook.go b/client/allocrunner/groupservice_hook.go
index 0325d1bae214..cd573517a190 100644
--- a/client/allocrunner/groupservice_hook.go
+++ b/client/allocrunner/groupservice_hook.go
@@ -46,20 +46,15 @@ type groupServiceHookConfig struct {
 }
 
 func newGroupServiceHook(cfg groupServiceHookConfig) *groupServiceHook {
-	var shutdownDelay time.Duration
 	tg := cfg.alloc.Job.LookupTaskGroup(cfg.alloc.TaskGroup)
 
-	if tg.ShutdownDelay != nil {
-		shutdownDelay = *tg.ShutdownDelay
-	}
-
 	h := &groupServiceHook{
 		allocID:        cfg.alloc.ID,
 		group:          cfg.alloc.TaskGroup,
 		restarter:      cfg.restarter,
 		consulClient:   cfg.consul,
 		taskEnvBuilder: cfg.taskEnvBuilder,
-		delay:          shutdownDelay,
+		delay:          tg.ShutdownDelay,
 	}
 	h.logger = cfg.logger.Named(h.Name())
 	h.services = cfg.alloc.Job.LookupTaskGroup(h.group).Services
@@ -116,6 +111,7 @@ func (h *groupServiceHook) Update(req *interfaces.RunnerUpdateRequest) error {
 	h.services = req.Alloc.Job.LookupTaskGroup(h.group).Services
 	h.canary = canary
 	h.taskEnvBuilder.UpdateTask(req.Alloc, nil)
+	h.delay = req.Alloc.Job.LookupTaskGroup(h.group).ShutdownDelay
 
 	// Create new task services struct with those new values
 	newWorkloadServices := h.getWorkloadServices()
diff --git a/client/allocrunner/groupservice_hook_test.go b/client/allocrunner/groupservice_hook_test.go
index 82ca09d7b686..164fbc3bb616 100644
--- a/client/allocrunner/groupservice_hook_test.go
+++ b/client/allocrunner/groupservice_hook_test.go
@@ -56,6 +56,34 @@ func TestGroupServiceHook_NoGroupServices(t *testing.T) {
 	require.Equal(t, "remove", ops[2].Op)
 }
 
+// TestGroupServiceHook_ShutdownDelayUpdate asserts calling group service hooks
+// update updates the hooks delay value.
+func TestGroupServiceHook_ShutdownDelayUpdate(t *testing.T) {
+	t.Parallel()
+
+	alloc := mock.Alloc()
+	alloc.Job.TaskGroups[0].ShutdownDelay = 10 * time.Second
+
+	logger := testlog.HCLogger(t)
+	consulClient := consul.NewMockConsulServiceClient(t, logger)
+
+	h := newGroupServiceHook(groupServiceHookConfig{
+		alloc:          alloc,
+		consul:         consulClient,
+		restarter:      agentconsul.NoopRestarter(),
+		taskEnvBuilder: taskenv.NewBuilder(mock.Node(), alloc, nil, alloc.Job.Region),
+		logger:         logger,
+	})
+	require.NoError(t, h.Prerun())
+
+	alloc.Job.TaskGroups[0].ShutdownDelay = 15 * time.Second
+	req := &interfaces.RunnerUpdateRequest{Alloc: alloc}
+	require.NoError(t, h.Update(req))
+
+	// Assert that update updated the delay value
+	require.Equal(t, h.delay, 15*time.Second)
+}
+
 // TestGroupServiceHook_GroupServices asserts group service hooks with group
 // services does not error.
 func TestGroupServiceHook_GroupServices(t *testing.T) {
diff --git a/command/agent/job_endpoint.go b/command/agent/job_endpoint.go
index 09bac74cc2ef..78c7f8c68805 100644
--- a/command/agent/job_endpoint.go
+++ b/command/agent/job_endpoint.go
@@ -777,6 +777,7 @@ func ApiTgToStructsTG(job *structs.Job, taskGroup *api.TaskGroup, tg *structs.Ta
 	tg.Name = *taskGroup.Name
 	tg.Count = *taskGroup.Count
 	tg.Meta = taskGroup.Meta
+	tg.ShutdownDelay = *taskGroup.ShutdownDelay
 	tg.Constraints = ApiConstraintsToStructs(taskGroup.Constraints)
 	tg.Affinities = ApiAffinitiesToStructs(taskGroup.Affinities)
 	tg.Networks = ApiNetworkResourceToStructs(taskGroup.Networks)
@@ -789,10 +790,6 @@ func ApiTgToStructsTG(job *structs.Job, taskGroup *api.TaskGroup, tg *structs.Ta
 		Mode:     *taskGroup.RestartPolicy.Mode,
 	}
 
-	if taskGroup.ShutdownDelay != nil {
-		tg.ShutdownDelay = taskGroup.ShutdownDelay
-	}
-
 	if taskGroup.ReschedulePolicy != nil {
 		tg.ReschedulePolicy = &structs.ReschedulePolicy{
 			Attempts:      *taskGroup.ReschedulePolicy.Attempts,
diff --git a/nomad/structs/diff_test.go b/nomad/structs/diff_test.go
index c51f38d0f80c..f653d1b30aed 100644
--- a/nomad/structs/diff_test.go
+++ b/nomad/structs/diff_test.go
@@ -3029,6 +3029,26 @@ func TestTaskGroupDiff(t *testing.T) {
 				},
 			},
 		},
+		{
+			TestCase: "TaskGroup shutdown_delay edited",
+			Old: &TaskGroup{
+				ShutdownDelay: 30 * time.Second,
+			},
+			New: &TaskGroup{
+				ShutdownDelay: 5 * time.Second,
+			},
+			Expected: &TaskGroupDiff{
+				Type: DiffTypeEdited,
+				Fields: []*FieldDiff{
+					{
+						Type: DiffTypeEdited,
+						Name: "ShutdownDelay",
+						Old:  "30000000000",
+						New:  "5000000000",
+					},
+				},
+			},
+		},
 	}
 
 	for i, c := range cases {
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index 453431d615da..1ea458a86f04 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -3790,7 +3790,7 @@ func (j *Job) Validate() error {
 			taskGroups[tg.Name] = idx
 		}
 
-		if tg.ShutdownDelay != nil && *tg.ShutdownDelay < 0 {
+		if tg.ShutdownDelay < 0 {
 			mErr.Errors = append(mErr.Errors, errors.New("ShutdownDelay must be a positive value"))
 		}
 
@@ -5202,7 +5202,7 @@ type TaskGroup struct {
 
 	// ShutdownDelay is the amount of time to wait between deregistering
 	// group services in consul and stopping tasks.
-	ShutdownDelay *time.Duration
+	ShutdownDelay time.Duration
 }
 
 func (tg *TaskGroup) Copy() *TaskGroup {
@@ -5250,10 +5250,6 @@ func (tg *TaskGroup) Copy() *TaskGroup {
 		}
 	}
 
-	if tg.ShutdownDelay != nil {
-		ntg.ShutdownDelay = tg.ShutdownDelay
-	}
-
 	return ntg
 }