interpolate environment for services in script checks

client/allocrunner/taskrunner/script_check_hook.go

@@ -174,7 +174,8 @@ func (h *scriptCheckHook) Stop(ctx context.Context, req *interfaces.TaskStopRequ
 
 func (h *scriptCheckHook) newScriptChecks() map[string]*scriptCheck {
 	scriptChecks := make(map[string]*scriptCheck)
-	for _, service := range h.task.Services {
+	interpolatedTaskServices := taskenv.InterpolateServices(h.taskEnv, h.task.Services)
+	for _, service := range interpolatedTaskServices {
 		for _, check := range service.Checks {
 			if check.Type != structs.ServiceCheckScript {
 				continue
@@ -204,7 +205,8 @@ func (h *scriptCheckHook) newScriptChecks() map[string]*scriptCheck {
 	// needs are entirely encapsulated within the group service hook which
 	// watches Consul for status changes.
 	tg := h.alloc.Job.LookupTaskGroup(h.alloc.TaskGroup)
-	for _, service := range tg.Services {
+	interpolatedGroupServices := taskenv.InterpolateServices(h.taskEnv, tg.Services)
+	for _, service := range interpolatedGroupServices {
 		for _, check := range service.Checks {
 			if check.Type != structs.ServiceCheckScript {
 				continue
@@ -293,37 +295,10 @@ func newScriptCheck(config *scriptCheckConfig) *scriptCheck {
 	sc.callback = newScriptCheckCallback(sc)
 	sc.logger = config.logger
 	sc.shutdownCh = config.shutdownCh
-
-	// the hash of the interior structs.ServiceCheck is used by the
-	// Consul client to get the ID to register for the check. So we
-	// update it here so that we have the same ID for UpdateTTL.
-
-	// TODO(tgross): this block is similar to one in service_hook
-	// and we can pull that out to a function so we know we're
-	// interpolating the same everywhere
-	sc.check.Name = config.taskEnv.ReplaceEnv(orig.Name)
-	sc.check.Type = config.taskEnv.ReplaceEnv(orig.Type)
 	sc.check.Command = sc.Command
 	sc.check.Args = sc.Args
-	sc.check.Path = config.taskEnv.ReplaceEnv(orig.Path)
-	sc.check.Protocol = config.taskEnv.ReplaceEnv(orig.Protocol)
-	sc.check.PortLabel = config.taskEnv.ReplaceEnv(orig.PortLabel)
-	sc.check.InitialStatus = config.taskEnv.ReplaceEnv(orig.InitialStatus)
-	sc.check.Method = config.taskEnv.ReplaceEnv(orig.Method)
-	sc.check.GRPCService = config.taskEnv.ReplaceEnv(orig.GRPCService)
-	if len(orig.Header) > 0 {
-		header := make(map[string][]string, len(orig.Header))
-		for k, vs := range orig.Header {
-			newVals := make([]string, len(vs))
-			for i, v := range vs {
-				newVals[i] = config.taskEnv.ReplaceEnv(v)
-			}
-			header[config.taskEnv.ReplaceEnv(k)] = newVals
-		}
-		sc.check.Header = header
-	}
+
 	if config.isGroup {
-		// TODO(tgross):
 		// group services don't have access to a task environment
 		// at creation, so their checks get registered before the
 		// check can be interpolated here. if we don't use the

client/allocrunner/taskrunner/script_check_hook_test.go

@@ -10,8 +10,11 @@ import (
 	"github.com/hashicorp/consul/api"
 	hclog "github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/nomad/client/allocrunner/taskrunner/interfaces"
+	"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/testlog"
+	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/stretchr/testify/require"
 )
@@ -217,3 +220,69 @@ func TestScript_Exec_Codes(t *testing.T) {
 		}
 	}
 }
+
+// TestScript_TaskEnvInterpolation asserts that script check hooks are
+// interpolated in the same way that services are
+func TestScript_TaskEnvInterpolation(t *testing.T) {
+
+	logger := testlog.HCLogger(t)
+	consulClient := consul.NewMockConsulServiceClient(t, logger)
+	exec, cancel := newBlockingScriptExec()
+	defer cancel()
+
+	alloc := mock.ConnectAlloc()
+	task := alloc.Job.TaskGroups[0].Tasks[0]
+
+	task.Services[0].Name = "${NOMAD_JOB_NAME}-${TASK}-${SVC_NAME}"
+	task.Services[0].Checks[0].Name = "${NOMAD_JOB_NAME}-${SVC_NAME}-check"
+	alloc.Job.Canonicalize() // need to re-canonicalize b/c the mock already did it
+
+	env := taskenv.NewBuilder(mock.Node(), alloc, task, "global").SetHookEnv(
+		"script_check",
+		map[string]string{"SVC_NAME": "frontend"}).Build()
+
+	svcHook := newServiceHook(serviceHookConfig{
+		alloc:  alloc,
+		task:   task,
+		consul: consulClient,
+		logger: logger,
+	})
+	// emulate prestart having been fired
+	svcHook.taskEnv = env
+
+	scHook := newScriptCheckHook(scriptCheckHookConfig{
+		alloc:        alloc,
+		task:         task,
+		consul:       consulClient,
+		logger:       logger,
+		shutdownWait: time.Hour, // heartbeater will never be called
+	})
+	// emulate prestart having been fired
+	scHook.taskEnv = env
+	scHook.driverExec = exec
+
+	expectedSvc := svcHook.getWorkloadServices().Services[0]
+	expected := agentconsul.MakeCheckID(agentconsul.MakeAllocServiceID(
+		alloc.ID, task.Name, expectedSvc), expectedSvc.Checks[0])
+
+	actual := scHook.newScriptChecks()
+	check, ok := actual[expected]
+	require.True(t, ok)
+	require.Equal(t, "my-job-frontend-check", check.check.Name)
+
+	// emulate an update
+	env = taskenv.NewBuilder(mock.Node(), alloc, task, "global").SetHookEnv(
+		"script_check",
+		map[string]string{"SVC_NAME": "backend"}).Build()
+	scHook.taskEnv = env
+	svcHook.taskEnv = env
+
+	expectedSvc = svcHook.getWorkloadServices().Services[0]
+	expected = agentconsul.MakeCheckID(agentconsul.MakeAllocServiceID(
+		alloc.ID, task.Name, expectedSvc), expectedSvc.Checks[0])
+
+	actual = scHook.newScriptChecks()
+	check, ok = actual[expected]
+	require.True(t, ok)
+	require.Equal(t, "my-job-backend-check", check.check.Name)
+}