From 7083400d75f26b0a6d2bb3cb3008eb631ce6708b Mon Sep 17 00:00:00 2001 From: "Ahmad N. F" Date: Mon, 24 Aug 2026 14:44:38 +0700 Subject: [PATCH] fix(job-expected-finish-time): fix run key by immediate downstream --- .../service/job_expectator_service.go | 19 +- .../service/job_expectator_service_test.go | 196 +++++++++++++++++- 2 files changed, 199 insertions(+), 16 deletions(-) diff --git a/core/scheduler/service/job_expectator_service.go b/core/scheduler/service/job_expectator_service.go index 664af962cf..fc7fff33a4 100644 --- a/core/scheduler/service/job_expectator_service.go +++ b/core/scheduler/service/job_expectator_service.go @@ -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 @@ -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 @@ -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, @@ -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 { diff --git a/core/scheduler/service/job_expectator_service_test.go b/core/scheduler/service/job_expectator_service_test.go index 0c97ee610c..66457f2ce6 100644 --- a/core/scheduler/service/job_expectator_service_test.go +++ b/core/scheduler/service/job_expectator_service_test.go @@ -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) @@ -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) @@ -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) @@ -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 @@ -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) @@ -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) @@ -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) @@ -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) @@ -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