address code review requests

Author: lgfa29

.changelog/14127.txt

@@ -1,5 +1,5 @@
 ```release-note:improvement
-client: add option to restart all tasks of an allocation, even if the task already run, such as non-sidecar prestart and poststart tasks.
+client: add option to restart all tasks of an allocation, regardless of lifecycle type or state.
 ```
 
 ```release-note:improvement

api/allocations.go

@@ -141,7 +141,9 @@ func (a *Allocations) GC(alloc *Allocation, q *QueryOptions) error {
 	return err
 }
 
-// Restart restarts an allocation.
+// Restart restarts the tasks that are currently running or a specific task if
+// taskName is provided. An error is returned if the task to be restarted is
+// not running.
 //
 // Note: for cluster topologies where API consumers don't have network access to
 // Nomad clients, set api.ClientConnTimeout to a small value (ex 1ms) to avoid
@@ -156,6 +158,12 @@ func (a *Allocations) Restart(alloc *Allocation, taskName string, q *QueryOption
 	return err
 }
 
+// RestartAllTasks restarts all tasks in the allocation, regardless of
+// lifecycle type or state. Tasks will restart following their lifecycle order.
+//
+// Note: for cluster topologies where API consumers don't have network access to
+// Nomad clients, set api.ClientConnTimeout to a small value (ex 1ms) to avoid
+// long pauses on this API call.
 func (a *Allocations) RestartAllTasks(alloc *Allocation, q *QueryOptions) error {
 	req := AllocationRestartRequest{
 		AllTasks: true,

client/allocrunner/alloc_runner.go

@@ -548,7 +548,8 @@ func (ar *allocRunner) handleTaskStateUpdates() {
 
 		if len(liveRunners) > 0 {
 			// if all live runners are sidecars - kill alloc
-			if killEvent == nil && hasSidecars && !hasNonSidecarTasks(liveRunners) {
+			hasLiveSidecars := hasSidecars && !hasNonSidecarTasks(liveRunners)
+			if killEvent == nil && hasLiveSidecars {
 				killEvent = structs.NewTaskEvent(structs.TaskMainDead)
 			}
 
@@ -573,11 +574,11 @@ func (ar *allocRunner) handleTaskStateUpdates() {
 				// Kill 'em all
 				states = ar.killTasks()
 
-				// Wait for TaskRunners to exit before continuing to
-				// prevent looping before TaskRunners have transitioned
-				// to Dead.
+				// Wait for TaskRunners to exit before continuing. This will
+				// prevent looping before TaskRunners have transitioned to
+				// Dead.
 				for _, tr := range liveRunners {
-					ar.logger.Info("killing task", "task", tr.Task().Name)
+					ar.logger.Info("waiting for task to exit", "task", tr.Task().Name)
 					select {
 					case <-tr.WaitCh():
 					case <-ar.waitCh:
@@ -1228,20 +1229,24 @@ func (ar *allocRunner) GetTaskEventHandler(taskName string) drivermanager.EventH
 }
 
 // Restart satisfies the WorkloadRestarter interface and restarts all tasks
-// that are currently running. Only the TaskRestartRunningSignal event type may
-// be used.
+// that are currently running.
+//
+// The event type will be set to TaskRestartRunningSignal to comply with
+// internal restart logic requirements.
 func (ar *allocRunner) Restart(ctx context.Context, event *structs.TaskEvent, failure bool) error {
 	if event.Type != structs.TaskRestartRunningSignal {
-		return fmt.Errorf("Invalid event %s for alloc restart request", event.Type)
+		event.Type = structs.TaskRestartRunningSignal
 	}
 	return ar.restartTasks(ctx, event, failure)
 }
 
-// RestartTask restarts the provided task. Only TaskRestartSignal event type
-// may be used.
+// RestartTask restarts the provided task.
+//
+// The event type will be set to TaskRestartSignal to comply with internal
+// restart logic requirements.
 func (ar *allocRunner) RestartTask(taskName string, event *structs.TaskEvent) error {
 	if event.Type != structs.TaskRestartSignal {
-		return fmt.Errorf("Invalid event %s for task restart request", event.Type)
+		event.Type = structs.TaskRestartSignal
 	}
 
 	tr, ok := ar.tasks[taskName]
@@ -1252,11 +1257,13 @@ func (ar *allocRunner) RestartTask(taskName string, event *structs.TaskEvent) er
 	return tr.Restart(context.TODO(), event, false)
 }
 
-// RestartRunning restarts all tasks that are currently running. Only the
-// TaskRestartRunningSignal event type may be used.
+// RestartRunning restarts all tasks that are currently running.
+//
+// The event type will be set to TaskRestartRunningSignal to comply with
+// internal restart logic requirements.
 func (ar *allocRunner) RestartRunning(event *structs.TaskEvent) error {
 	if event.Type != structs.TaskRestartRunningSignal {
-		return fmt.Errorf("Invalid event %s for running tasks restart request", event.Type)
+		event.Type = structs.TaskRestartRunningSignal
 	}
 	return ar.restartTasks(context.TODO(), event, false)
 }

client/allocrunner/alloc_runner_test.go

@@ -1676,11 +1676,7 @@ func TestAllocRunner_MoveAllocDir(t *testing.T) {
 	ar.Run()
 	defer destroy(ar)
 
-	testutil.WaitForResult(func() (bool, error) {
-		return ar.AllocState().ClientStatus == structs.AllocClientStatusComplete, nil
-	}, func(_ error) {
-		t.Fatalf("expected alloc to be complete")
-	})
+	WaitForClientState(t, ar, structs.AllocClientStatusComplete)
 
 	// Step 2. Modify its directory
 	task := alloc.Job.TaskGroups[0].Tasks[0]
@@ -1708,11 +1704,7 @@ func TestAllocRunner_MoveAllocDir(t *testing.T) {
 	ar2.Run()
 	defer destroy(ar2)
 
-	testutil.WaitForResult(func() (bool, error) {
-		return ar.AllocState().ClientStatus == structs.AllocClientStatusComplete, nil
-	}, func(_ error) {
-		t.Fatalf("expected alloc to be complete")
-	})
+	WaitForClientState(t, ar, structs.AllocClientStatusComplete)
 
 	// Ensure that data from ar was moved to ar2
 	dataFile = filepath.Join(ar2.allocDir.SharedDir, "data", "data_file")

client/allocrunner/testing.go

@@ -4,6 +4,7 @@
 package allocrunner
 
 import (
+	"fmt"
 	"sync"
 	"testing"
 
@@ -20,6 +21,7 @@ import (
 	"github.com/hashicorp/nomad/client/state"
 	"github.com/hashicorp/nomad/client/vaultclient"
 	"github.com/hashicorp/nomad/nomad/structs"
+	"github.com/hashicorp/nomad/testutil"
 	"github.com/stretchr/testify/require"
 )
 
@@ -104,3 +106,13 @@ func TestAllocRunnerFromAlloc(t *testing.T, alloc *structs.Allocation) (*allocRu
 
 	return ar, cleanup
 }
+
+func WaitForClientState(t *testing.T, ar *allocRunner, state string) {
+	testutil.WaitForResult(func() (bool, error) {
+		got := ar.AllocState().ClientStatus
+		return got == state,
+			fmt.Errorf("expected alloc runner to be in state %s, got %s", state, got)
+	}, func(err error) {
+		require.NoError(t, err)
+	})
+}

command/alloc_restart.go

@@ -36,7 +36,8 @@ Restart Specific Options:
 
   -all-tasks
     If set, all tasks in the allocation will be restarted, even the ones that
-    already ran. This option cannot be used if a task is defined.
+    already ran. This option cannot be used with '-task' or the '<task>'
+    argument.
 
   -task <task-name>
     Specify the individual task to restart. If task name is given with both an