Switch to custom task state enum

eremetic-framework · Aug 17, 2016 · 0713c06 · 0713c06
1 parent 560ab7f
commit 0713c06
Show file tree

Hide file tree

Showing 11 changed files with 58 additions and 51 deletions.
diff --git a/database/database_test.go b/database/database_test.go
@@ -8,7 +8,6 @@ import (
 	"time"
 
 	"github.com/klarna/eremetic/types"
-	mesos "github.com/mesos/mesos-go/mesosproto"
 	. "github.com/smartystreets/goconvey/convey"
 )
 
@@ -20,7 +19,7 @@ func setup() error {
 func TestDatabase(t *testing.T) {
 	status := []types.Status{
 		types.Status{
-			Status: mesos.TaskState_TASK_RUNNING.String(),
+			Status: types.TaskState_TASK_RUNNING,
 			Time:   time.Now().Unix(),
 		},
 	}
@@ -106,15 +105,15 @@ func TestDatabase(t *testing.T) {
 			ID: "1234",
 			Status: []types.Status{
 				types.Status{
-					Status: mesos.TaskState_TASK_STAGING.String(),
+					Status: types.TaskState_TASK_STAGING,
 					Time:   time.Now().Unix(),
 				},
 				types.Status{
-					Status: mesos.TaskState_TASK_RUNNING.String(),
+					Status: types.TaskState_TASK_RUNNING,
 					Time:   time.Now().Unix(),
 				},
 				types.Status{
-					Status: mesos.TaskState_TASK_FINISHED.String(),
+					Status: types.TaskState_TASK_FINISHED,
 					Time:   time.Now().Unix(),
 				},
 			},
@@ -125,11 +124,11 @@ func TestDatabase(t *testing.T) {
 			ID: "2345",
 			Status: []types.Status{
 				types.Status{
-					Status: mesos.TaskState_TASK_STAGING.String(),
+					Status: types.TaskState_TASK_STAGING,
 					Time:   time.Now().Unix(),
 				},
 				types.Status{
-					Status: mesos.TaskState_TASK_RUNNING.String(),
+					Status: types.TaskState_TASK_RUNNING,
 					Time:   time.Now().Unix(),
 				},
 			},

diff --git a/handler/handler.go b/handler/handler.go
@@ -96,7 +96,7 @@ func NotifyCallback(task *types.EremeticTask) {
 
 	cbData := callbackData{
 		Time:   task.Status[len(task.Status)-1].Time,
-		Status: task.Status[len(task.Status)-1].Status,
+		Status: string(task.Status[len(task.Status)-1].Status),
 		TaskID: task.ID,
 	}
 

diff --git a/handler/handler_test.go b/handler/handler_test.go
@@ -13,7 +13,6 @@ import (
 	"github.com/gorilla/mux"
 	"github.com/klarna/eremetic/database"
 	"github.com/klarna/eremetic/types"
-	mesos "github.com/mesos/mesos-go/mesosproto"
 	. "github.com/smartystreets/goconvey/convey"
 )
 
@@ -36,7 +35,7 @@ func TestHandling(t *testing.T) {
 	scheduler := &mockScheduler{}
 	status := []types.Status{
 		types.Status{
-			Status: mesos.TaskState_TASK_RUNNING.String(),
+			Status: types.TaskState_TASK_RUNNING,
 			Time:   time.Now().Unix(),
 		},
 	}

diff --git a/scheduler/reconcile_test.go b/scheduler/reconcile_test.go
@@ -8,7 +8,6 @@ import (
 
 	"github.com/klarna/eremetic/database"
 	"github.com/klarna/eremetic/types"
-	mesos "github.com/mesos/mesos-go/mesosproto"
 	. "github.com/smartystreets/goconvey/convey"
 	"github.com/stretchr/testify/mock"
 )
@@ -41,7 +40,7 @@ func TestReconcile(t *testing.T) {
 					panic("mock error")
 				}
 				t.UpdateStatus(types.Status{
-					Status: mesos.TaskState_TASK_RUNNING.String(),
+					Status: types.TaskState_TASK_RUNNING,
 					Time:   time.Now().Unix() + 1,
 				})
 				database.PutTask(&t)
@@ -51,7 +50,7 @@ func TestReconcile(t *testing.T) {
 				ID: "1234",
 				Status: []types.Status{
 					types.Status{
-						Status: mesos.TaskState_TASK_STAGING.String(),
+						Status: types.TaskState_TASK_STAGING,
 						Time:   time.Now().Unix(),
 					},
 				},
@@ -73,7 +72,7 @@ func TestReconcile(t *testing.T) {
 				ID: "1234",
 				Status: []types.Status{
 					types.Status{
-						Status: mesos.TaskState_TASK_STAGING.String(),
+						Status: types.TaskState_TASK_STAGING,
 						Time:   time.Now().Unix(),
 					},
 				},

diff --git a/scheduler/scheduler.go b/scheduler/scheduler.go
@@ -123,6 +123,7 @@ loop:
 // StatusUpdate takes care of updating the status
 func (s *eremeticScheduler) StatusUpdate(driver sched.SchedulerDriver, status *mesos.TaskStatus) {
 	id := status.TaskId.GetValue()
+	newState := types.TaskState(status.State.String())
 
 	log.Debugf("Received task status [%s] for task [%s]", status.State.String(), id)
 
@@ -138,29 +139,29 @@ func (s *eremeticScheduler) StatusUpdate(driver sched.SchedulerDriver, status *m
 		}
 	}
 
-	if !task.IsRunning() && *status.State == mesos.TaskState_TASK_RUNNING {
+	if !task.IsRunning() && newState == types.TaskState_TASK_RUNNING {
 		TasksRunning.Inc()
 	}
 
-	if types.IsTerminal(status.State.String()) {
-		TasksTerminated.With(prometheus.Labels{"status": status.State.String()}).Inc()
+	if types.IsTerminal(newState) {
+		TasksTerminated.With(prometheus.Labels{"status": string(newState)}).Inc()
 		if task.WasRunning() {
 			TasksRunning.Dec()
 		}
 	}
 
 	task.UpdateStatus(types.Status{
-		Status: status.State.String(),
+		Status: newState,
 		Time:   time.Now().Unix(),
 	})
 
-	if *status.State == mesos.TaskState_TASK_FAILED && !task.WasRunning() {
+	if newState == types.TaskState_TASK_FAILED && !task.WasRunning() {
 		if task.Retry >= maxRetries {
 			log.Warnf("giving up on %s after %d retry attempts", id, task.Retry)
 		} else {
 			log.Infof("task %s was never running. re-scheduling", id)
 			task.UpdateStatus(types.Status{
-				Status: mesos.TaskState_TASK_STAGING.String(),
+				Status: types.TaskState_TASK_STAGING,
 				Time:   time.Now().Unix(),
 			})
 			task.Retry++
@@ -171,7 +172,7 @@ func (s *eremeticScheduler) StatusUpdate(driver sched.SchedulerDriver, status *m
 		}
 	}
 
-	if types.IsTerminal(status.State.String()) {
+	if types.IsTerminal(newState) {
 		handler.NotifyCallback(&task)
 	}
 

diff --git a/scheduler/scheduler_test.go b/scheduler/scheduler_test.go
@@ -101,7 +101,7 @@ func TestScheduler(t *testing.T) {
 					})
 					task, _ := database.ReadTask(id)
 					So(len(task.Status), ShouldEqual, 1)
-					So(task.Status[0].Status, ShouldEqual, mesos.TaskState_TASK_RUNNING.String())
+					So(task.Status[0].Status, ShouldEqual, types.TaskState_TASK_RUNNING)
 
 					s.StatusUpdate(nil, &mesos.TaskStatus{
 						TaskId: &mesos.TaskID{
@@ -112,8 +112,8 @@ func TestScheduler(t *testing.T) {
 					task, _ = database.ReadTask(id)
 
 					So(len(task.Status), ShouldEqual, 2)
-					So(task.Status[0].Status, ShouldEqual, mesos.TaskState_TASK_RUNNING.String())
-					So(task.Status[1].Status, ShouldEqual, mesos.TaskState_TASK_FAILED.String())
+					So(task.Status[0].Status, ShouldEqual, types.TaskState_TASK_RUNNING)
+					So(task.Status[1].Status, ShouldEqual, types.TaskState_TASK_FAILED)
 				})
 
 				Convey("Failing immediatly", func() {
@@ -125,9 +125,9 @@ func TestScheduler(t *testing.T) {
 						State: mesos.TaskState_TASK_FAILED.Enum(),
 					})
 					task, _ := database.ReadTask(id)
-					So(len(task.Status), ShouldEqual, 2)
-					So(task.Status[0].Status, ShouldEqual, mesos.TaskState_TASK_FAILED.String())
-					So(task.Status[1].Status, ShouldEqual, mesos.TaskState_TASK_STAGING.String())
+					So(task.Status, ShouldHaveLength, 2)
+					So(task.Status[0].Status, ShouldEqual, types.TaskState_TASK_FAILED)
+					So(task.Status[1].Status, ShouldEqual, types.TaskState_TASK_STAGING)
 
 					select {
 					case c := <-s.tasks:

diff --git a/scheduler/task.go b/scheduler/task.go
@@ -38,7 +38,7 @@ func createEremeticTask(request types.Request) (types.EremeticTask, error) {
 
 	status := []types.Status{
 		types.Status{
-			Status: mesos.TaskState_TASK_STAGING.String(),
+			Status: types.TaskState_TASK_STAGING,
 			Time:   time.Now().Unix(),
 		},
 	}

diff --git a/scheduler/task_test.go b/scheduler/task_test.go
@@ -14,7 +14,7 @@ func TestTask(t *testing.T) {
 
 	status := []types.Status{
 		types.Status{
-			Status: mesos.TaskState_TASK_RUNNING.String(),
+			Status: types.TaskState_TASK_RUNNING,
 			Time:   time.Now().Unix(),
 		},
 	}

diff --git a/types/states.go b/types/states.go
@@ -1,8 +1,21 @@
 package types
 
+type TaskState string
+
+const (
+	TaskState_TASK_STAGING  TaskState = "TASK_STAGING"
+	TaskState_TASK_STARTING TaskState = "TASK_STARTING"
+	TaskState_TASK_RUNNING  TaskState = "TASK_RUNNING"
+	TaskState_TASK_FINISHED TaskState = "TASK_FINISHED"
+	TaskState_TASK_FAILED   TaskState = "TASK_FAILED"
+	TaskState_TASK_KILLED   TaskState = "TASK_KILLED"
+	TaskState_TASK_LOST     TaskState = "TASK_LOST"
+	TaskState_TASK_ERROR    TaskState = "TASK_ERROR"
+)
+
 // IsTerminal takes a string representation of a state and returns whether it
 // is terminal or not.
-func IsTerminal(state string) bool {
+func IsTerminal(state TaskState) bool {
 	switch state {
 	case "TASK_LOST", "TASK_KILLED", "TASK_FAILED", "TASK_FINISHED":
 		return true

diff --git a/types/task.go b/types/task.go
@@ -2,13 +2,11 @@ package types
 
 import (
 	"time"
-
-	mesos "github.com/mesos/mesos-go/mesosproto"
 )
 
 type Status struct {
-	Time   int64  `json:"time"`
-	Status string `json:"status"`
+	Time   int64     `json:"time"`
+	Status TaskState `json:"status"`
 }
 
 type EremeticTask struct {
@@ -32,7 +30,7 @@ type EremeticTask struct {
 
 func (task *EremeticTask) WasRunning() bool {
 	for _, s := range task.Status {
-		if s.Status == mesos.TaskState_TASK_RUNNING.String() {
+		if s.Status == TaskState_TASK_RUNNING {
 			return true
 		}
 	}
@@ -52,7 +50,7 @@ func (task *EremeticTask) IsRunning() bool {
 		return false
 	}
 	st := task.Status[len(task.Status)-1]
-	return st.Status == mesos.TaskState_TASK_RUNNING.String()
+	return st.Status == TaskState_TASK_RUNNING
 }
 
 func (task *EremeticTask) LastUpdated() time.Time {

diff --git a/types/types_test.go b/types/types_test.go
@@ -4,38 +4,36 @@ import (
 	"fmt"
 	"testing"
 
-	mesos "github.com/mesos/mesos-go/mesosproto"
-
 	. "github.com/smartystreets/goconvey/convey"
 )
 
 func TestTypes(t *testing.T) {
 	Convey("states", t, func() {
-		terminalStates := []mesos.TaskState{
-			mesos.TaskState_TASK_FINISHED,
-			mesos.TaskState_TASK_FAILED,
-			mesos.TaskState_TASK_KILLED,
-			mesos.TaskState_TASK_LOST,
+		terminalStates := []TaskState{
+			TaskState_TASK_FINISHED,
+			TaskState_TASK_FAILED,
+			TaskState_TASK_KILLED,
+			TaskState_TASK_LOST,
 		}
 
-		nonTerminalStates := []mesos.TaskState{
-			mesos.TaskState_TASK_RUNNING,
-			mesos.TaskState_TASK_STAGING,
-			mesos.TaskState_TASK_STARTING,
+		nonTerminalStates := []TaskState{
+			TaskState_TASK_RUNNING,
+			TaskState_TASK_STAGING,
+			TaskState_TASK_STARTING,
 		}
 
 		Convey("IsTerminal", func() {
 			for _, state := range terminalStates {
-				test := fmt.Sprintf("Should be true for %s", state.String())
+				test := fmt.Sprintf("Should be true for %s", state)
 				Convey(test, func() {
-					So(IsTerminal(state.String()), ShouldBeTrue)
+					So(IsTerminal(state), ShouldBeTrue)
 				})
 			}
 
 			for _, state := range nonTerminalStates {
-				test := fmt.Sprintf("Should be false for %s", state.String())
+				test := fmt.Sprintf("Should be false for %s", state)
 				Convey(test, func() {
-					So(IsTerminal(state.String()), ShouldBeFalse)
+					So(IsTerminal(state), ShouldBeFalse)
 				})
 			}
 		})