package sources //go:generate mockgen --source=./job_progress.go --destination=mock_job_progress_test.go --package=sources import ( "context" "errors" "fmt" "sync" "sync/atomic" "time" ) type JobProgressHook interface { // Start and End marks the overall start and end time for this job. Start(JobProgressRef, time.Time) End(JobProgressRef, time.Time) // StartEnumerating and EndEnumerating marks the start and end time for // calling the source's Enumerate method. If the source does not // support enumeration these methods will never be called. StartEnumerating(JobProgressRef, time.Time) EndEnumerating(JobProgressRef, time.Time) // StartUnitChunking and EndUnitChunking marks the start and end time // for calling the source's ChunkUnit method for a given unit. If the // source does not support enumeration these methods will never be // called. StartUnitChunking(JobProgressRef, SourceUnit, time.Time) EndUnitChunking(JobProgressRef, SourceUnit, time.Time) // ReportError is called when any general error is encountered, usually // from enumeration. ReportError(JobProgressRef, error) // ReportUnit is called when a unit has been enumerated. If the source // does not support enumeration this method will never be called. ReportUnit(JobProgressRef, SourceUnit) // ReportChunk is called when a chunk has been produced for the given // unit. The unit will be nil if the source does not support // enumeration. ReportChunk(JobProgressRef, SourceUnit, *Chunk) // Finish marks the job as done. Finish(JobProgressRef) } // JobProgressRef is a wrapper of a JobProgress for read-only access to its state. type JobProgressRef struct { SourceID int64 JobID int64 jobProgress *JobProgress } // Snapshot returns a snapshot of the job's current metrics. func (r *JobProgressRef) Snapshot() JobProgressMetrics { if r.jobProgress == nil { return JobProgressMetrics{} } return r.jobProgress.Snapshot() } // Done returns a channel that will block until the job has completed. func (r *JobProgressRef) Done() <-chan struct{} { if r.jobProgress == nil { // Return a closed channel so it does not block. ch := make(chan struct{}) close(ch) return ch } return r.jobProgress.Done() } // Fatal is a wrapper around error to differentiate non-fatal errors from fatal // ones. A fatal error is typically from a finished context or any error // returned from a source's Init, Chunks, Enumerate, or ChunkUnit methods. type Fatal struct{ error } func (f Fatal) Error() string { return fmt.Sprintf("fatal: %s", f.error.Error()) } func (f Fatal) Unwrap() error { return f.error } // ChunkError is a custom error type for errors encountered during chunking of // a specific unit. type ChunkError struct { unit SourceUnit err error } func (f ChunkError) Error() string { return fmt.Sprintf("error chunking unit %q: %s", f.unit.SourceUnitID(), f.err.Error()) } func (f ChunkError) Unwrap() error { return f.err } // JobProgress aggregates information about a run of a Source. type JobProgress struct { // Unique identifiers for this job. SourceID int64 JobID int64 // Tracks whether the job is finished or not. ctx context.Context cancel context.CancelFunc // Metrics. metrics JobProgressMetrics metricsLock sync.Mutex // Coarse grained hooks for adding extra functionality when events trigger. hooks []JobProgressHook } // JobProgressMetrics tracks the metrics of a job. type JobProgressMetrics struct { StartTime time.Time EndTime time.Time TotalUnits uint64 FinishedUnits uint64 TotalChunks uint64 Errors []error DoneEnumerating bool } // WithHooks adds hooks to be called when an event triggers. func WithHooks(hooks ...JobProgressHook) func(*JobProgress) { return func(jp *JobProgress) { jp.hooks = append(jp.hooks, hooks...) } } // NewJobProgress creates a new job report for the given source and job ID. func NewJobProgress(sourceID, jobID int64, opts ...func(*JobProgress)) *JobProgress { ctx, cancel := context.WithCancel(context.Background()) jp := &JobProgress{ SourceID: sourceID, JobID: jobID, ctx: ctx, cancel: cancel, } for _, opt := range opts { opt(jp) } return jp } // executeHooks is a helper method to execute all the hooks for the given // closure. func (jp *JobProgress) executeHooks(todo func(hook JobProgressHook)) { for _, hook := range jp.hooks { // TODO: Non-blocking? todo(hook) } } // TODO: Comment all this mess. They are mostly implementing JobProgressHook but // without the JobProgressRef parameter. func (jp *JobProgress) Start(start time.Time) { jp.metricsLock.Lock() jp.metrics.StartTime = start jp.metricsLock.Unlock() jp.executeHooks(func(hook JobProgressHook) { hook.Start(jp.Ref(), start) }) } func (jp *JobProgress) End(end time.Time) { jp.metricsLock.Lock() jp.metrics.EndTime = end jp.metricsLock.Unlock() jp.executeHooks(func(hook JobProgressHook) { hook.End(jp.Ref(), end) }) } func (jp *JobProgress) Finish() { jp.cancel() jp.executeHooks(func(hook JobProgressHook) { hook.Finish(jp.Ref()) }) } func (jp *JobProgress) Done() <-chan struct{} { return jp.ctx.Done() } func (jp *JobProgress) ReportUnit(unit SourceUnit) { atomic.AddUint64(&jp.metrics.TotalUnits, 1) jp.executeHooks(func(hook JobProgressHook) { hook.ReportUnit(jp.Ref(), unit) }) } func (jp *JobProgress) ReportChunk(unit SourceUnit, chunk *Chunk) { atomic.AddUint64(&jp.metrics.TotalChunks, 1) jp.executeHooks(func(hook JobProgressHook) { hook.ReportChunk(jp.Ref(), unit, chunk) }) } func (jp *JobProgress) StartUnitChunking(unit SourceUnit, start time.Time) { // TODO: Record time. jp.executeHooks(func(hook JobProgressHook) { hook.StartUnitChunking(jp.Ref(), unit, start) }) } func (jp *JobProgress) EndUnitChunking(unit SourceUnit, end time.Time) { // TODO: Record time. atomic.AddUint64(&jp.metrics.FinishedUnits, 1) jp.executeHooks(func(hook JobProgressHook) { hook.EndUnitChunking(jp.Ref(), unit, end) }) } func (jp *JobProgress) StartEnumerating(start time.Time) { // TODO: Record time. jp.executeHooks(func(hook JobProgressHook) { hook.StartEnumerating(jp.Ref(), start) }) } func (jp *JobProgress) EndEnumerating(end time.Time) { jp.metricsLock.Lock() // TODO: Record time. jp.metrics.DoneEnumerating = true jp.metricsLock.Unlock() jp.executeHooks(func(hook JobProgressHook) { hook.EndEnumerating(jp.Ref(), end) }) } // Snapshot safely gets the job's current metrics. func (jp *JobProgress) Snapshot() JobProgressMetrics { jp.metricsLock.Lock() defer jp.metricsLock.Unlock() return jp.metrics } // ReportError adds a non-nil error to the aggregate of errors // encountered during scanning. func (jp *JobProgress) ReportError(err error) { if err == nil { return } jp.metricsLock.Lock() jp.metrics.Errors = append(jp.metrics.Errors, err) jp.metricsLock.Unlock() jp.executeHooks(func(hook JobProgressHook) { hook.ReportError(jp.Ref(), err) }) } // Ref provides a read-only reference to the JobProgress. func (jp *JobProgress) Ref() JobProgressRef { return JobProgressRef{ SourceID: jp.SourceID, JobID: jp.JobID, jobProgress: jp, } } // EnumerationErrors joins all errors encountered during initialization or // enumeration. func (m JobProgressMetrics) EnumerationError() error { return errors.Join(m.Errors...) } // ChunkErrors joins all errors encountered during chunking. func (m JobProgressMetrics) ChunkError() error { var aggregate []error for _, err := range m.Errors { var chunkErr ChunkError if ok := errors.As(err, &chunkErr); ok { aggregate = append(aggregate, err) } } return errors.Join(aggregate...) } // FatalError returns the first Fatal error, if any, encountered in the scan. func (m JobProgressMetrics) FatalError() error { for _, err := range m.Errors { var fatalErr Fatal if found := errors.As(err, &fatalErr); found { return fatalErr } } return nil } // FatalErrors returns all of the encountered fatal errors joined together. func (m JobProgressMetrics) FatalErrors() error { var aggregate []error for _, err := range m.Errors { var fatalErr Fatal if found := errors.As(err, &fatalErr); found { aggregate = append(aggregate, fatalErr) } } return errors.Join(aggregate...) } func (m JobProgressMetrics) PercentComplete() int { num := m.FinishedUnits den := m.TotalUnits if num == 0 && den == 0 { return 0 } return int(num * 100 / den) }