job-search/internal/util/worker.go

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()
	}
}