Refactors task execution logic

Author: grussorusso

internal/workflow/choice_task.go

@@ -32,18 +32,14 @@ func (c *ChoiceTask) GetAlternatives() []TaskId {
 	return c.AlternativeNextTasks
 }
 
-func (c *ChoiceTask) execute(progress *Progress, input *PartialData, r *Request) (*PartialData, *Progress, bool, error) {
-
-	outputData := NewPartialData(ReqId(r.Id), c.GetId(), nil) // partial initialization of outputData
-
-	// NOTE: we do not call task.CheckInput() as this task has no signature to match against
+func (c *ChoiceTask) Evaluate(input *PartialData, r *Request) (TaskId, error) {
 
 	// simply evaluate the Conditions and set the matching one
 	matchedCondition := -1
 	for i, condition := range c.Conditions {
 		ok, err := condition.Evaluate(input.Data)
 		if err != nil {
-			return nil, progress, false, fmt.Errorf("error while testing condition: %v", err)
+			return "", fmt.Errorf("error while testing condition: %v", err)
 		}
 		if ok {
 			matchedCondition = i
@@ -52,64 +48,10 @@ func (c *ChoiceTask) execute(progress *Progress, input *PartialData, r *Request)
 	}
 
 	if matchedCondition < 0 {
-		return nil, progress, false, fmt.Errorf("no condition is met")
-	}
-
-	nextTask := c.AlternativeNextTasks[matchedCondition]
-	outputData.ForTask = nextTask
-	outputData.Data = input.Data
-
-	// we skip all branch that will not be executed
-	tasksToSkip := c.GetTasksToSkip(r.W, matchedCondition)
-	for _, t := range tasksToSkip {
-		progress.Skip(t.GetId())
-	}
-
-	progress.Complete(c.GetId())
-	err := progress.AddReadyTask(nextTask)
-	if err != nil {
-		return nil, progress, false, err
+		return "", fmt.Errorf("no condition is met")
 	}
-	return outputData, progress, true, nil
-}
 
-// VisitBranch returns all tasks of a branch under a choice task; branch number starts from 0
-func (c *ChoiceTask) VisitBranch(workflow *Workflow, branch int) []Task {
-	branchTasks := make([]Task, 0)
-	if len(c.AlternativeNextTasks) <= branch {
-		fmt.Printf("fail to get branch %d\n", branch)
-		return branchTasks
-	}
-	taskId := c.AlternativeNextTasks[branch]
-	return Visit(workflow, taskId, true)
-}
-
-// GetTasksToSkip skips all tasks that are in a branch that will not be executed.
-// If a skipped branch contains one or more tasks that in use by the current branch, the task
-// should NOT be skipped (Tested in TestParsingChoiceDagWithDataTestExpr)
-func (c *ChoiceTask) GetTasksToSkip(workflow *Workflow, matchedCondition int) []Task {
-	toSkip := make([]Task, 0)
-
-	toNotSkip := c.VisitBranch(workflow, matchedCondition)
-	for i := 0; i < len(c.AlternativeNextTasks); i++ {
-		if i == matchedCondition {
-			continue
-		}
-		branchTasks := c.VisitBranch(workflow, i)
-		for _, task := range branchTasks {
-			shouldBeSkipped := true
-			for _, taskNotSkipped := range toNotSkip {
-				if task.Equals(taskNotSkipped) {
-					shouldBeSkipped = false
-					break
-				}
-			}
-			if shouldBeSkipped {
-				toSkip = append(toSkip, task)
-			}
-		}
-	}
-	return toSkip
+	return c.AlternativeNextTasks[matchedCondition], nil
 }
 
 func (c *ChoiceTask) String() string {

internal/workflow/end_task.go

@@ -17,11 +17,6 @@ func NewEndTask() *EndTask {
 	}
 }
 
-func (e *EndTask) execute(progress *Progress, partialData *PartialData) (*PartialData, *Progress, bool, error) {
-	progress.Complete(e.Id)
-	return partialData, progress, false, nil // false because we want to stop when reaching the end
-}
-
 func (e *EndTask) String() string {
 	return fmt.Sprintf("[EndTask]")
 }

internal/workflow/fail_task.go

@@ -21,18 +21,6 @@ func NewFailureTask(error, cause string) *FailureTask {
 	return &fail
 }
 
-func (f *FailureTask) execute(progress *Progress, r *Request) (*PartialData, *Progress, bool, error) {
-
-	output := make(map[string]interface{})
-	output[f.Error] = f.Cause
-	outputData := NewPartialData(ReqId(r.Id), f.GetNext(), output)
-
-	progress.Complete(f.GetId())
-
-	shouldContinueExecution := f.GetType() != Fail && f.GetType() != Succeed
-	return outputData, progress, shouldContinueExecution, nil
-}
-
 func (f *FailureTask) GetNext() TaskId {
 	return f.NextTask
 }
@@ -48,3 +36,9 @@ func (f *FailureTask) SetNext(nextTask Task) error {
 func (f *FailureTask) String() string {
 	return fmt.Sprintf("[Fail: %s]", f.Error)
 }
+
+func (f *FailureTask) execute(input *PartialData, r *Request) (map[string]interface{}, error) {
+	output := make(map[string]interface{})
+	output[f.Error] = f.Cause
+	return output, nil
+}

internal/workflow/function_task.go

@@ -35,26 +35,18 @@ func (s *FunctionTask) SetNext(nextTask Task) error {
 	return nil
 }
 
-func (s *FunctionTask) execute(progress *Progress, input *PartialData, r *Request) (*PartialData, *Progress, bool, error) {
+func (s *FunctionTask) execute(input *PartialData, r *Request) (map[string]interface{}, error) {
 
 	err := s.CheckInput(input.Data)
 	if err != nil {
-		return nil, progress, false, err
+		return nil, err
 	}
 	output, err := s.exec(r, input.Data)
 	if err != nil {
-		return nil, progress, false, err
+		return nil, err
 	}
 
-	nextTask := s.GetNext()
-	outputData := NewPartialData(ReqId(r.Id), nextTask, output)
-
-	progress.Complete(s.Id)
-	err = progress.AddReadyTask(nextTask)
-	if err != nil {
-		return nil, progress, false, err
-	}
-	return outputData, progress, true, nil
+	return output, nil
 }
 
 func (s *FunctionTask) exec(compRequest *Request, params ...map[string]interface{}) (map[string]interface{}, error) {

internal/workflow/pass_task.go

@@ -31,3 +31,7 @@ func (p *PassTask) SetNext(nextTask Task) error {
 func (p *PassTask) String() string {
 	return "[ Pass ]"
 }
+
+func (p *PassTask) execute(input *PartialData, r *Request) (map[string]interface{}, error) {
+	return input.Data, nil
+}

internal/workflow/progress.go

@@ -110,14 +110,6 @@ func (p *Progress) Complete(id TaskId) {
 	}
 }
 
-func (p *Progress) AddReadyTask(id TaskId) error {
-	if !p.IsReady(id) {
-		return fmt.Errorf("the task is not ready")
-	}
-	p.ReadyToExecute = append(p.ReadyToExecute, id)
-	return nil
-}
-
 // TODO: skip on cascade next nodes
 func (p *Progress) Skip(id TaskId) {
 	p.Status[id] = Skipped

internal/workflow/start_task.go

@@ -29,16 +29,8 @@ func (s *StartTask) SetNext(nextTask Task) error {
 	return nil
 }
 
-func (s *StartTask) execute(progress *Progress, partialData *PartialData) (*PartialData, *Progress, bool, error) {
-
-	// TODO: move this logic into workflow "handleCompletion(output)"
-	progress.Complete(s.GetId())
-
-	err := progress.AddReadyTask(s.GetNext())
-	if err != nil {
-		return nil, progress, false, err
-	}
-	return partialData, progress, true, nil
+func (s *StartTask) execute(input *PartialData, r *Request) (map[string]interface{}, error) {
+	return input.Data, nil
 }
 
 func (s *StartTask) String() string {

internal/workflow/succeed_task.go

@@ -31,3 +31,7 @@ func (s *SuccessTask) GetNext() TaskId {
 func (s *SuccessTask) String() string {
 	return "[Succeed]"
 }
+
+func (s *SuccessTask) execute(input *PartialData, r *Request) (map[string]interface{}, error) {
+	return input.Data, nil
+}

internal/workflow/task.go

@@ -25,12 +25,15 @@ type UnaryTask interface {
 
 	// SetNext connects the output of this task to another Task
 	SetNext(nextTask Task) error
+
+	execute(data *PartialData, r *Request) (map[string]interface{}, error)
 }
 
 type ConditionalTask interface {
 	Task
 	AddAlternative(nextTask Task) error
 	GetAlternatives() []TaskId
+	Evaluate(data *PartialData, r *Request) (TaskId, error)
 }
 
 type baseTask struct {

internal/workflow/workflow.go

@@ -169,28 +169,9 @@ func Visit(workflow *Workflow, taskId TaskId, excludeEnd bool) []Task {
 	return tasks
 }
 
-func (workflow *Workflow) doNothingExec(progress *Progress, input *PartialData, task UnaryTask, r *Request) (*PartialData, *Progress, bool, error) {
-
-	output := input.Data
-	outputData := NewPartialData(ReqId(r.Id), task.GetNext(), output)
-
-	progress.Complete(task.GetId())
-
-	shouldContinueExecution := task.GetType() != Fail && task.GetType() != Succeed
-	if shouldContinueExecution {
-		err := progress.AddReadyTask(task.GetNext())
-		if err != nil {
-			return nil, progress, false, nil
-		}
-	}
-
-	return outputData, progress, shouldContinueExecution, nil
-}
-
 func (workflow *Workflow) Execute(r *Request, input *PartialData, progress *Progress) (*PartialData, *Progress, bool, error) {
-	var output *PartialData
 	var err error
-	shouldContinue := true
+	var outputData *PartialData
 
 	var nextTasks []TaskId
 
@@ -213,24 +194,57 @@ func (workflow *Workflow) Execute(r *Request, input *PartialData, progress *Prog
 		}
 
 		switch task := n.(type) {
-		case *FunctionTask:
-			output, progress, shouldContinue, err = task.execute(progress, input, r)
-		case *ChoiceTask:
-			output, progress, shouldContinue, err = task.execute(progress, input, r)
-		case *StartTask:
-			output, progress, shouldContinue, err = task.execute(progress, input)
-		case *PassTask:
-			output, progress, shouldContinue, err = workflow.doNothingExec(progress, input, task, r)
-		case *FailureTask:
-			output, progress, shouldContinue, err = task.execute(progress, r)
-		case *SuccessTask:
-			output, progress, shouldContinue, err = workflow.doNothingExec(progress, input, task, r)
+		case UnaryTask:
+			output, err := task.execute(input, r)
+			if err != nil {
+				progress.Fail(n.GetId())
+				return nil, progress, false, err
+			}
+			progress.Complete(task.GetId())
+
+			nextTask := task.GetNext()
+			outputData = NewPartialData(ReqId(r.Id), nextTask, output)
+			if progress.IsReady(nextTask) {
+				progress.ReadyToExecute = append(progress.ReadyToExecute, nextTask)
+			}
+
+		case ConditionalTask:
+			nextTaskId, err := task.Evaluate(input, r)
+			if err != nil {
+				progress.Fail(n.GetId())
+				return nil, progress, false, err
+			}
+
+			// we skip all tasks that will not be executed
+			toSkip := make([]Task, 0)
+			toNotSkip := Visit(workflow, nextTaskId, false)
+			for _, a := range task.GetAlternatives() {
+				if a == nextTaskId {
+					continue
+				}
+				branchTasks := Visit(workflow, a, false)
+				for _, otherTask := range branchTasks {
+					if !slices.Contains(toNotSkip, otherTask) {
+						toSkip = append(toSkip, otherTask)
+					}
+				}
+			}
+			for _, t := range toSkip {
+				progress.Skip(t.GetId())
+			}
+			progress.Complete(task.GetId())
+
+			outputData = NewPartialData(ReqId(r.Id), nextTaskId, input.Data)
+			if progress.IsReady(nextTaskId) {
+				progress.ReadyToExecute = append(progress.ReadyToExecute, nextTaskId)
+			}
 		case *EndTask:
-			output, progress, shouldContinue, err = task.execute(progress, input)
+			progress.Complete(task.GetId())
+			outputData = input
 		}
 		if err != nil {
 			progress.Fail(n.GetId())
-			return output, progress, false, err
+			return nil, progress, false, err
 		}
 	} else {
 		err = SaveProgress(progress)
@@ -244,7 +258,7 @@ func (workflow *Workflow) Execute(r *Request, input *PartialData, progress *Prog
 		return nil, progress, false, nil
 	}
 
-	return output, progress, shouldContinue, nil
+	return outputData, progress, true, nil
 }
 
 // GetUniqueFunctions returns a list with the function names used in the Workflow. The returned function names are unique and in alphabetical order
@@ -440,7 +454,7 @@ func (workflow *Workflow) Invoke(r *Request) error {
 				return fmt.Errorf("failed workflow execution: %v", err)
 			}
 
-			if !shouldContinue && pd != nil {
+			if len(progress.ReadyToExecute) == 0 && pd != nil {
 				r.ExecReport.Result = pd.Data
 			}
 		}