func: add tests for new allocation methods

nomad/structs/group.go

@@ -131,11 +131,9 @@ func (ds *DisconnectStrategy) Canonicalize() {
 // ReconcileStrategy returns the strategy to be used when reconciling allocations
 // after a client reconnects. Best score is the default one.
 func (ds *DisconnectStrategy) ReconcileStrategy() string {
-
-	strategy := ReconcileOptionBestScore
-	if ds != nil || (ds != nil && ds.Reconcile != "") {
-		strategy = ds.Reconcile
+	if ds == nil || (ds != nil && ds.Reconcile == "") {
+		return ReconcileOptionBestScore
 	}
 
-	return strategy
+	return ds.Reconcile
 }

nomad/structs/group_test.go

@@ -117,6 +117,41 @@ func TestDisconnectStategy_Validate(t *testing.T) {
 	}
 }
 
+func TestReconcileStrategy(t *testing.T) {
+	ci.Parallel(t)
+
+	cases := []struct {
+		name            string
+		disconnectBlock *DisconnectStrategy
+		expected        string
+	}{
+		{
+			name:            "nil_disconnect_default_to_best_score",
+			disconnectBlock: nil,
+			expected:        ReconcileOptionBestScore,
+		},
+		{
+			name:            "empty_reconcile_default_to_best_score",
+			disconnectBlock: &DisconnectStrategy{},
+			expected:        ReconcileOptionBestScore,
+		},
+		{
+			name: "longest_running",
+			disconnectBlock: &DisconnectStrategy{
+				Reconcile: ReconcileOptionLongestRunning,
+			},
+			expected: ReconcileOptionLongestRunning,
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			rs := c.disconnectBlock.ReconcileStrategy()
+			must.Eq(t, c.expected, rs)
+		})
+	}
+}
+
 func TestJobConfig_Validate_StopAferClient_Disconnect(t *testing.T) {
 	ci.Parallel(t)
 	// Setup a system Job with Disconnect.StopOnClientAfter set, which is invalid

nomad/structs/structs.go

@@ -11449,8 +11449,14 @@ func (a *Allocation) GetLeaderTask() Task {
 	tg := a.Job.LookupTaskGroup(a.TaskGroup)
 
 	task := Task{}
+
+	if tg == nil {
+		return task
+	}
+
 	switch len(tg.Tasks) {
 	case 0:
+		return task
 
 	case 1:
 		task = *tg.Tasks[0]
@@ -11470,12 +11476,18 @@ func (a *Allocation) GetLeaderTask() Task {
 // LatestStartOfTask returns the time of the last start event for the given task
 // using the allocations TaskStates. If the task has not started, the zero time
 // will be returned.
-func (a *Allocation) LatestStartOfTask(task string) time.Time {
+func (a *Allocation) LatestStartOfTask(taskName string) time.Time {
 	t := time.Time{}
+
+	task := a.TaskStates[taskName]
+	if task == nil {
+		return t
+	}
+
 	// TaskStates are appended to the list and we only need the latest
 	// transition, so traverse from the end until we find one.
-	for i := len(a.TaskStates[task].Events) - 1; i >= 0; i-- {
-		e := a.TaskStates[task].Events[i]
+	for i := len(task.Events) - 1; i >= 0; i-- {
+		e := task.Events[i]
 		if e.Type == TaskStarted {
 			t = time.Unix(0, e.Time).UTC()
 			break

nomad/structs/structs_test.go

@@ -5891,6 +5891,251 @@ func TestAllocation_NeedsToReconnect(t *testing.T) {
 	}
 }
 
+func TestAllocation_ShouldBeReplaced(t *testing.T) {
+	ci.Parallel(t)
+
+	testAlloc := MockAlloc()
+
+	testCases := []struct {
+		name            string
+		taskGroup       string
+		disconnectGroup *DisconnectStrategy
+		expected        *bool
+	}{
+		{
+			name:      "missing_task_group",
+			taskGroup: "missing-task-group",
+			expected:  nil,
+		},
+		{
+			name:            "missing_disconnect_group",
+			taskGroup:       "web",
+			disconnectGroup: nil,
+			expected:        nil,
+		},
+		{
+			name:            "empty_disconnect_group",
+			taskGroup:       "web",
+			disconnectGroup: &DisconnectStrategy{},
+			expected:        nil,
+		},
+		{
+			name:      "replace_enabled",
+			taskGroup: "web",
+			disconnectGroup: &DisconnectStrategy{
+				Replace: pointer.Of(true),
+			},
+			expected: pointer.Of(true),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			alloc := testAlloc.Copy()
+			alloc.TaskGroup = tc.taskGroup
+
+			alloc.Job.TaskGroups[0].Disconnect = tc.disconnectGroup
+
+			got := alloc.ShouldBeReplaced()
+
+			require.Equal(t, tc.expected, got)
+		})
+	}
+}
+
+func TestAllocation_GetLeaderTask(t *testing.T) {
+	ci.Parallel(t)
+
+	testAlloc := MockAlloc()
+	testAlloc.Job.TaskGroups = []*TaskGroup{
+		{
+			Name:  "no-tasks",
+			Tasks: []*Task{},
+		},
+		{
+			Name: "one-task",
+			Tasks: []*Task{
+				{
+					Name: "task1",
+				},
+			},
+		},
+		{
+			Name: "multiple-tasks-no-leader",
+			Tasks: []*Task{
+				{
+					Name: "task1",
+				},
+				{
+					Name: "task2",
+				},
+				{
+					Name: "task3",
+				},
+			},
+		},
+		{
+			Name: "multiple-tasks-one-leader",
+			Tasks: []*Task{
+				{
+					Name: "task1",
+				},
+				{
+					Name:   "task2",
+					Leader: true,
+				},
+				{
+					Name: "task3",
+				},
+			},
+		},
+	}
+
+	testCases := []struct {
+		name      string
+		taskGroup string
+		expected  Task
+	}{
+		{
+			name:      "missing_task_group",
+			taskGroup: "missing-task-group",
+			expected:  Task{},
+		},
+		{
+			name:      "empty_tasks",
+			taskGroup: "no-tasks",
+			expected:  Task{},
+		},
+		{
+			name:      "one_task",
+			taskGroup: "one-task",
+			expected: Task{
+				Name: "task1",
+			},
+		},
+		{
+			name:      "multiple_tasks_no_leader",
+			taskGroup: "multiple-tasks-no-leader",
+			expected:  Task{},
+		},
+		{
+			name:      "multiple_tasks_one_leader",
+			taskGroup: "multiple-tasks-one-leader",
+			expected: Task{
+				Name:   "task2",
+				Leader: true,
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {})
+		ta := testAlloc.Copy()
+		ta.TaskGroup = tc.taskGroup
+
+		got := ta.GetLeaderTask()
+		must.Eq(t, tc.expected, got)
+	}
+}
+
+func TestAllocation_LatestStartOfTask(t *testing.T) {
+	ci.Parallel(t)
+	testNow := time.Now()
+
+	alloc := MockAlloc()
+	alloc.TaskStates = map[string]*TaskState{
+		"empty-events": {
+			Events: []*TaskEvent{},
+		},
+		"no-start": {
+			Events: []*TaskEvent{
+				{
+					Type: TaskReceived,
+					Time: testNow.UnixNano(),
+				},
+			},
+		},
+		"multiple-starts": {
+			Events: []*TaskEvent{
+				{
+					Type: TaskReceived,
+					Time: testNow.Add(-30 * time.Minute).UnixNano(),
+				},
+				{
+					Type: TaskStarted,
+					Time: testNow.Add(-20 * time.Minute).UnixNano(),
+				},
+				{
+					Type: TaskKilled,
+					Time: testNow.Add(-10 * time.Minute).UnixNano(),
+				},
+				{
+					Type: TaskStarted,
+					Time: testNow.Add(-5 * time.Minute).UnixNano(),
+				},
+			},
+		},
+		"disconnects": {
+			Events: []*TaskEvent{
+				{
+					Type: TaskReceived,
+					Time: testNow.Add(-30 * time.Minute).UnixNano(),
+				},
+				{
+					Type: TaskStarted,
+					Time: testNow.Add(-20 * time.Minute).UnixNano(),
+				},
+				{
+					Type: TaskClientReconnected,
+					Time: testNow.Add(-10 * time.Minute).UnixNano(),
+				},
+			},
+		},
+	}
+
+	testCases := []struct {
+		name     string
+		taskName string
+		expected time.Time
+	}{
+		{
+			name:     "missing_task",
+			taskName: "missing-task",
+			expected: time.Time{},
+		},
+		{
+			name:     "no_events",
+			taskName: "empty-events",
+			expected: time.Time{},
+		},
+		{
+			name:     "task_hasn't_started",
+			taskName: "no-start",
+			expected: time.Time{},
+		},
+		{
+			name:     "task_has_multiple_starts",
+			taskName: "multiple-starts",
+			expected: testNow.Add(-5 * time.Minute).UTC(),
+		},
+		{
+			name:     "task_has_disconnects",
+			taskName: "disconnects",
+			expected: testNow.Add(-20 * time.Minute).UTC(),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			alloc.TaskGroup = "web"
+			got := alloc.LatestStartOfTask(tc.taskName)
+
+			must.Eq(t, tc.expected, got)
+		})
+	}
+
+}
+
 func TestAllocation_Canonicalize_Old(t *testing.T) {
 	ci.Parallel(t)
 

scheduler/reconnecting_picker/reconnecting_picker.go

@@ -99,6 +99,7 @@ func (rp *ReconnectingPicker) pickLongestRunning(original *structs.Allocation, r
 		lt = *original.Job.LookupTaskGroup(original.TaskGroup).Tasks[0]
 	}
 
+	// If the replacement has a later start time, keep the original.
 	if original.LatestStartOfTask(lt.Name).Sub(replacement.LatestStartOfTask(lt.Name)) < 0 {
 		return original
 	}