Skip to content

Commit

Permalink
internal/workflow: add expansions
Browse files Browse the repository at this point in the history 
One of the upcoming use cases for relui, tagging x/ repos, has a
dependency graph we can't know without doing significant work. This is
in conflict with one of the workflow package's design principles, that
workflows be defined ahead of time. So here's one idea for bending that
rule.

Introduce the concept of an expansion task. An expansion task runs the
same as a normal task, but in addition to its inputs it receives (a copy
of) the workflow definition and can add tasks and outputs to it as
usual. Once the task finishes running, the workflow swaps in the new
definition and continues. Other tasks can run concurrently.

One significant difference is that the expansion task must run every
time the workflow is started or resumed, so they really should be pure
functions with no network access or side effects. To encourage this,
they don't receive a Context.

Introducing this required some refactoring of the existing Run loop; we
need to allow these apparently-unused tasks to run, and deal with new
tasks popping up mid-run. But overall I think it's not too huge a
distortion to the design.

For golang/go#48523.

Change-Id: I92c164883e06474fa951abba5fa18fd78eacb8ce
Reviewed-on: https://go-review.googlesource.com/c/build/+/423037
Run-TryBot: Heschi Kreinick <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>
Reviewed-by: Jenny Rakoczy <[email protected]>
  • Loading branch information
@heschi
heschi committed Aug 25, 2022
1 parent 7cf3efd commit bf01fcc
Show file tree
Hide file tree
Showing 2 changed files with 285 additions and 62 deletions.
240 changes: 185 additions & 55 deletions internal/workflow/workflow.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,7 @@
//
// Workflows are a set of tasks and actions that produce and consume Values.
// Tasks don't run until the workflow is started, so Values represent data that
// doesn't exist yet, and can't be used directly. Each value has a dynamic type,
// which must match its uses.
// doesn't exist yet, and can't be used directly.
//
// To wrap an existing Go object in a Value, use Const. To define a
// parameter that will be set when the workflow is started, use Param.
Expand All @@ -32,6 +31,13 @@
// definitions to create an ordering dependency that doesn't correspond to a
// function argument.
//
// Expansions are a third type of function that adds to a running workflow
// definition rather than producing an output. Unlike Actions and Tasks, they
// execute multiple times and must produce exactly the same workflow
// modifications each time. As such, they should be pure functions of their
// inputs. Producing different modifications, or running multiple expansions
// concurrently, is an error that will corrupt the workflow's state.
//
// Once a Definition is complete, call Start to set its parameters and
// instantiate it into a Workflow. Call Run to execute the workflow until
// completion.
Expand Down Expand Up @@ -75,6 +81,17 @@ func (d *Definition) name(name string) string {
return d.namePrefix + name
}

func (d *Definition) shallowClone() *Definition {
clone := New()
for k, v := range d.tasks {
clone.tasks[k] = v
}
for k, v := range d.outputs {
clone.outputs[k] = v
}
return clone
}

