package util import ( "context" "errors" "fmt" "log/slog" "time" ) // DefaultWorkerShutdownTimeout is used when Worker.ShutdownTimeout is zero. const DefaultWorkerShutdownTimeout = 10 * time.Second // ErrWorkerNoWork tells the Worker that a run succeeded but found nothing to // do. Tasks return it, or an error wrapping it, to have the run treated as a // success (checkpoint advances) while being logged distinctly. For example: // // var ErrNoNewReleases = fmt.Errorf("%w: no new releases", ErrWorkerNoWork) var ErrWorkerNoWork = errors.New("util: no work") // Task processes everything that has changed since the given checkpoint. // // Returning nil or an error matching ErrWorkerNoWork advances the checkpoint to // the time this run started. Any other error leaves the checkpoint unchanged, // so the next run retries from the same point. type Task func(ctx context.Context, since time.Time) error // Worker runs a Task periodically until its context is cancelled. type Worker struct { // Name is attached to every log line as "source". Name string // Interval is how often Task runs. Must be greater than zero. Interval time.Duration // ShutdownTimeout is how long an in-flight Task may keep running after // the parent context is cancelled. 0 or less means // DefaultWorkerShutdownTimeout. ShutdownTimeout time.Duration // Task is the work to perform on each tick. Task Task // Logger is what the worker logs to; nil means slog.Default(), looked up // when Run starts. Logger *slog.Logger } // Run executes Task immediately, then once per Interval, until ctx is // cancelled. It blocks until shutdown has finished. func (w Worker) Run(ctx context.Context, since time.Time) error { if w.Task == nil { return fmt.Errorf("util: worker %q: no task configured", w.Name) } if w.Interval <= 0 { return fmt.Errorf( "util: worker %q: interval must be greater than 0, got %s", w.Name, w.Interval, ) } timeout := w.ShutdownTimeout if timeout <= 0 { timeout = DefaultWorkerShutdownTimeout } logger := w.Logger if logger == nil { logger = slog.Default() } logger = logger.With("source", w.Name) logger.Info("starting worker", "since", since, "interval", w.Interval) // workCtx is NOT cancelled immediately when ctx is. It gives an in-flight // task up to `timeout` to finish before it's cancelled. workCtx, cancelWork := context.WithCancel(context.WithoutCancel(ctx)) defer cancelWork() go func() { select { case <-ctx.Done(): case <-workCtx.Done(): return } select { case <-time.After(timeout): logger.Error("shutdown timeout exceeded, cancelling in-flight task") cancelWork() case <-workCtx.Done(): } }() ticker := time.NewTicker(w.Interval) defer ticker.Stop() if ctx.Err() == nil { since = w.execute(workCtx, logger, since) } for { // Check for shutdown first so we don't start new work after a // signal even if a tick is also ready. if ctx.Err() != nil { logger.Warn("shutdown signal received, exiting") return nil } select { case <-ctx.Done(): logger.Warn("shutdown signal received, exiting") return nil case <-ticker.C: since = w.execute(workCtx, logger, since) } } } // execute runs one iteration of the task, logs the outcome and returns the // checkpoint to use for the next run. func (w Worker) execute( ctx context.Context, logger *slog.Logger, since time.Time, ) time.Time { start := time.Now() logger.Info("running task", "since", since) err := w.Task(ctx, since) elapsed := time.Since(start) switch { case err == nil: logger.Info("task completed", "elapsed", elapsed) return start case errors.Is(err, ErrWorkerNoWork): logger.Info( "task completed with nothing to do", "reason", err, "elapsed", elapsed, ) return start case errors.Is(err, context.Canceled): logger.Warn("task cancelled", "elapsed", elapsed) return since default: logger.Error("task failed", "error", err, "elapsed", elapsed) return since } } // AwaitTask runs task in its own goroutine and returns its result, or ctx.Err() // if ctx is cancelled first. This lets a Task stop waiting on a call that does // not respond promptly to cancellation. // // The result channel is buffered, so if ctx wins the race the goroutine can // still deliver its result and exit instead of blocking forever. func AwaitTask[T any](ctx context.Context, task func() (T, error)) (T, error) { type outcome struct { value T err error } results := make(chan outcome, 1) go func() { value, err := task() results <- outcome{ value: value, err: err, } }() select { case result := <-results: return result.value, result.err case <-ctx.Done(): var zero T return zero, ctx.Err() } } // AwaitTaskE is AwaitTask for a function that returns only an error. func AwaitTaskE(ctx context.Context, task func() error) error { results := make(chan error, 1) go func() { err := task() results <- err }() select { case err := <-results: return err case <-ctx.Done(): return ctx.Err() } }