diff --git a/client/allocrunner/taskrunner/script_check_hook.go b/client/allocrunner/taskrunner/script_check_hook.go index b40e92301aaa..f10486ac526c 100644 --- a/client/allocrunner/taskrunner/script_check_hook.go +++ b/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 diff --git a/client/allocrunner/taskrunner/script_check_hook_test.go b/client/allocrunner/taskrunner/script_check_hook_test.go index 26219f943730..9fba3a67f667 100644 --- a/client/allocrunner/taskrunner/script_check_hook_test.go +++ b/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) +}