type definitionState struct {
parameters []MetaParameter // Ordered according to registration, unique parameter names.
tasks map[string]*taskDefinition
Expand Down Expand Up @@ -324,7 +341,7 @@ func Task5[C context.Context, I1, I2, I3, I4, I5, O1 any](d *Definition, name st
return addTask[O1](d, name, f, []metaValue{i1, i2, i3, i4, i5}, opts)
}

func addTask[O1 any](d *Definition, name string, f interface{}, inputs []metaValue, opts []TaskOption) *taskResult[O1] {
func addFunc(d *Definition, name string, f interface{}, inputs []metaValue, opts []TaskOption) *taskDefinition {
name = d.name(name)
td := &taskDefinition{name: name, f: f, args: inputs}
for _, input := range inputs {
Expand All @@ -334,12 +351,22 @@ func addTask[O1 any](d *Definition, name string, f interface{}, inputs []metaVal
td.deps = append(td.deps, opt.(*after).deps...)
}
d.tasks[name] = td
return td
}

func addTask[O1 any](d *Definition, name string, f interface{}, inputs []metaValue, opts []TaskOption) *taskResult[O1] {
td := addFunc(d, name, f, inputs, opts)
return &taskResult[O1]{td}
}

func addAction(d *Definition, name string, f interface{}, inputs []metaValue, opts []TaskOption) *dependency {
tr := addTask[interface{}](d, name, f, inputs, opts)
return &dependency{tr.task}
td := addFunc(d, name, f, inputs, opts)
return &dependency{td}
}

func addExpansion(d *Definition, name string, f interface{}, inputs []metaValue, opts []TaskOption) {
td := addFunc(d, name, f, inputs, opts)
td.isExpansion = true
}

// ActionN adds an Action to the workflow definition. Its behavior and
Expand Down Expand Up @@ -377,6 +404,39 @@ func (d *dependency) dependencies() []*taskDefinition {
return []*taskDefinition{d.task}
}

// ExpandN adds a workflow expansion task to the workflow definition.
// Expansion tasks run similarly to normal tasks, but instead of computing
// a result, they can add to the workflow definition.
//
// Unlike normal tasks, expansions may run multiple times and must produce
// the exact same changes to the definition each time.
//
// Running more than one expansion concurrently is an error and will corrupt
// the workflow.
func Expand0(d *Definition, name string, f func(*Definition) error, opts ...TaskOption) {
addExpansion(d, name, f, nil, opts)
}

func Expand1[I1 any](d *Definition, name string, f func(*Definition, I1) error, i1 Value[I1], opts ...TaskOption) {
addExpansion(d, name, f, []metaValue{i1}, opts)
}

func Expand2[I1, I2 any](d *Definition, name string, f func(*Definition, I1, I2) error, i1 Value[I1], i2 Value[I2], opts ...TaskOption) {
addExpansion(d, name, f, []metaValue{i1, i2}, opts)
}

func Expand3[I1, I2, I3 any](d *Definition, name string, f func(*Definition, I1, I2, I3) error, i1 Value[I1], i2 Value[I2], i3 Value[I3], opts ...TaskOption) {
addExpansion(d, name, f, []metaValue{i1, i2, i3}, opts)
}

func Expand4[I1, I2, I3, I4 any](d *Definition, name string, f func(*Definition, I1, I2, I3, I4) error, i1 Value[I1], i2 Value[I2], i3 Value[I3], i4 Value[I4], opts ...TaskOption) {
addExpansion(d, name, f, []metaValue{i1, i2, i3, i4}, opts)
}

func Expand5[I1, I2, I3, I4, I5 any](d *Definition, name string, f func(*Definition, I1, I2, I3, I4, I5) error, i1 Value[I1], i2 Value[I2], i3 Value[I3], i4 Value[I4], i5 Value[I5], opts ...TaskOption) {
addExpansion(d, name, f, []metaValue{i1, i2, i3, i4, i5}, opts)
}

// A TaskContext is a context.Context, plus workflow-related features.
type TaskContext struct {
disableRetries bool
Expand Down Expand Up @@ -439,10 +499,11 @@ type Logger interface {
}

type taskDefinition struct {
name string
args []metaValue
deps []*taskDefinition
f interface{}
name string
isExpansion bool
args []metaValue
deps []*taskDefinition
f interface{}
}

type taskResult[T any] struct {
Expand Down Expand Up @@ -477,16 +538,25 @@ type Workflow struct {

def *Definition
tasks map[*taskDefinition]*taskState
// pendingStates stores states that haven't been loaded because their
// tasks didn't exist at Resume time.
pendingStates map[string]*TaskState
}

type taskState struct {
def *taskDefinition
started bool
finished bool
def *taskDefinition
created bool
started bool
finished bool
err error

// normal tasks
result interface{}
serializedResult []byte
err error
retryCount int

// workflow expansion
expanded *Definition
}

func (t *taskState) toExported() *TaskState {
Expand Down Expand Up @@ -534,7 +604,7 @@ func (w *Workflow) validate() error {
used[output] = true
}
for _, task := range w.def.tasks {
if !used[task] {
if !used[task] && !task.isExpansion {
return fmt.Errorf("task %v is not referenced and should be deleted", task.name)
}
}
Expand Down Expand Up @@ -568,39 +638,55 @@ func (w *Workflow) validate() error {
func Resume(def *Definition, state *WorkflowState, taskStates map[string]*TaskState) (*Workflow, error) {
w := &Workflow{
ID: state.ID,
def: def,
params: state.Params,
tasks: map[*taskDefinition]*taskState{},
retryCommands: make(chan retryCommand, len(def.tasks)),
def: def,
tasks: map[*taskDefinition]*taskState{},
pendingStates: taskStates,
}
if err := w.validate(); err != nil {
return nil, err
}
for _, taskDef := range def.tasks {
tState, ok := taskStates[taskDef.name]
if !ok {
return nil, fmt.Errorf("task state for %q not found", taskDef.name)
var err error
w.tasks[taskDef], err = loadTaskState(w.pendingStates, taskDef, false)
if err != nil {
return nil, fmt.Errorf("loading state for %v: %v", taskDef.name, err)
}
state := &taskState{
def: taskDef,
started: tState.Finished, // Can't resume tasks, so either it's new or done.
finished: tState.Finished,
serializedResult: tState.SerializedResult,
retryCount: tState.RetryCount,
}
if state.serializedResult != nil {
result, err := unmarshalNew(reflect.ValueOf(taskDef.f).Type().Out(0), tState.SerializedResult)
if err != nil {
return nil, fmt.Errorf("failed to unmarshal result of %v: %v", taskDef.name, err)
}
state.result = result
}
return w, nil
}

func loadTaskState(states map[string]*TaskState, def *taskDefinition, allowMissing bool) (*taskState, error) {
tState, ok := states[def.name]
if !ok {
if !allowMissing {
return nil, fmt.Errorf("task state not found")
}
if tState.Error != "" {
state.err = fmt.Errorf("serialized error: %v", tState.Error) // untyped, but hopefully that doesn't matter.
tState = &TaskState{}
}
// Can't resume tasks, so either it's new or done.
// Expansions need to run every time.
finished := tState.Finished && !def.isExpansion
state := &taskState{
def: def,
created: ok,
started: finished,
finished: finished,
serializedResult: tState.SerializedResult,
retryCount: tState.RetryCount,
}
if state.serializedResult != nil {
result, err := unmarshalNew(reflect.ValueOf(def.f).Type().Out(0), tState.SerializedResult)
if err != nil {
return nil, fmt.Errorf("failed to unmarshal result: %v", err)
}
w.tasks[taskDef] = state
state.result = result
}
return w, nil
if tState.Error != "" {
state.err = fmt.Errorf("serialized error: %v", tState.Error) // untyped, but hopefully that doesn't matter.
}
return state, nil
}

func unmarshalNew(t reflect.Type, data []byte) (interface{}, error) {
Expand All @@ -620,29 +706,25 @@ func (w *Workflow) Run(ctx context.Context, listener Listener) (map[string]inter
listener = &defaultListener{}
}

for _, task := range w.tasks {
listener.TaskStateChanged(w.ID, task.def.name, task.toExported())
}

stateChan := make(chan taskState, 2*len(w.def.tasks))
doneOnce := ctx.Done()
for {
// If we have all the outputs, the workflow is done.
outValues := map[string]interface{}{}
for outName, outDef := range w.def.outputs {
if task := w.tasks[outDef]; task.finished && task.err == nil {
outValues[outName] = task.result
}
}
if len(outValues) == len(w.def.outputs) {
return outValues, nil
}

running := 0
allDone := true
for _, task := range w.tasks {
if !task.created {
task.created = true
listener.TaskStateChanged(w.ID, task.def.name, task.toExported())
}
if task.started && !task.finished {
running++
}
if !task.finished || task.err != nil {
allDone = false
}
}
if allDone {
break
}

if ctx.Err() == nil {
Expand All @@ -658,9 +740,13 @@ func (w *Workflow) Run(ctx context.Context, listener Listener) (map[string]inter
task.started = true
running++
listener.TaskStateChanged(w.ID, task.def.name, task.toExported())
go func(task taskState) {
stateChan <- runTask(ctx, w.ID, listener, task, args)
}(*task)
taskCopy := *task
if task.def.isExpansion {
defCopy := w.def.shallowClone()
go func() { stateChan <- runExpansion(defCopy, taskCopy, args) }()
} else {
go func() { stateChan <- runTask(ctx, w.ID, listener, taskCopy, args) }()
}
}
}

Expand All @@ -675,6 +761,9 @@ func (w *Workflow) Run(ctx context.Context, listener Listener) (map[string]inter

select {
case state := <-stateChan:
if state.def.isExpansion && state.finished && state.err == nil {
state.err = w.expand(state.expanded)
}
listener.TaskStateChanged(w.ID, state.def.name, state.toExported())
w.tasks[state.def] = &state
case retry := <-w.retryCommands:
Expand All @@ -689,14 +778,20 @@ func (w *Workflow) Run(ctx context.Context, listener Listener) (map[string]inter
break
}
listener.Logger(w.ID, def.name).Printf("Manual retry requested")
stateChan <- taskState{def: def}
stateChan <- taskState{def: def, created: true}
retry.reply <- nil
// Don't get stuck when cancellation comes in after all tasks have
// finished, but also don't busy wait if something's still running.
case <-doneOnce:
doneOnce = nil
}
}

outs := map[string]interface{}{}
for name, def := range w.def.outputs {
outs[name] = w.tasks[def].result
}
return outs, nil
}

func (w *Workflow) taskArgs(def *taskDefinition) ([]reflect.Value, bool) {
Expand Down Expand Up @@ -753,12 +848,47 @@ func runTask(ctx context.Context, workflowID uuid.UUID, listener Listener, state
tctx.Printf("task failed, will retry (%v of %v): %v", state.retryCount+1, MaxRetries, state.err)
state = taskState{
def: state.def,
created: true,
retryCount: state.retryCount + 1,
}
}
return state
}

func runExpansion(d *Definition, state taskState, args []reflect.Value) taskState {
in := append([]reflect.Value{reflect.ValueOf(d)}, args...)
fv := reflect.ValueOf(state.def.f)
out := fv.Call(in)
state.finished = true
if out[0].IsNil() {
state.expanded = d
} else {
state.err = out[0].Interface().(error)
}
return state
}

func (w *Workflow) expand(expanded *Definition) error {
origDef := w.def
w.def = expanded
if err := w.validate(); err != nil {
w.def = origDef
return err
}
for _, def := range w.def.tasks {
if _, ok := w.tasks[def]; ok {
continue
}
// w.tasks[def] = &taskState{def: def}
var err error
w.tasks[def], err = loadTaskState(w.pendingStates, def, true)
if err != nil {
return err
}
}
return nil
}

type defaultListener struct{}

func (s *defaultListener) TaskStateChanged(_ uuid.UUID, _ string, _ *TaskState) error {
Expand Down
Loading

0 comments on commit bf01fcc

Please sign in to comment.