diff --git a/client/allocrunner/alloc_runner_unix_test.go b/client/allocrunner/alloc_runner_unix_test.go index ecb387d34e86..e0608efc75cd 100644 --- a/client/allocrunner/alloc_runner_unix_test.go +++ b/client/allocrunner/alloc_runner_unix_test.go @@ -116,8 +116,9 @@ func TestAllocRunner_Restore_RunningTerminal(t *testing.T) { // Assert consul was cleaned up: // 2 removals (canary+noncanary) during prekill // 2 removals (canary+noncanary) during exited + // 2 removals (canary+noncanary) during stop consulOps := conf2.Consul.(*consul.MockConsulServiceClient).GetOps() - require.Len(t, consulOps, 4) + require.Len(t, consulOps, 6) for _, op := range consulOps { require.Equal(t, "remove", op.Op) } diff --git a/client/allocrunner/taskrunner/service_hook.go b/client/allocrunner/taskrunner/service_hook.go index c5def6b29a3c..a89110417107 100644 --- a/client/allocrunner/taskrunner/service_hook.go +++ b/client/allocrunner/taskrunner/service_hook.go @@ -16,6 +16,11 @@ import ( "github.com/hashicorp/nomad/plugins/drivers" ) +var _ interfaces.TaskPoststartHook = &serviceHook{} +var _ interfaces.TaskPreKillHook = &serviceHook{} +var _ interfaces.TaskExitedHook = &serviceHook{} +var _ interfaces.TaskStopHook = &serviceHook{} + type serviceHookConfig struct { alloc *structs.Allocation task *structs.Task @@ -179,6 +184,13 @@ func (h *serviceHook) deregister() { } +func (h *serviceHook) Stop(ctx context.Context, req *interfaces.TaskStopRequest, resp *interfaces.TaskStopResponse) error { + h.mu.Lock() + defer h.mu.Unlock() + h.deregister() + return nil +} + func (h *serviceHook) getTaskServices() *agentconsul.TaskServices { // Interpolate with the task's environment interpolatedServices := interpolateServices(h.taskEnv, h.services) diff --git a/client/allocrunner/taskrunner/task_runner_test.go b/client/allocrunner/taskrunner/task_runner_test.go index 7d2bd5275e60..234aa70bd303 100644 --- a/client/allocrunner/taskrunner/task_runner_test.go +++ b/client/allocrunner/taskrunner/task_runner_test.go @@ -1991,7 +1991,7 @@ func TestTaskRunner_UnregisterConsul_Retries(t *testing.T) { consul := conf.Consul.(*consulapi.MockConsulServiceClient) consulOps := consul.GetOps() - require.Len(t, consulOps, 6) + require.Len(t, consulOps, 8) // Initial add require.Equal(t, "add", consulOps[0].Op) @@ -2006,6 +2006,10 @@ func TestTaskRunner_UnregisterConsul_Retries(t *testing.T) { // Removing canary and non-canary entries on retry require.Equal(t, "remove", consulOps[4].Op) require.Equal(t, "remove", consulOps[5].Op) + + // Removing canary and non-canary entries on stop + require.Equal(t, "remove", consulOps[1].Op) + require.Equal(t, "remove", consulOps[2].Op) } // testWaitForTaskToStart waits for the task to be running or fails the test