Merge pull request #633 from mvo5/feature/task-wait-for-wait-tasks

state: add task.{WaitFor,WaitTasks}

2 files changed, 98 insertions, 24 deletions

diff --git a/overlord/state/task.go b/overlord/state/task.go
index 9ab485fa4f..83a6e75812 100644
--- a/overlord/state/task.go
+++ b/overlord/state/task.go
@@ -33,44 +33,48 @@ type progress struct {
 //
 // See Change for more details.
 type Task struct {
-	state *State
-	id string
-	kind string
-	summary string
-	status Status
-	progress progress
-	data customData
+	state *State
+	id string
+	kind string
+	summary string
+	status Status
+	progress progress
+	data customData
+	waitTasks taskIDsSet
 }
 
 func newTask(state *State, id, kind, summary string) *Task {
	return &Task{
-	state: state,
-	id: id,
-	kind: kind,
-	summary: summary,
-	data: make(customData),
+	state: state,
+	id: id,
+	kind: kind,
+	summary: summary,
+	data: make(customData),
+	waitTasks: make(taskIDsSet),
	}
 }
 
 type marshalledTask struct {
-	ID string `json:"id"`
-	Kind string `json:"kind"`
-	Summary string `json:"summary"`
-	Status Status `json:"status"`
-	Progress progress `json:"progress"`
-	Data map[string]*json.RawMessage `json:"data"`
+	ID string `json:"id"`
+	Kind string `json:"kind"`
+	Summary string `json:"summary"`
+	Status Status `json:"status"`
+	Progress progress `json:"progress"`
+	Data map[string]*json.RawMessage `json:"data"`
+	WaitTasks taskIDsSet `json:"wait-tasks"`
 }
 
 // MarshalJSON makes Task a json.Marshaller
 func (t *Task) MarshalJSON() ([]byte, error) {
	t.state.ensureLocked()
	return json.Marshal(marshalledTask{
-	ID: t.id,
-	Kind: t.kind,
-	Summary: t.summary,
-	Status: t.status,
-	Progress: t.progress,
-	Data: t.data,
+	ID: t.id,
+	Kind: t.kind,
+	Summary: t.summary,
+	Status: t.status,
+	Progress: t.progress,
+	Data: t.data,
+	WaitTasks: t.waitTasks,
	})
 }
 
@@ -90,6 +94,7 @@ func (t *Task) UnmarshalJSON(data []byte) error {
	t.status = unmarshalled.Status
	t.progress = unmarshalled.Progress
	t.data = unmarshalled.Data
+	t.waitTasks = unmarshalled.WaitTasks
	return nil
 }
 
@@ -164,3 +169,17 @@ func (t *Task) Get(key string, value interface{}) error {
	t.state.ensureLocked()
	return t.data.get(key, value)
 }
+
+// WaitFor registers another task as a requirement for t to make progress
+// and sets the status as WaitingStatus.
+func (t *Task) WaitFor(another *Task) {
+	t.state.ensureLocked()
+	t.status = WaitingStatus
+	t.waitTasks.add(another.ID())
+}
+
+// WaitTasks returns the list of tasks registered for t to wait for.
+func (t *Task) WaitTasks() []*Task {
+	t.state.ensureLocked()
+	return t.waitTasks.tasks(t.state)
+}
diff --git a/overlord/state/task_test.go b/overlord/state/task_test.go
index 06b4acb48f..509ee0e7b7 100644
--- a/overlord/state/task_test.go
+++ b/overlord/state/task_test.go
@@ -20,9 +20,12 @@
 package state_test
 
 import (
+	"fmt"
+
	. "gopkg.in/check.v1"
 
	"github.com/ubuntu-core/snappy/overlord/state"
+	"github.com/ubuntu-core/snappy/testutil"
 )
 
 type taskSuite struct{}
@@ -192,3 +195,55 @@ func (ts *taskSuite) TestState(c *C) {
 
	c.Assert(t.State(), Equals, st)
 }
+
+func (ts *taskSuite) TestTaskMarshalsWaitFor(c *C) {
+	st := state.New(nil)
+	st.Lock()
+	defer st.Unlock()
+
+	chg := st.NewChange("install", "...")
+	t1 := chg.NewTask("download", "1...")
+	t2 := chg.NewTask("install", "2...")
+	t2.WaitFor(t1)
+
+	d, err := t2.MarshalJSON()
+	c.Assert(err, IsNil)
+
+	needle := fmt.Sprintf(`"wait-tasks":["%s"`, t1.ID())
+	c.Assert(string(d), testutil.Contains, needle)
+}
+
+func (ts *taskSuite) TestTaskWaitFor(c *C) {
+	st := state.New(nil)
+	st.Lock()
+	defer st.Unlock()
+
+	chg := st.NewChange("install", "...")
+	t1 := chg.NewTask("download", "1...")
+	t2 := chg.NewTask("install", "2...")
+	t2.WaitFor(t1)
+
+	c.Assert(t2.WaitTasks(), DeepEquals, []*state.Task{t1})
+	c.Assert(t2.Status(), Equals, state.WaitingStatus)
+}
+
+func (cs *taskSuite) TestWaitForNeedsLocked(c *C) {
+	st := state.New(nil)
+	st.Lock()
+	chg := st.NewChange("install", "...")
+	t1 := chg.NewTask("download", "1...")
+	t2 := chg.NewTask("install", "2...")
+	st.Unlock()
+
+	c.Assert(func() { t2.WaitFor(t1) }, PanicMatches, "internal error: accessing state without lock")
+}
+
+func (cs *taskSuite) TestWaitTasksNeedsLocked(c *C) {
+	st := state.New(nil)
+	st.Lock()
+	chg := st.NewChange("install", "...")
+	t := chg.NewTask("download", "1...")
+	st.Unlock()
+
+	c.Assert(func() { t.WaitTasks() }, PanicMatches, "internal error: accessing state without lock")
+}