Fix bug that the count of the default task group is set to the prewarming pool size

nomad/job.go

@@ -12,7 +12,9 @@ const (
 	TemplateJobPrefix     = "template"
 	TaskGroupName         = "default-group"
 	TaskName              = "default-task"
+	TaskCount             = 1
 	TaskDriver            = "docker"
+	TaskCommand           = "sleep"
 	ConfigTaskGroupName   = "config"
 	ConfigTaskName        = "config"
 	ConfigTaskDriver      = "exec"
@@ -25,6 +27,7 @@ const (
 )
 
 var (
+	TaskArgs                     = []string{"infinity"}
 	ErrorConfigTaskGroupNotFound = errors.New("config task group not found in job")
 )
 
@@ -79,7 +82,7 @@ func CreateTemplateJob(
 	job.ID = &id
 	job.Name = &id
 
-	var taskGroup = createTaskGroup(&job, TaskGroupName, prewarmingPoolSize)
+	var taskGroup = createTaskGroup(&job, TaskGroupName)
 	configureTask(taskGroup, TaskName, cpuLimit, memoryLimit, image, networkAccess, exposedPorts)
 	storeTemplateConfiguration(&job, prewarmingPoolSize)
 
@@ -96,15 +99,15 @@ func (a *APIClient) RegisterRunnerJob(template *nomadApi.Job) error {
 	return a.MonitorEvaluation(evalID, context.Background())
 }
 
-func createTaskGroup(job *nomadApi.Job, name string, prewarmingPoolSize uint) *nomadApi.TaskGroup {
+func createTaskGroup(job *nomadApi.Job, name string) *nomadApi.TaskGroup {
 	var taskGroup *nomadApi.TaskGroup
 	if len(job.TaskGroups) == 0 {
-		taskGroup = nomadApi.NewTaskGroup(name, int(prewarmingPoolSize))
+		taskGroup = nomadApi.NewTaskGroup(name, TaskCount)
 		job.TaskGroups = []*nomadApi.TaskGroup{taskGroup}
 	} else {
 		taskGroup = job.TaskGroups[0]
 		taskGroup.Name = &name
-		count := 1
+		count := TaskCount
 		taskGroup.Count = &count
 	}
 	return taskGroup
@@ -173,7 +176,7 @@ func configureTask(
 	integerCPULimit := int(cpuLimit)
 	integerMemoryLimit := int(memoryLimit)
 	if task.Resources == nil {
-		task.Resources = &nomadApi.Resources{}
+		task.Resources = nomadApi.DefaultResources()
 	}
 	task.Resources.CPU = &integerCPULimit
 	task.Resources.MemoryMB = &integerMemoryLimit
@@ -182,6 +185,8 @@ func configureTask(
 		task.Config = make(map[string]interface{})
 	}
 	task.Config["image"] = image
+	task.Config["command"] = TaskCommand
+	task.Config["args"] = TaskArgs
 
 	configureNetwork(taskGroup, networkAccess, exposedPorts)
 }
@@ -206,12 +211,12 @@ func findOrCreateConfigTaskGroup(job *nomadApi.Job) *nomadApi.TaskGroup {
 		taskGroup = nomadApi.NewTaskGroup(ConfigTaskGroupName, 0)
 		job.AddTaskGroup(taskGroup)
 	}
-	createDummyTaskIfNotPresent(taskGroup)
+	createConfigTaskIfNotPresent(taskGroup)
 	return taskGroup
 }
 
-// createDummyTaskIfNotPresent ensures that a dummy task is in the task group so that the group is accepted by Nomad.
+// createConfigTaskIfNotPresent ensures that a dummy task is in the task group so that the group is accepted by Nomad.
-func createDummyTaskIfNotPresent(taskGroup *nomadApi.TaskGroup) {
+func createConfigTaskIfNotPresent(taskGroup *nomadApi.TaskGroup) {
 	var task *nomadApi.Task
 	for _, t := range taskGroup.Tasks {
 		if t.Name == ConfigTaskName {

nomad/job_test.go

@@ -15,7 +15,7 @@ import (
 )
 
 func createTestTaskGroup() *nomadApi.TaskGroup {
-	return nomadApi.NewTaskGroup("taskGroup", 42)
+	return nomadApi.NewTaskGroup("taskGroup", 1)
 }
 
 func createTestTask() *nomadApi.Task {
@@ -23,9 +23,12 @@ func createTestTask() *nomadApi.Task {
 }
 
 func createTestResources() *nomadApi.Resources {
+	result := nomadApi.DefaultResources()
 	expectedCPULimit := 1337
 	expectedMemoryLimit := 42
-	return &nomadApi.Resources{CPU: &expectedCPULimit, MemoryMB: &expectedMemoryLimit}
+	result.CPU = &expectedCPULimit
+	result.MemoryMB = &expectedMemoryLimit
+	return result
 }
 
 func TestCreateTaskGroupCreatesNewTaskGroupWhenJobHasNoTaskGroup(t *testing.T) {
@@ -33,7 +36,7 @@ func TestCreateTaskGroupCreatesNewTaskGroupWhenJobHasNoTaskGroup(t *testing.T) {
 
 	if assert.Equal(t, 0, len(job.TaskGroups)) {
 		expectedTaskGroup := createTestTaskGroup()
-		taskGroup := createTaskGroup(job, *expectedTaskGroup.Name, uint(*expectedTaskGroup.Count))
+		taskGroup := createTaskGroup(job, *expectedTaskGroup.Name)
 
 		assert.Equal(t, *expectedTaskGroup, *taskGroup)
 		assert.Equal(t, []*nomadApi.TaskGroup{taskGroup}, job.TaskGroups, "it should add the task group to the job")
@@ -48,9 +51,7 @@ func TestCreateTaskGroupOverwritesOptionsWhenJobHasTaskGroup(t *testing.T) {
 	job.TaskGroups = newTaskGroupList
 
 	newName := *existingTaskGroup.Name + "longerName"
-	newCount := *existingTaskGroup.Count + 42
+	taskGroup := createTaskGroup(job, newName)
-
-	taskGroup := createTaskGroup(job, newName, uint(newCount))
 
 	// create a new copy to avoid changing the original one as it is a pointer
 	expectedTaskGroup := *existingTaskGroup
@@ -166,7 +167,10 @@ func TestConfigureTaskWhenNoTaskExists(t *testing.T) {
 	expectedTask := nomadApi.NewTask("task", TaskDriver)
 	expectedTask.Resources = expectedResources
 	expectedImage := "python:latest"
-	expectedTask.Config = map[string]interface{}{"image": expectedImage, "network_mode": "none"}
+	expectedCommand := "sleep"
+	expectedArgs := []string{"infinity"}
+	expectedTask.Config = map[string]interface{}{
+		"image": expectedImage, "command": expectedCommand, "args": expectedArgs, "network_mode": "none"}
 	expectedTaskGroup.Tasks = []*nomadApi.Task{expectedTask}
 	expectedTaskGroup.Networks = []*nomadApi.NetworkResource{}
 
@@ -203,11 +207,11 @@ func TestConfigureTaskWhenTaskExists(t *testing.T) {
 }
 
 func TestCreateTemplateJobSetsAllGivenArguments(t *testing.T) {
-	testJob := helpers.CreateTemplateJob()
+	base, testJob := helpers.CreateTemplateJob()
 	prewarmingPoolSize, err := strconv.Atoi(testJob.TaskGroups[1].Meta[ConfigMetaPoolSizeKey])
 	require.NoError(t, err)
 	job := CreateTemplateJob(
-		helpers.CreateTemplateJob(),
+		base,
 		tests.DefaultJobID,
 		uint(prewarmingPoolSize),
 		uint(*testJob.TaskGroups[0].Tasks[0].Resources.CPU),

runner/manager_test.go

@@ -257,7 +257,7 @@ func (s *ManagerTestSuite) TestUpdateRunnersRemovesIdleAndUsedRunner() {
 }
 
 func (s *ManagerTestSuite) TestUpdateEnvironmentRemovesIdleRunnersWhenScalingDown() {
-	job := helpers.CreateTemplateJob()
+	_, job := helpers.CreateTemplateJob()
 	initialRunners := uint(40)
 	updatedRunners := uint(10)
 	err := s.nomadRunnerManager.registerEnvironment(anotherEnvironmentID, initialRunners, job, true)

tests/helpers/test_helpers.go

@@ -141,17 +141,14 @@ func HttpPutJSON(url string, body interface{}) (response *http.Response, err err
 	return HttpPut(url, reader)
 }
 
-func CreateTemplateJob() (job *nomadApi.Job) {
+func CreateTemplateJob() (base, job *nomadApi.Job) {
+	base = nomadApi.NewBatchJob(tests.DefaultJobID, tests.DefaultJobID, "region-name", 100)
 	job = nomadApi.NewBatchJob(tests.DefaultJobID, tests.DefaultJobID, "region-name", 100)
 	configTaskGroup := nomadApi.NewTaskGroup("config", 0)
 	configTaskGroup.Meta = make(map[string]string)
 	configTaskGroup.Meta["prewarmingPoolSize"] = "0"
 	configTask := nomadApi.NewTask("config", "exec")
-	configTask.Config = map[string]interface{}{
+	configTask.Config = map[string]interface{}{"command": "true"}
-		"command": "true",
-		"image":   "python:latest",
-	}
-	configTask.Resources = nomadApi.DefaultResources()
 	configTaskGroup.AddTask(configTask)
 
 	defaultTaskGroup := nomadApi.NewTaskGroup("default-group", 1)
@@ -167,5 +164,5 @@ func CreateTemplateJob() (job *nomadApi.Job) {
 	defaultTaskGroup.AddTask(defaultTask)
 
 	job.TaskGroups = []*nomadApi.TaskGroup{defaultTaskGroup, configTaskGroup}
-	return job
+	return base, job
 }