Refactor createTaskInfo

The function was growing very big and hard to get an overview of.
eremetic-framework · Sep 1, 2016 · 73c1b0f · 73c1b0f
1 parent a81d29f
commit 73c1b0f
Showing 1 changed file with 40 additions and 23 deletions.
diff --git a/scheduler/task.go b/scheduler/task.go
@@ -14,6 +14,31 @@ func createTaskInfo(task types.EremeticTask, offer *mesos.Offer) (types.Eremetic
 	task.AgentIP = offer.GetUrl().GetAddress().GetIp()
 	task.AgentPort = offer.GetUrl().GetAddress().GetPort()
 
+	volumes := collectVolumes(task)
+	commandInfo := collectCommandInfo(task)
+
+	taskInfo := &mesos.TaskInfo{
+		TaskId:  &mesos.TaskID{Value: proto.String(task.ID)},
+		SlaveId: offer.SlaveId,
+		Name:    proto.String(task.Name),
+		Command: commandInfo,
+		Container: &mesos.ContainerInfo{
+			Type: mesos.ContainerInfo_DOCKER.Enum(),
+			Docker: &mesos.ContainerInfo_DockerInfo{
+				Image:          proto.String(task.Image),
+				ForcePullImage: proto.Bool(task.ForcePullImage),
+			},
+			Volumes: volumes,
+		},
+		Resources: []*mesos.Resource{
+			mesosutil.NewScalarResource("cpus", task.TaskCPUs),
+			mesosutil.NewScalarResource("mem", task.TaskMem),
+		},
+	}
+	return task, taskInfo
+}
+
+func collectEnvironment(task types.EremeticTask) []*mesos.Environment_Variable {
 	var environment []*mesos.Environment_Variable
 	for k, v := range task.Environment {
 		environment = append(environment, &mesos.Environment_Variable{
@@ -33,6 +58,10 @@ func createTaskInfo(task types.EremeticTask, offer *mesos.Offer) (types.Eremetic
 		Value: proto.String(task.ID),
 	})
 
+	return environment
+}
+
+func collectVolumes(task types.EremeticTask) []*mesos.Volume {
 	var volumes []*mesos.Volume
 	for _, v := range task.Volumes {
 		volumes = append(volumes, &mesos.Volume{
@@ -42,6 +71,10 @@ func createTaskInfo(task types.EremeticTask, offer *mesos.Offer) (types.Eremetic
 		})
 	}
 
+	return volumes
+}
+
+func collectURIs(task types.EremeticTask) []*mesos.CommandInfo_URI {
 	var uris []*mesos.CommandInfo_URI
 	for _, v := range task.FetchURIs {
 		uris = append(uris, &mesos.CommandInfo_URI{
@@ -52,12 +85,16 @@ func createTaskInfo(task types.EremeticTask, offer *mesos.Offer) (types.Eremetic
 		})
 	}
 
+	return uris
+}
+
+func collectCommandInfo(task types.EremeticTask) *mesos.CommandInfo {
 	commandInfo := &mesos.CommandInfo{
 		User: proto.String(task.User),
 		Environment: &mesos.Environment{
-			Variables: environment,
+			Variables: collectEnvironment(task),
 		},
-		Uris: uris,
+		Uris: collectURIs(task),
 	}
 
 	if task.Command != "" {
@@ -66,25 +103,5 @@ func createTaskInfo(task types.EremeticTask, offer *mesos.Offer) (types.Eremetic
 		commandInfo.Shell = proto.Bool(false)
 	}
 
-	taskInfo := &mesos.TaskInfo{
-		TaskId: &mesos.TaskID{
-			Value: proto.String(task.ID),
-		},
-		SlaveId: offer.SlaveId,
-		Name:    proto.String(task.Name),
-		Command: commandInfo,
-		Container: &mesos.ContainerInfo{
-			Type: mesos.ContainerInfo_DOCKER.Enum(),
-			Docker: &mesos.ContainerInfo_DockerInfo{
-				Image:          proto.String(task.Image),
-				ForcePullImage: proto.Bool(task.ForcePullImage),
-			},
-			Volumes: volumes,
-		},
-		Resources: []*mesos.Resource{
-			mesosutil.NewScalarResource("cpus", task.TaskCPUs),
-			mesosutil.NewScalarResource("mem", task.TaskMem),
-		},
-	}
-	return task, taskInfo
+	return commandInfo
 }