From b81a0018b4c52cc4e04a92f98eed71170288a2ba Mon Sep 17 00:00:00 2001
From: Drew Bailey <2614075+drewbailey@users.noreply.github.com>
Date: Sun, 5 Apr 2020 17:30:25 -0400
Subject: [PATCH] Group shutdown delay fixes

Group shutdown delay updates were not properly handled in Update hook.
This commit also ensures that plan output is displayed.
---
 client/allocrunner/groupservice_hook.go      | 10 ++++++-
 client/allocrunner/groupservice_hook_test.go | 29 ++++++++++++++++++++
 nomad/structs/diff.go                        | 12 ++++++++
 nomad/structs/diff_test.go                   | 21 ++++++++++++++
 4 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/client/allocrunner/groupservice_hook.go b/client/allocrunner/groupservice_hook.go
index 0325d1bae214..3c915e90dcb0 100644
--- a/client/allocrunner/groupservice_hook.go
+++ b/client/allocrunner/groupservice_hook.go
@@ -112,11 +112,19 @@ func (h *groupServiceHook) Update(req *interfaces.RunnerUpdateRequest) error {
 	}
 
 	// Update group service hook fields
+	tg := req.Alloc.Job.LookupTaskGroup(h.group)
 	h.networks = networks
-	h.services = req.Alloc.Job.LookupTaskGroup(h.group).Services
+	h.services = tg.Services
 	h.canary = canary
 	h.taskEnvBuilder.UpdateTask(req.Alloc, nil)
 
+	// Update shutdown delay
+	var shutdown time.Duration
+	if tg.ShutdownDelay != nil {
+		shutdown = *tg.ShutdownDelay
+	}
+	h.delay = shutdown
+
 	// 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..8920c5284a8a 100644
--- a/client/allocrunner/groupservice_hook_test.go
+++ b/client/allocrunner/groupservice_hook_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/hashicorp/nomad/client/consul"
 	"github.com/hashicorp/nomad/client/taskenv"
 	agentconsul "github.com/hashicorp/nomad/command/agent/consul"
+	"github.com/hashicorp/nomad/helper"
 	"github.com/hashicorp/nomad/helper/testlog"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
@@ -56,6 +57,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 = helper.TimeToPtr(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 = helper.TimeToPtr(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/nomad/structs/diff.go b/nomad/structs/diff.go
index 5e6eaf306544..251ceee5d3b0 100644
--- a/nomad/structs/diff.go
+++ b/nomad/structs/diff.go
@@ -224,6 +224,18 @@ func (tg *TaskGroup) Diff(other *TaskGroup, contextual bool) (*TaskGroupDiff, er
 		newPrimitiveFlat = flatmap.Flatten(other, filter, true)
 	}
 
+	// ShutdownDelay diff
+	if tg.ShutdownDelay == nil {
+		oldPrimitiveFlat["ShutdownDelay"] = "nil"
+	} else {
+		oldPrimitiveFlat["ShutdownDelay"] = fmt.Sprintf("%d", *tg.ShutdownDelay)
+	}
+	if other.ShutdownDelay == nil {
+		newPrimitiveFlat["ShutdownDelay"] = "nil"
+	} else {
+		newPrimitiveFlat["ShutdownDelay"] = fmt.Sprintf("%d", *other.ShutdownDelay)
+	}
+
 	// Diff the primitive fields.
 	diff.Fields = fieldDiffs(oldPrimitiveFlat, newPrimitiveFlat, false)
 
diff --git a/nomad/structs/diff_test.go b/nomad/structs/diff_test.go
index c51f38d0f80c..06ac79ad1e88 100644
--- a/nomad/structs/diff_test.go
+++ b/nomad/structs/diff_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/hashicorp/nomad/helper"
 	"github.com/stretchr/testify/require"
 )
 
@@ -3029,6 +3030,26 @@ func TestTaskGroupDiff(t *testing.T) {
 				},
 			},
 		},
+		{
+			TestCase: "TaskGroup shutdown_delay edited",
+			Old: &TaskGroup{
+				ShutdownDelay: helper.TimeToPtr(30 * time.Second),
+			},
+			New: &TaskGroup{
+				ShutdownDelay: helper.TimeToPtr(5 * time.Second),
+			},
+			Expected: &TaskGroupDiff{
+				Type: DiffTypeEdited,
+				Fields: []*FieldDiff{
+					{
+						Type: DiffTypeEdited,
+						Name: "ShutdownDelay",
+						Old:  "30000000000",
+						New:  "5000000000",
+					},
+				},
+			},
+		},
 	}
 
 	for i, c := range cases {