Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 12 additions & 7 deletions core/scheduler/service/job_expectator_service.go
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ func (s *JobExpectatorService) GenerateExpectedFinishTimes(ctx context.Context,
continue
}
s.l.Debug("calculating expected finish time for job", "job", jobSchedule.JobName, "scheduled_at", jobSchedule.ScheduledAt)
err := s.PopulateExpectedFinishTime(jobSchedule, jobsWithLineageMap[jobSchedule.JobName], jobRunExpectedFinishTimeDetail, jobDurationsEstimation, referenceTime)
err := s.PopulateExpectedFinishTime(jobSchedule.JobName, jobsWithLineageMap[jobSchedule.JobName], jobRunExpectedFinishTimeDetail, jobDurationsEstimation, referenceTime)
if err != nil {
s.l.Error(fmt.Sprintf("failed to populate expected finish time for job [%s]: %s", jobSchedule.JobName, err.Error()))
return nil, err
Expand Down Expand Up @@ -150,9 +150,13 @@ func (s *JobExpectatorService) GenerateExpectedFinishTimes(ctx context.Context,
return finalJobRunExpectedFinishTimes, nil
}

func (s *JobExpectatorService) PopulateExpectedFinishTime(jobTarget *scheduler.JobSchedule, currentJobWithLineage *scheduler.JobLineageSummary, jobRunExpectedFinishTimes map[scheduler.JobSchedule]FinishTimeDetail, jobDurationsEstimation map[scheduler.JobName]*time.Duration, referenceTime time.Time) error {
// selfParent is the immediate downstream job that led to currentJobWithLineage in the current
// traversal - currentJobWithLineage's own run is keyed by that name in JobRuns, since
// LineageResolver.BuildLineage keys a node's JobRuns by its immediate downstream
// to support lineages where a shared upstream carries a distinct run per downstream path.
func (s *JobExpectatorService) PopulateExpectedFinishTime(selfParent scheduler.JobName, currentJobWithLineage *scheduler.JobLineageSummary, jobRunExpectedFinishTimes map[scheduler.JobSchedule]FinishTimeDetail, jobDurationsEstimation map[scheduler.JobName]*time.Duration, referenceTime time.Time) error {
// pre condition check
if currentJobWithLineage == nil || currentJobWithLineage.JobRuns[jobTarget.JobName] == nil {
if currentJobWithLineage == nil || currentJobWithLineage.GetRunForJob(selfParent) == nil {
// TODO: add metric to track how many times this happens
s.l.Error(fmt.Sprintf("[critical] no job run found for job [%s], skipping expected finish time calculation", currentJobWithLineage.JobName))
return nil
Expand All @@ -162,7 +166,7 @@ func (s *JobExpectatorService) PopulateExpectedFinishTime(jobTarget *scheduler.J
return nil
}

currentJobRun := currentJobWithLineage.JobRuns[jobTarget.JobName]
currentJobRun := currentJobWithLineage.GetRunForJob(selfParent)
currentJobScheduleKey := scheduler.JobSchedule{
// TODO: add project name as well, PR: https://github.com/goto/optimus/pull/501
JobName: currentJobWithLineage.JobName,
Expand Down Expand Up @@ -231,17 +235,18 @@ func (s *JobExpectatorService) PopulateExpectedFinishTime(jobTarget *scheduler.J
FinishTime: maxUpstreamExpectedFinishTime.Add(*estimatedDuration),
}
for _, upstream := range currentJobWithLineage.Upstreams {
err := s.PopulateExpectedFinishTime(jobTarget, upstream, jobRunExpectedFinishTimes, jobDurationsEstimation, referenceTime)
err := s.PopulateExpectedFinishTime(currentJobWithLineage.JobName, upstream, jobRunExpectedFinishTimes, jobDurationsEstimation, referenceTime)
if err != nil {
return err
}
if upstream.JobRuns[jobTarget.JobName] == nil {
upstreamJobRun := upstream.GetRunForJob(currentJobWithLineage.JobName)
if upstream.JobRuns[currentJobWithLineage.JobName] == nil {
s.l.Debug(fmt.Sprintf("no upstream job run found for job, skipping upstream in expected finish time calculation [job: %s, upstream_job: %s]", currentJobWithLineage.JobName, upstream.JobName))
continue
}
upstreamScheduleKey := scheduler.JobSchedule{
JobName: upstream.JobName,
ScheduledAt: upstream.JobRuns[jobTarget.JobName].ScheduledAt,
ScheduledAt: upstreamJobRun.ScheduledAt,
}
upstreamExpectedFinishTime, ok := jobRunExpectedFinishTimes[upstreamScheduleKey]
if !ok {
Expand Down
196 changes: 187 additions & 9 deletions core/scheduler/service/job_expectator_service_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -429,7 +429,7 @@ func TestPopulateExpectedFinishTime(t *testing.T) {
jobDurationEstimation[jobTarget.JobName] = func() *time.Duration { d := 30 * time.Minute; return &d }()

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)

// then
assert.NoError(t, err)
Expand Down Expand Up @@ -477,7 +477,7 @@ func TestPopulateExpectedFinishTime(t *testing.T) {
jobDurationEstimation[jobTarget.JobName] = func() *time.Duration { d := 30 * time.Minute; return &d }()

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)

// then
assert.NoError(t, err)
Expand Down Expand Up @@ -523,7 +523,7 @@ func TestPopulateExpectedFinishTime(t *testing.T) {
// no duration estimation added

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)

// then
assert.NoError(t, err)
Expand Down Expand Up @@ -575,7 +575,7 @@ func TestPopulateExpectedFinishTime(t *testing.T) {
}

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)
// then
assert.NoError(t, err)
// should not be updated
Expand Down Expand Up @@ -623,7 +623,7 @@ func TestPopulateExpectedFinishTime(t *testing.T) {
jobDurationEstimation[jobTarget.JobName] = func() *time.Duration { d := 30 * time.Minute; return &d }()

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)

// then
assert.NoError(t, err)
Expand Down Expand Up @@ -672,7 +672,7 @@ func TestPopulateExpectedFinishTime(t *testing.T) {
jobDurationEstimation[jobTarget.JobName] = func() *time.Duration { d := 30 * time.Minute; return &d }()

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)

// then
assert.NoError(t, err)
Expand Down Expand Up @@ -719,7 +719,7 @@ func TestPopulateExpectedFinishTime(t *testing.T) {
jobDurationEstimation[jobTarget.JobName] = func() *time.Duration { d := 30 * time.Minute; return &d }()

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)

// then
assert.NoError(t, err)
Expand Down Expand Up @@ -781,7 +781,7 @@ func TestPopulateExpectedFinishTime(t *testing.T) {
jobDurationEstimation[jobUpstreamWithLineage.JobName] = func() *time.Duration { d := 45 * time.Minute; return &d }()

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)

// then
assert.NoError(t, err)
Expand Down Expand Up @@ -848,13 +848,191 @@ func TestPopulateExpectedFinishTime(t *testing.T) {
jobDurationEstimation[jobUpstreamWithLineage.JobName] = func() *time.Duration { d := 45 * time.Minute; return &d }()

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)

// then
assert.NoError(t, err)
expectedExpectedFinishTime := referenceTime.Add(10 * time.Minute).Add(30 * time.Minute)
assert.Equal(t, expectedExpectedFinishTime, jobRunExpectedFinishTime[*jobTarget].FinishTime)
})

t.Run("when upstream chain is 2 levels deep, should propagate expected finish time through the full chain", func(t *testing.T) {
// given
jobRunExpectationDetailsRepo := NewJobRunExpectationDetailsRepository(t)
jobDetailsGetter := NewJobDetailsGetter(t)
jobLineageFetcher := NewJobLineageFetcher(t)
durationEstimator := NewDurationEstimator(t)

jobExpectatorService := service.NewJobExpectatorService(
l,
10,
jobRunExpectationDetailsRepo,
jobDetailsGetter,
jobLineageFetcher,
durationEstimator,
)
jobRunExpectedFinishTime := map[scheduler.JobSchedule]service.FinishTimeDetail{}
jobDurationEstimation := map[scheduler.JobName]*time.Duration{}

scheduledAtA := referenceTime.Add(10 * time.Minute)
scheduledAtB := referenceTime.Add(1 * time.Hour)
scheduledAtC := referenceTime.Add(-2 * time.Hour)
jobCEndTime := referenceTime.Add(90 * time.Minute) // job-C finished after job-B's own scheduled_at

jobTarget := &scheduler.JobSchedule{
JobName: scheduler.JobName("job-A"),
ScheduledAt: scheduledAtA,
}
jobBName := scheduler.JobName("job-B")
jobCName := scheduler.JobName("job-C")

// job-C's run is keyed by job-B (its immediate downstream), not by the root job-A,
// matching LineageResolver.BuildLineage's diamond-safe keying convention.
jobCWithLineage := &scheduler.JobLineageSummary{
JobName: jobCName,
IsEnabled: true,
JobRuns: map[scheduler.JobName]*scheduler.JobRunSummary{
jobBName: {
JobName: jobCName,
ScheduledAt: scheduledAtC,
JobEndTime: &jobCEndTime,
},
},
Upstreams: []*scheduler.JobLineageSummary{},
}
jobBWithLineage := &scheduler.JobLineageSummary{
JobName: jobBName,
IsEnabled: true,
JobRuns: map[scheduler.JobName]*scheduler.JobRunSummary{
jobTarget.JobName: {
JobName: jobBName,
ScheduledAt: scheduledAtB,
},
},
Upstreams: []*scheduler.JobLineageSummary{jobCWithLineage},
}
currentJobWithLineage := &scheduler.JobLineageSummary{
JobName: jobTarget.JobName,
IsEnabled: true,
JobRuns: map[scheduler.JobName]*scheduler.JobRunSummary{
jobTarget.JobName: {
JobName: jobTarget.JobName,
ScheduledAt: scheduledAtA,
},
},
Upstreams: []*scheduler.JobLineageSummary{jobBWithLineage},
}

jobDurationEstimation[jobTarget.JobName] = func() *time.Duration { d := 30 * time.Minute; return &d }()
jobDurationEstimation[jobBName] = func() *time.Duration { d := 45 * time.Minute; return &d }()

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)

// then
assert.NoError(t, err)
// job-C already finished; its finish time is its own end time
assert.Equal(t, jobCEndTime, jobRunExpectedFinishTime[scheduler.JobSchedule{JobName: jobCName, ScheduledAt: scheduledAtC}].FinishTime)
// job-B hasn't started; its expected finish time is max(its own scheduled_at, job-C's finish time) + job-B duration
expectedBFinish := jobCEndTime.Add(45 * time.Minute)
assert.Equal(t, expectedBFinish, jobRunExpectedFinishTime[scheduler.JobSchedule{JobName: jobBName, ScheduledAt: scheduledAtB}].FinishTime)
// job-A hasn't started; its expected finish time is max(its own scheduled_at, job-B's finish time) + job-A duration,
// proving job-C's contribution propagated two levels up through job-B
expectedAFinish := expectedBFinish.Add(30 * time.Minute)
assert.Equal(t, expectedAFinish, jobRunExpectedFinishTime[*jobTarget].FinishTime)
})

t.Run("when an upstream job occurs twice in the lineage, once shallow and once deep (diamond), should reuse the shared node's cached finish time instead of recomputing or conflicting", func(t *testing.T) {
// given
jobRunExpectationDetailsRepo := NewJobRunExpectationDetailsRepository(t)
jobDetailsGetter := NewJobDetailsGetter(t)
jobLineageFetcher := NewJobLineageFetcher(t)
durationEstimator := NewDurationEstimator(t)

jobExpectatorService := service.NewJobExpectatorService(
l,
10,
jobRunExpectationDetailsRepo,
jobDetailsGetter,
jobLineageFetcher,
durationEstimator,
)
jobRunExpectedFinishTime := map[scheduler.JobSchedule]service.FinishTimeDetail{}
jobDurationEstimation := map[scheduler.JobName]*time.Duration{}

scheduledAtA := referenceTime.Add(30 * time.Minute)
scheduledAtB := referenceTime.Add(1 * time.Hour)
scheduledAtX := referenceTime.Add(-3 * time.Hour)
jobXEndTime := referenceTime.Add(3 * time.Hour)

jobTarget := &scheduler.JobSchedule{
JobName: scheduler.JobName("job-A"),
ScheduledAt: scheduledAtA,
}
jobBName := scheduler.JobName("job-B")
jobXName := scheduler.JobName("job-X")

// job-X is a diamond: it is job-A's DIRECT upstream (shallow) and also job-B's upstream
// (deep, via job-A -> job-B -> job-X). LineageResolver.BuildLineage's buildLineageTree
// memoizes by job name, so both edges point at the very same *JobLineageSummary object,
// and calculateAllUpstreamRuns adds one JobRuns entry per immediate-downstream path -
// here both paths resolve to the same actual run, so both keys share the same run pointer.
jobXRun := &scheduler.JobRunSummary{
JobName: jobXName,
ScheduledAt: scheduledAtX,
JobEndTime: &jobXEndTime,
}
jobXWithLineage := &scheduler.JobLineageSummary{
JobName: jobXName,
IsEnabled: true,
JobRuns: map[scheduler.JobName]*scheduler.JobRunSummary{
jobTarget.JobName: jobXRun, // shallow path: job-A -> job-X
jobBName: jobXRun, // deep path: job-A -> job-B -> job-X (same run)
},
Upstreams: []*scheduler.JobLineageSummary{},
}
jobBWithLineage := &scheduler.JobLineageSummary{
JobName: jobBName,
IsEnabled: true,
JobRuns: map[scheduler.JobName]*scheduler.JobRunSummary{
jobTarget.JobName: {
JobName: jobBName,
ScheduledAt: scheduledAtB,
},
},
Upstreams: []*scheduler.JobLineageSummary{jobXWithLineage},
}
currentJobWithLineage := &scheduler.JobLineageSummary{
JobName: jobTarget.JobName,
IsEnabled: true,
JobRuns: map[scheduler.JobName]*scheduler.JobRunSummary{
jobTarget.JobName: {
JobName: jobTarget.JobName,
ScheduledAt: scheduledAtA,
},
},
// job-X listed before job-B, so the shallow edge is visited first
Upstreams: []*scheduler.JobLineageSummary{jobXWithLineage, jobBWithLineage},
}

jobDurationEstimation[jobTarget.JobName] = func() *time.Duration { d := 30 * time.Minute; return &d }()
jobDurationEstimation[jobBName] = func() *time.Duration { d := 45 * time.Minute; return &d }()

// when
err := jobExpectatorService.PopulateExpectedFinishTime(jobTarget.JobName, currentJobWithLineage, jobRunExpectedFinishTime, jobDurationEstimation, referenceTime)

// then
assert.NoError(t, err)
// job-X's finish time is computed once (from whichever path is visited first) and shared
assert.Equal(t, jobXEndTime, jobRunExpectedFinishTime[scheduler.JobSchedule{JobName: jobXName, ScheduledAt: scheduledAtX}].FinishTime)
// job-B's expected finish incorporates job-X's finish time via the deep path
expectedBFinish := jobXEndTime.Add(45 * time.Minute)
assert.Equal(t, expectedBFinish, jobRunExpectedFinishTime[scheduler.JobSchedule{JobName: jobBName, ScheduledAt: scheduledAtB}].FinishTime)
// job-A's expected finish is max(its own scheduled_at, job-X's direct finish, job-B's finish) + job-A duration;
// job-B's finish (which itself folds in job-X) dominates, proving the diamond didn't get double-counted or dropped
expectedAFinish := expectedBFinish.Add(30 * time.Minute)
assert.Equal(t, expectedAFinish, jobRunExpectedFinishTime[*jobTarget].FinishTime)
})
}

// jobRunExpectationDetailsRepository is an autogenerated mock type for the jobRunExpectationDetailsRepository type
Expand Down
Loading