scheduler/generic_sched_test: several scheduling cases

scheduler/generic_sched_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/hashicorp/nomad/helper/uuid"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
+	"github.com/hashicorp/nomad/testutil"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -2768,6 +2769,167 @@ func TestServiceSched_NodeDown(t *testing.T) {
 	}
 }
 
+func TestServiceSched_StopAfterClientDisconnect(t *testing.T) {
+	cases := []struct {
+		stop        time.Duration
+		when        time.Time
+		rescheduled bool
+	}{
+		{
+			rescheduled: true,
+		},
+		{
+			stop:        1 * time.Second,
+			rescheduled: false,
+		},
+		{
+			stop:        1 * time.Second,
+			when:        time.Now().UTC().Add(-10 * time.Second),
+			rescheduled: true,
+		},
+		{
+			stop:        1 * time.Second,
+			when:        time.Now().UTC().Add(10 * time.Second),
+			rescheduled: false,
+		},
+	}
+
+	for i, tc := range cases {
+		t.Run(fmt.Sprintf(""), func(t *testing.T) {
+			h := NewHarness(t)
+
+			// Node, which is down
+			node := mock.Node()
+			node.Status = structs.NodeStatusDown
+			require.NoError(t, h.State.UpsertNode(h.NextIndex(), node))
+
+			// Job with allocations and stop_after_client_disconnect
+			job := mock.Job()
+			job.TaskGroups[0].Count = 1
+			job.TaskGroups[0].StopAfterClientDisconnect = &tc.stop
+			require.NoError(t, h.State.UpsertJob(h.NextIndex(), job))
+
+			// Alloc for the running group
+			alloc := mock.Alloc()
+			alloc.Job = job
+			alloc.JobID = job.ID
+			alloc.NodeID = node.ID
+			alloc.Name = fmt.Sprintf("my-job.web[%d]", i)
+			alloc.DesiredStatus = structs.AllocDesiredStatusRun
+			alloc.ClientStatus = structs.AllocClientStatusRunning
+			if !tc.when.IsZero() {
+				alloc.AllocStates = []*structs.AllocState{{
+					Field: "ClientStatus",
+					Value: structs.AllocClientStatusLost,
+					Time:  tc.when,
+				}}
+			}
+			allocs := []*structs.Allocation{alloc}
+			require.NoError(t, h.State.UpsertAllocs(h.NextIndex(), allocs))
+
+			// Create a mock evaluation to deal with drain
+			evals := []*structs.Evaluation{{
+				Namespace:   structs.DefaultNamespace,
+				ID:          uuid.Generate(),
+				Priority:    50,
+				TriggeredBy: structs.EvalTriggerNodeDrain,
+				JobID:       job.ID,
+				NodeID:      node.ID,
+				Status:      structs.EvalStatusPending,
+			}}
+			eval := evals[0]
+			require.NoError(t, h.State.UpsertEvals(h.NextIndex(), evals))
+
+			// Process the evaluation
+			err := h.Process(NewServiceScheduler, eval)
+			require.NoError(t, err)
+			require.Equal(t, h.Evals[0].Status, structs.EvalStatusComplete)
+			require.Len(t, h.Plans, 1, "plan")
+
+			// Followup eval created
+			require.True(t, len(h.CreateEvals) > 0)
+			e := h.CreateEvals[0]
+			require.Equal(t, eval.ID, e.PreviousEval)
+
+			if tc.rescheduled {
+				require.Equal(t, "blocked", e.Status)
+			} else {
+				require.Equal(t, "pending", e.Status)
+				require.NotEmpty(t, e.WaitUntil)
+			}
+
+			// This eval is still being inserted in the state store
+			ws := memdb.NewWatchSet()
+			testutil.WaitForResult(func() (bool, error) {
+				found, err := h.State.EvalByID(ws, e.ID)
+				if err != nil {
+					return false, err
+				}
+				if found == nil {
+					return false, nil
+				}
+				return true, nil
+			}, func(err error) {
+				require.NoError(t, err)
+			})
+
+			alloc, err = h.State.AllocByID(ws, alloc.ID)
+			require.NoError(t, err)
+
+			// Allocations have been transitioned to lost
+			require.Equal(t, structs.AllocDesiredStatusStop, alloc.DesiredStatus)
+			require.Equal(t, structs.AllocClientStatusLost, alloc.ClientStatus)
+			// At least 1, 2 if we manually set the tc.when
+			require.NotEmpty(t, alloc.AllocStates)
+
+			if tc.rescheduled {
+				// Register a new node, leave it up, process the followup eval
+				node = mock.Node()
+				require.NoError(t, h.State.UpsertNode(h.NextIndex(), node))
+				require.NoError(t, h.Process(NewServiceScheduler, eval))
+
+				as, err := h.State.AllocsByJob(ws, job.Namespace, job.ID, false)
+				require.NoError(t, err)
+
+				testutil.WaitForResult(func() (bool, error) {
+					as, err := h.State.AllocsByJob(ws, job.Namespace, job.ID, false)
+					if err != nil {
+						return false, err
+					}
+					return len(as) == 2, nil
+				}, func(err error) {
+					require.NoError(t, err)
+				})
+
+				a2 := as[0]
+				if a2.ID == alloc.ID {
+					a2 = as[1]
+				}
+
+				require.Equal(t, structs.AllocClientStatusPending, a2.ClientStatus)
+				require.Equal(t, structs.AllocDesiredStatusRun, a2.DesiredStatus)
+				require.Equal(t, node.ID, a2.NodeID)
+
+				// No blocked evals
+				require.Empty(t, h.ReblockEvals)
+				require.Len(t, h.CreateEvals, 1)
+				require.Equal(t, h.CreateEvals[0].ID, e.ID)
+			} else {
+				// No new alloc was created
+				as, err := h.State.AllocsByJob(ws, job.Namespace, job.ID, false)
+				require.NoError(t, err)
+
+				require.Len(t, as, 1)
+				old := as[0]
+
+				require.Equal(t, alloc.ID, old.ID)
+				require.Equal(t, structs.AllocClientStatusLost, old.ClientStatus)
+				require.Equal(t, structs.AllocDesiredStatusStop, old.DesiredStatus)
+			}
+		})
+	}
+}
+
 func TestServiceSched_NodeUpdate(t *testing.T) {
 	h := NewHarness(t)