provide -no-shutdown-delay flag for job/alloc stop (#11596)

Some operators use very long group/task `shutdown_delay` settings to
safely drain network connections to their workloads after service
deregistration. But during incident response, they may want to cause
that drain to be skipped so they can quickly shed load.

Provide a `-no-shutdown-delay` flag on the `nomad alloc stop` and
`nomad job stop` commands that bypasses the delay. This sets a new
desired transition state on the affected allocations that the
allocation/task runner will identify during pre-kill on the client.

Note (as documented here) that using this flag will almost always
result in failed inbound network connections for workloads as the
tasks will exit before clients receive updated service discovery
information and won't be gracefully drained.

.changelog/11596.txt

@@ -0,0 +1,3 @@
+```release-note:improvement
+cli: provide `-no-shutdown-delay` option to `job stop` and `alloc stop` commands to ignore `shutdown_delay`
+```

api/jobs.go

@@ -299,6 +299,11 @@ type DeregisterOptions struct {
 	// is useful when an operator wishes to push through a job deregistration
 	// in busy clusters with a large evaluation backlog.
 	EvalPriority int
+
+	// NoShutdownDelay, if set to true, will override the group and
+	// task shutdown_delay configuration and ignore the delay for any
+	// allocations stopped as a result of this Deregister call.
+	NoShutdownDelay bool
 }
 
 // DeregisterOpts is used to remove an existing job. See DeregisterOptions
@@ -312,8 +317,8 @@ func (j *Jobs) DeregisterOpts(jobID string, opts *DeregisterOptions, q *WriteOpt
 	// Protect against nil opts. url.Values expects a string, and so using
 	// fmt.Sprintf is the best way to do this.
 	if opts != nil {
-		endpoint += fmt.Sprintf("?purge=%t&global=%t&eval_priority=%v",
-			opts.Purge, opts.Global, opts.EvalPriority)
+		endpoint += fmt.Sprintf("?purge=%t&global=%t&eval_priority=%v&no_shutdown_delay=%t",
+			opts.Purge, opts.Global, opts.EvalPriority, opts.NoShutdownDelay)
 	}
 
 	wm, err := j.client.delete(endpoint, &resp, q)

client/allocrunner/alloc_runner.go

@@ -171,6 +171,9 @@ type allocRunner struct {
 
 	taskHookCoordinator *taskHookCoordinator
 
+	shutdownDelayCtx      context.Context
+	shutdownDelayCancelFn context.CancelFunc
+
 	// rpcClient is the RPC Client that should be used by the allocrunner and its
 	// hooks to communicate with Nomad Servers.
 	rpcClient RPCer
@@ -230,6 +233,10 @@ func NewAllocRunner(config *Config) (*allocRunner, error) {
 
 	ar.taskHookCoordinator = newTaskHookCoordinator(ar.logger, tg.Tasks)
 
+	shutdownDelayCtx, shutdownDelayCancel := context.WithCancel(context.Background())
+	ar.shutdownDelayCtx = shutdownDelayCtx
+	ar.shutdownDelayCancelFn = shutdownDelayCancel
+
 	// Initialize the runners hooks.
 	if err := ar.initRunnerHooks(config.ClientConfig); err != nil {
 		return nil, err
@@ -265,6 +272,7 @@ func (ar *allocRunner) initTaskRunners(tasks []*structs.Task) error {
 			DriverManager:        ar.driverManager,
 			ServersContactedCh:   ar.serversContactedCh,
 			StartConditionMetCtx: ar.taskHookCoordinator.startConditionForTask(task),
+			ShutdownDelayCtx:     ar.shutdownDelayCtx,
 		}
 
 		if ar.cpusetManager != nil {
@@ -824,6 +832,10 @@ func (ar *allocRunner) Update(update *structs.Allocation) {
 	default:
 	}
 
+	if update.DesiredTransition.ShouldIgnoreShutdownDelay() {
+		ar.shutdownDelayCancelFn()
+	}
+
 	// Queue the new update
 	ar.allocUpdatedCh <- update
 }

client/allocrunner/alloc_runner_hooks.go

@@ -159,6 +159,7 @@ func (ar *allocRunner) initRunnerHooks(config *clientconfig.Config) error {
 			taskEnvBuilder:      envBuilder,
 			networkStatusGetter: ar,
 			logger:              hookLogger,
+			shutdownDelayCtx:    ar.shutdownDelayCtx,
 		}),
 		newConsulGRPCSocketHook(hookLogger, alloc, ar.allocDir, config.ConsulConfig),
 		newConsulHTTPSocketHook(hookLogger, alloc, ar.allocDir, config.ConsulConfig),

client/allocrunner/groupservice_hook.go

@@ -1,6 +1,7 @@
 package allocrunner
 
 import (
+	"context"
 	"sync"
 	"time"
 
@@ -29,9 +30,9 @@ type groupServiceHook struct {
 	consulClient        consul.ConsulServiceAPI
 	consulNamespace     string
 	prerun              bool
-	delay               time.Duration
 	deregistered        bool
 	networkStatusGetter networkStatusGetter
+	shutdownDelayCtx    context.Context
 
 	logger log.Logger
 
@@ -41,6 +42,7 @@ type groupServiceHook struct {
 	networks       structs.Networks
 	ports          structs.AllocatedPorts
 	taskEnvBuilder *taskenv.Builder
+	delay          time.Duration
 
 	// Since Update() may be called concurrently with any other hook all
 	// hook methods must be fully serialized
@@ -54,6 +56,7 @@ type groupServiceHookConfig struct {
 	restarter           agentconsul.WorkloadRestarter
 	taskEnvBuilder      *taskenv.Builder
 	networkStatusGetter networkStatusGetter
+	shutdownDelayCtx    context.Context
 	logger              log.Logger
 }
 
@@ -76,6 +79,7 @@ func newGroupServiceHook(cfg groupServiceHookConfig) *groupServiceHook {
 		networkStatusGetter: cfg.networkStatusGetter,
 		logger:              cfg.logger.Named(groupServiceHookName),
 		services:            cfg.alloc.Job.LookupTaskGroup(cfg.alloc.TaskGroup).Services,
+		shutdownDelayCtx:    cfg.shutdownDelayCtx,
 	}
 
 	if cfg.alloc.AllocatedResources != nil {
@@ -187,9 +191,12 @@ func (h *groupServiceHook) preKillLocked() {
 
 	h.logger.Debug("delay before killing tasks", "group", h.group, "shutdown_delay", h.delay)
 
-	// Wait for specified shutdown_delay
+	select {
+	// Wait for specified shutdown_delay unless ignored
 	// This will block an agent from shutting down.
-	<-time.After(h.delay)
+	case <-time.After(h.delay):
+	case <-h.shutdownDelayCtx.Done():
+	}
 }
 
 func (h *groupServiceHook) Postrun() error {

client/allocrunner/taskrunner/task_runner.go

@@ -112,6 +112,11 @@ type TaskRunner struct {
 	killErr     error
 	killErrLock sync.Mutex
 
+	// shutdownDelayCtx is a context from the alloc runner which will
+	// tell us to exit early from shutdown_delay
+	shutdownDelayCtx      context.Context
+	shutdownDelayCancelFn context.CancelFunc
+
 	// Logger is the logger for the task runner.
 	logger log.Logger
 
@@ -287,6 +292,13 @@ type Config struct {
 
 	// startConditionMetCtx is done when TR should start the task
 	StartConditionMetCtx <-chan struct{}
+
+	// ShutdownDelayCtx is a context from the alloc runner which will
+	// tell us to exit early from shutdown_delay
+	ShutdownDelayCtx context.Context
+
+	// ShutdownDelayCancelFn should only be used in testing.
+	ShutdownDelayCancelFn context.CancelFunc
 }
 
 func NewTaskRunner(config *Config) (*TaskRunner, error) {
@@ -342,6 +354,8 @@ func NewTaskRunner(config *Config) (*TaskRunner, error) {
 		maxEvents:              defaultMaxEvents,
 		serversContactedCh:     config.ServersContactedCh,
 		startConditionMetCtx:   config.StartConditionMetCtx,
+		shutdownDelayCtx:       config.ShutdownDelayCtx,
+		shutdownDelayCancelFn:  config.ShutdownDelayCancelFn,
 	}
 
 	// Create the logger based on the allocation ID
@@ -895,6 +909,8 @@ func (tr *TaskRunner) handleKill(resultCh <-chan *drivers.ExitResult) *drivers.E
 		select {
 		case result := <-resultCh:
 			return result
+		case <-tr.shutdownDelayCtx.Done():
+			break
 		case <-time.After(delay):
 		}
 	}
@@ -1478,3 +1494,9 @@ func (tr *TaskRunner) DriverCapabilities() (*drivers.Capabilities, error) {
 func (tr *TaskRunner) SetAllocHookResources(res *cstructs.AllocHookResources) {
 	tr.allocHookResources = res
 }
+
+// shutdownDelayCancel is used for testing only and cancels the
+// shutdownDelayCtx
+func (tr *TaskRunner) shutdownDelayCancel() {
+	tr.shutdownDelayCancelFn()
+}

client/allocrunner/taskrunner/task_runner_test.go

@@ -14,6 +14,10 @@ import (
 	"time"
 
 	"github.com/golang/snappy"
+	"github.com/kr/pretty"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
 	"github.com/hashicorp/nomad/client/allocdir"
 	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
 	"github.com/hashicorp/nomad/client/config"
@@ -26,16 +30,14 @@ import (
 	agentconsul "github.com/hashicorp/nomad/command/agent/consul"
 	mockdriver "github.com/hashicorp/nomad/drivers/mock"
 	"github.com/hashicorp/nomad/drivers/rawexec"
+	"github.com/hashicorp/nomad/helper"
 	"github.com/hashicorp/nomad/helper/testlog"
 	"github.com/hashicorp/nomad/helper/uuid"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/hashicorp/nomad/plugins/device"
 	"github.com/hashicorp/nomad/plugins/drivers"
 	"github.com/hashicorp/nomad/testutil"
-	"github.com/kr/pretty"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
 type MockTaskStateUpdater struct {
@@ -94,26 +96,30 @@ func testTaskRunnerConfig(t *testing.T, alloc *structs.Allocation, taskName stri
 		cleanup()
 	}
 
+	shutdownDelayCtx, shutdownDelayCancelFn := context.WithCancel(context.Background())
+
 	// Create a closed channel to mock TaskHookCoordinator.startConditionForTask.
 	// Closed channel indicates this task is not blocked on prestart hooks.
 	closedCh := make(chan struct{})
 	close(closedCh)
 
 	conf := &Config{
-		Alloc:                alloc,
-		ClientConfig:         clientConf,
-		Task:                 thisTask,
-		TaskDir:              taskDir,
-		Logger:               clientConf.Logger,
-		Consul:               consulapi.NewMockConsulServiceClient(t, logger),
-		ConsulSI:             consulapi.NewMockServiceIdentitiesClient(),
-		Vault:                vaultclient.NewMockVaultClient(),
-		StateDB:              cstate.NoopDB{},
-		StateUpdater:         NewMockTaskStateUpdater(),
-		DeviceManager:        devicemanager.NoopMockManager(),
-		DriverManager:        drivermanager.TestDriverManager(t),
-		ServersContactedCh:   make(chan struct{}),
-		StartConditionMetCtx: closedCh,
+		Alloc:                 alloc,
+		ClientConfig:          clientConf,
+		Task:                  thisTask,
+		TaskDir:               taskDir,
+		Logger:                clientConf.Logger,
+		Consul:                consulapi.NewMockConsulServiceClient(t, logger),
+		ConsulSI:              consulapi.NewMockServiceIdentitiesClient(),
+		Vault:                 vaultclient.NewMockVaultClient(),
+		StateDB:               cstate.NoopDB{},
+		StateUpdater:          NewMockTaskStateUpdater(),
+		DeviceManager:         devicemanager.NoopMockManager(),
+		DriverManager:         drivermanager.TestDriverManager(t),
+		ServersContactedCh:    make(chan struct{}),
+		StartConditionMetCtx:  closedCh,
+		ShutdownDelayCtx:      shutdownDelayCtx,
+		ShutdownDelayCancelFn: shutdownDelayCancelFn,
 	}
 	return conf, trCleanup
 }
@@ -996,6 +1002,82 @@ WAIT:
 	}
 }
 
+// TestTaskRunner_NoShutdownDelay asserts services are removed from
+// Consul and tasks are killed without waiting for ${shutdown_delay}
+// when the alloc has the NoShutdownDelay transition flag set.
+func TestTaskRunner_NoShutdownDelay(t *testing.T) {
+	t.Parallel()
+
+	// don't set this too high so that we don't block the test runner
+	// on shutting down the agent if the test fails
+	maxTestDuration := time.Duration(testutil.TestMultiplier()*10) * time.Second
+	maxTimeToFailDuration := time.Duration(testutil.TestMultiplier()) * time.Second
+
+	alloc := mock.Alloc()
+	alloc.DesiredTransition = structs.DesiredTransition{NoShutdownDelay: helper.BoolToPtr(true)}
+	task := alloc.Job.TaskGroups[0].Tasks[0]
+	task.Services[0].Tags = []string{"tag1"}
+	task.Services = task.Services[:1] // only need 1 for this test
+	task.Driver = "mock_driver"
+	task.Config = map[string]interface{}{
+		"run_for": "1000s",
+	}
+	task.ShutdownDelay = maxTestDuration
+
+	tr, conf, cleanup := runTestTaskRunner(t, alloc, task.Name)
+	defer cleanup()
+
+	mockConsul := conf.Consul.(*consulapi.MockConsulServiceClient)
+
+	testWaitForTaskToStart(t, tr)
+
+	testutil.WaitForResult(func() (bool, error) {
+		ops := mockConsul.GetOps()
+		if n := len(ops); n != 1 {
+			return false, fmt.Errorf("expected 1 consul operation. Found %d", n)
+		}
+		return ops[0].Op == "add", fmt.Errorf("consul operation was not a registration: %#v", ops[0])
+	}, func(err error) {
+		t.Fatalf("err: %v", err)
+	})
+
+	testCtx, cancel := context.WithTimeout(context.Background(), maxTimeToFailDuration)
+	defer cancel()
+
+	killed := make(chan error)
+	go func() {
+		tr.shutdownDelayCancel()
+		err := tr.Kill(testCtx, structs.NewTaskEvent("test"))
+		killed <- err
+	}()
+
+	// Wait for first de-registration call. Note that unlike
+	// TestTaskRunner_ShutdownDelay, we're racing with task exit
+	// and can't assert that we only get the first deregistration op
+	// (from serviceHook.PreKill).
+	testutil.WaitForResult(func() (bool, error) {
+		ops := mockConsul.GetOps()
+		if n := len(ops); n < 2 {
+			return false, fmt.Errorf("expected at least 2 consul operations.")
+		}
+		return ops[1].Op == "remove", fmt.Errorf(
+			"consul operation was not a deregistration: %#v", ops[1])
+	}, func(err error) {
+		t.Fatalf("err: %v", err)
+	})
+
+	// Wait for the task to exit
+	select {
+	case <-tr.WaitCh():
+	case <-time.After(maxTimeToFailDuration):
+		t.Fatalf("task kill did not ignore shutdown delay")
+		return
+	}
+
+	err := <-killed
+	require.NoError(t, err, "killing task returned unexpected error")
+}
+
 // TestTaskRunner_Dispatch_Payload asserts that a dispatch job runs and the
 // payload was written to disk.
 func TestTaskRunner_Dispatch_Payload(t *testing.T) {

command/agent/alloc_endpoint.go

@@ -138,8 +138,18 @@ func (s *HTTPServer) allocStop(allocID string, resp http.ResponseWriter, req *ht
 		return nil, CodedError(405, ErrInvalidMethod)
 	}
 
+	noShutdownDelay := false
+	if noShutdownDelayQS := req.URL.Query().Get("no_shutdown_delay"); noShutdownDelayQS != "" {
+		var err error
+		noShutdownDelay, err = strconv.ParseBool(noShutdownDelayQS)
+		if err != nil {
+			return nil, fmt.Errorf("no_shutdown_delay
 value is not a boolean: %v", err)
+		}
+	}
+
 	sr := &structs.AllocStopRequest{
-		AllocID: allocID,
+		AllocID:         allocID,
+		NoShutdownDelay: noShutdownDelay,
 	}
 	s.parseWriteRequest(req, &sr.WriteRequest)
 

command/agent/job_endpoint.go

@@ -453,6 +453,18 @@ func (s *HTTPServer) jobDelete(resp http.ResponseWriter, req *http.Request,
 		return nil, err
 	}
 
+	// Identify the no_shutdown_delay query param and parse.
+	noShutdownDelayStr := req.URL.Query().Get("no_shutdown_delay")
+	var noShutdownDelay bool
+	if noShutdownDelayStr != "" {
+		var err error
+		noShutdownDelay, err = strconv.ParseBool(noShutdownDelayStr)
+		if err != nil {
+			return nil, fmt.Errorf("Failed to parse value of %qq (%v) as a bool: %v", "no_shutdown_delay", noShutdownDelayStr, err)
+		}
+	}
+	args.NoShutdownDelay = noShutdownDelay
+
 	// Validate the evaluation priority if the user supplied a non-default
 	// value. It's more efficient to do it here, within the agent rather than
 	// sending a bad request for the server to reject.