Skip to content

Commit 2ed229d

Browse files
authored
Merge pull request #3 from Azure/haitao/dependon-fix
dependon fix
2 parents 011f502 + ee6ddb5 commit 2ed229d

File tree

2 files changed

+74
-33
lines changed

2 files changed

+74
-33
lines changed

job.go

+59-27
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,7 @@ func NewJob(name string) *Job {
4242
}
4343

4444
j.rootJob = &StepInfo[interface{}]{
45-
name: "[Start]",
46-
dependOn: []string{},
45+
name: "[Start]",
4746
task: asynctask.Start(context.Background(), func(fctx context.Context) (*interface{}, error) {
4847
fmt.Println("RootJob Added")
4948
// this will pause all steps from starting, until Start() method is called.
@@ -53,14 +52,14 @@ func NewJob(name string) *Job {
5352
}),
5453
}
5554

56-
j.Steps[j.rootJob.Name()] = j.rootJob
57-
j.stepsDag.Add(j.rootJob.Name())
55+
j.Steps[j.rootJob.GetName()] = j.rootJob
56+
j.stepsDag.Add(j.rootJob.GetName())
5857

5958
return j
6059
}
6160

6261
func InputParam[T any](bCtx context.Context, j *Job, stepName string, value *T) *StepInfo[T] {
63-
step := newStepInfo[T](stepName, []string{j.rootJob.Name()})
62+
step := newStepInfo[T](stepName)
6463

6564
instrumentedFunc := func(ctx context.Context) (*T, error) {
6665
j.rootJob.Wait(ctx)
@@ -69,15 +68,18 @@ func InputParam[T any](bCtx context.Context, j *Job, stepName string, value *T)
6968
step.task = asynctask.Start(bCtx, instrumentedFunc)
7069

7170
j.Steps[stepName] = step
72-
j.registerStepInGraph(stepName, j.rootJob.Name())
71+
j.registerStepInGraph(stepName, j.rootJob.GetName())
7372

7473
return step
7574
}
7675

77-
func AddStep[T any](bCtx context.Context, j *Job, stepName string, stepFunc asynctask.AsyncFunc[T], dependOn []string, optionDecorators ...ExecutionOptionPreparer) (*StepInfo[T], error) {
78-
// manually specified the dependencies, without consume the result.
76+
func AddStep[T any](bCtx context.Context, j *Job, stepName string, stepFunc asynctask.AsyncFunc[T], optionDecorators ...ExecutionOptionPreparer) (*StepInfo[T], error) {
77+
step := newStepInfo[T](stepName, optionDecorators...)
78+
79+
// also consider specified the dependencies from ExecutionOptionPreparer, without consume the result.
80+
var precedingStepNames = step.DependsOn()
7981
var precedingTasks []asynctask.Waitable
80-
for _, stepName := range dependOn {
82+
for _, stepName := range precedingStepNames {
8183
if step, ok := j.Steps[stepName]; ok {
8284
precedingTasks = append(precedingTasks, step.Waitable())
8385
} else {
@@ -87,11 +89,10 @@ func AddStep[T any](bCtx context.Context, j *Job, stepName string, stepFunc asyn
8789

8890
// if a step have no preceding tasks, link it to our rootJob as preceding task, so it won't start yet.
8991
if len(precedingTasks) == 0 {
92+
precedingStepNames = append(precedingStepNames, j.rootJob.GetName())
9093
precedingTasks = append(precedingTasks, j.rootJob.Waitable())
9194
}
9295

93-
step := newStepInfo[T](stepName, dependOn, optionDecorators...)
94-
9596
// instrument to :
9697
// replaceRuntimeContext,
9798
// trackStepState
@@ -115,30 +116,45 @@ func AddStep[T any](bCtx context.Context, j *Job, stepName string, stepFunc asyn
115116
step.task = asynctask.Start(bCtx, instrumentedFunc)
116117

117118
j.Steps[stepName] = step
118-
if len(dependOn) > 0 {
119-
j.registerStepInGraph(stepName, dependOn...)
120-
} else {
121-
j.registerStepInGraph(stepName, j.rootJob.Name())
122-
}
119+
j.registerStepInGraph(stepName, precedingStepNames...)
123120

124121
return step, nil
125122
}
126123

127124
func StepAfter[T, S any](bCtx context.Context, j *Job, stepName string, parentStep *StepInfo[T], stepFunc asynctask.ContinueFunc[T, S], optionDecorators ...ExecutionOptionPreparer) (*StepInfo[S], error) {
128125
// check parentStepT is in this job
129-
if get, ok := j.Steps[parentStep.Name()]; !ok || get != parentStep {
130-
return nil, fmt.Errorf("step [%s] not found in job", parentStep.Name())
126+
if get, ok := j.Steps[parentStep.GetName()]; !ok || get != parentStep {
127+
return nil, fmt.Errorf("step [%s] not found in job", parentStep.GetName())
128+
}
129+
130+
step := newStepInfo[S](stepName, append(optionDecorators, ExecuteAfter(parentStep))...)
131+
132+
// also consider specified the dependencies from ExecutionOptionPreparer, without consume the result.
133+
var precedingStepNames = step.DependsOn()
134+
var precedingTasks []asynctask.Waitable
135+
for _, stepName := range precedingStepNames {
136+
if step, ok := j.Steps[stepName]; ok {
137+
precedingTasks = append(precedingTasks, step.Waitable())
138+
} else {
139+
return nil, fmt.Errorf("step [%s] not found", stepName)
140+
}
131141
}
132142

133-
step := newStepInfo[S](stepName, []string{parentStep.Name()}, optionDecorators...)
143+
// if a step have no preceding tasks, link it to our rootJob as preceding task, so it won't start yet.
144+
if len(precedingTasks) == 0 {
145+
precedingTasks = append(precedingTasks, j.rootJob.Waitable())
146+
}
134147

135148
// instrument to :
136149
// replaceRuntimeContext
137150
// trackStepState
138151
// retryHandling (TODO)
139152
// errorHandling (TODO)
140153
// timeoutHandling (TODO)
141-
instrumentedFunc := func(_ context.Context, t *T) (*S, error) {
154+
instrumentedFunc := func(ctx context.Context, t *T) (*S, error) {
155+
if err := asynctask.WaitAll(ctx, &asynctask.WaitAllOptions{}, precedingTasks...); err != nil {
156+
return nil, err
157+
}
142158
step.state = StepStateRunning
143159
result, err := stepFunc(j.runtimeCtx, t)
144160
if err != nil {
@@ -152,7 +168,7 @@ func StepAfter[T, S any](bCtx context.Context, j *Job, stepName string, parentSt
152168
step.task = asynctask.ContinueWith(bCtx, parentStep.task, instrumentedFunc)
153169

154170
j.Steps[stepName] = step
155-
j.registerStepInGraph(stepName, parentStep.Name())
171+
j.registerStepInGraph(stepName, precedingStepNames...)
156172
if err := j.stepsDag.Validate(); err != nil {
157173
return nil, fmt.Errorf("cycle dependency detected: %s", err)
158174
}
@@ -161,15 +177,30 @@ func StepAfter[T, S any](bCtx context.Context, j *Job, stepName string, parentSt
161177

162178
func StepAfterBoth[T, S, R any](bCtx context.Context, j *Job, stepName string, parentStepT *StepInfo[T], parentStepS *StepInfo[S], stepFunc asynctask.AfterBothFunc[T, S, R], optionDecorators ...ExecutionOptionPreparer) (*StepInfo[R], error) {
163179
// check parentStepT is in this job
164-
if get, ok := j.Steps[parentStepT.Name()]; !ok || get != parentStepT {
165-
return nil, fmt.Errorf("step [%s] not found in job", parentStepT.Name())
180+
if get, ok := j.Steps[parentStepT.GetName()]; !ok || get != parentStepT {
181+
return nil, fmt.Errorf("step [%s] not found in job", parentStepT.GetName())
166182
}
167-
if get, ok := j.Steps[parentStepS.Name()]; !ok || get != parentStepS {
168-
return nil, fmt.Errorf("step [%s] not found in job", parentStepS.Name())
183+
if get, ok := j.Steps[parentStepS.GetName()]; !ok || get != parentStepS {
184+
return nil, fmt.Errorf("step [%s] not found in job", parentStepS.GetName())
169185
}
170186

171-
step := newStepInfo[R](stepName, []string{parentStepT.Name(), parentStepS.Name()}, optionDecorators...)
187+
step := newStepInfo[R](stepName, append(optionDecorators, ExecuteAfter(parentStepT), ExecuteAfter(parentStepS))...)
188+
189+
// also consider specified the dependencies from ExecutionOptionPreparer, without consume the result.
190+
var precedingStepNames = step.DependsOn()
191+
var precedingTasks []asynctask.Waitable
192+
for _, stepName := range precedingStepNames {
193+
if step, ok := j.Steps[stepName]; ok {
194+
precedingTasks = append(precedingTasks, step.Waitable())
195+
} else {
196+
return nil, fmt.Errorf("step [%s] not found", stepName)
197+
}
198+
}
172199

200+
// if a step have no preceding tasks, link it to our rootJob as preceding task, so it won't start yet.
201+
if len(precedingTasks) == 0 {
202+
precedingTasks = append(precedingTasks, j.rootJob.Waitable())
203+
}
173204
// instrument to :
174205
// replaceRuntimeContext
175206
// trackStepState
@@ -190,7 +221,8 @@ func StepAfterBoth[T, S, R any](bCtx context.Context, j *Job, stepName string, p
190221
step.task = asynctask.AfterBoth(bCtx, parentStepT.task, parentStepS.task, instrumentedFunc)
191222

192223
j.Steps[stepName] = step
193-
j.registerStepInGraph(stepName, parentStepT.Name(), parentStepS.Name())
224+
j.registerStepInGraph(stepName, precedingStepNames...)
225+
194226
return step, nil
195227
}
196228

step.go

+15-6
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@ type StepExecutionOptions struct {
1818
Timeout time.Duration
1919
ErrorPolicy StepErrorPolicy
2020
RetryPolicy StepRetryPolicy
21+
22+
// dependencies that are not input.
23+
DependOn []string
2124
}
2225

2326
type StepErrorPolicy struct{}
@@ -26,8 +29,15 @@ type StepRetryPolicy struct{}
2629

2730
type ExecutionOptionPreparer func(*StepExecutionOptions) *StepExecutionOptions
2831

32+
func ExecuteAfter(step StepMeta) ExecutionOptionPreparer {
33+
return func(options *StepExecutionOptions) *StepExecutionOptions {
34+
options.DependOn = append(options.DependOn, step.GetName())
35+
return options
36+
}
37+
}
38+
2939
type StepMeta interface {
30-
Name() string
40+
GetName() string
3141
GetState() StepState
3242
DependsOn() []string
3343
Wait(context.Context) error
@@ -39,14 +49,13 @@ type StepInfo[T any] struct {
3949
name string
4050
task *asynctask.Task[T]
4151
state StepState
42-
dependOn []string
4352
executionOptions *StepExecutionOptions
53+
job *Job
4454
}
4555

46-
func newStepInfo[T any](stepName string, dependOn []string, optionDecorators ...ExecutionOptionPreparer) *StepInfo[T] {
56+
func newStepInfo[T any](stepName string, optionDecorators ...ExecutionOptionPreparer) *StepInfo[T] {
4757
step := &StepInfo[T]{
4858
name: stepName,
49-
dependOn: dependOn,
5059
state: StepStatePending,
5160
executionOptions: &StepExecutionOptions{},
5261
}
@@ -61,7 +70,7 @@ func newStepInfo[T any](stepName string, dependOn []string, optionDecorators ...
6170
// compiler check
6271
var _ StepMeta = &StepInfo[string]{}
6372

64-
func (si *StepInfo[T]) Name() string {
73+
func (si *StepInfo[T]) GetName() string {
6574
return si.name
6675
}
6776

@@ -70,7 +79,7 @@ func (si *StepInfo[T]) GetState() StepState {
7079
}
7180

7281
func (si *StepInfo[T]) DependsOn() []string {
73-
return si.dependOn
82+
return si.executionOptions.DependOn
7483
}
7584

7685
func (si *StepInfo[T]) Wait(ctx context.Context) error {

0 commit comments

Comments
 (0)