pocket-id/backend/internal/job/scheduler.go

package job

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"time"

	backoff "github.com/cenkalti/backoff/v5"
	"github.com/go-co-op/gocron/v2"
	"github.com/google/uuid"
	"go.opentelemetry.io/otel/trace"

	"github.com/pocket-id/pocket-id/backend/internal/service"
	"github.com/pocket-id/pocket-id/backend/internal/tracing"
)

type Scheduler struct {
	scheduler gocron.Scheduler
}

func NewScheduler() (*Scheduler, error) {
	scheduler, err := gocron.NewScheduler()
	if err != nil {
		return nil, fmt.Errorf("failed to create a new scheduler: %w", err)
	}

	return &Scheduler{
		scheduler: scheduler,
	}, nil
}

func (s *Scheduler) RemoveJob(name string) error {
	jobs := s.scheduler.Jobs()

	var errs []error
	for _, job := range jobs {
		if job.Name() == name {
			err := s.scheduler.RemoveJob(job.ID())
			if err != nil {
				errs = append(errs, fmt.Errorf("failed to dequeue job %q with ID %q: %w", name, job.ID().String(), err))
			}
		}
	}

	return errors.Join(errs...)
}

// Run the scheduler.
// This function blocks until the context is canceled.
func (s *Scheduler) Run(ctx context.Context) error {
	slog.Info("Starting job scheduler")
	s.scheduler.Start()

	// Block until context is canceled
	<-ctx.Done()

	err := s.scheduler.Shutdown()
	if err != nil {
		slog.Error("Error shutting down job scheduler", slog.Any("error", err))
	} else {
		slog.Info("Job scheduler shut down")
	}

	return nil
}

func (s *Scheduler) RegisterJob(ctx context.Context, name string, def gocron.JobDefinition, jobFn func(ctx context.Context) error, opts service.RegisterJobOpts) error {
	// Wrap the job in a handler that adds tracing and logging
	jobFn = jobWithObservability(name, jobFn)

	// If a BackOff strategy is provided, wrap the job with retry logic
	if opts.BackOff != nil {
		jobFn = jobWithBackOff(jobFn, opts.BackOff)
	}

	jobOptions := []gocron.JobOption{
		gocron.WithContext(ctx),
		gocron.WithName(name),
	}

	if opts.RunImmediately {
		jobOptions = append(jobOptions, gocron.JobOption(gocron.WithStartImmediately()))
	}

	jobOptions = append(jobOptions, opts.ExtraOptions...)

	_, err := s.scheduler.NewJob(def, gocron.NewTask(jobFn), jobOptions...)
	if err != nil {
		return fmt.Errorf("failed to register job %q: %w", name, err)
	}

	return nil
}

type (
	jobNameKey struct{}
	jobIDKey   struct{}
	jobFn      = func(ctx context.Context) error
)

func jobWithObservability(jobName string, job jobFn) jobFn {
	return func(ctx context.Context) error {
		// Generate a random job ID
		jobID := uuid.NewString()

		// Save in the context
		ctx = context.WithValue(ctx, jobNameKey{}, jobName)
		ctx = context.WithValue(ctx, jobIDKey{}, jobID)

		// Create a new context with the span
		var err error
		ctx, span := tracing.Start(ctx, "pocketid.job."+jobName,
			trace.WithSpanKind(trace.SpanKindInternal),
			trace.WithAttributes(
				tracing.JobID(jobID),
			),
		)
		defer tracing.End(span, err)

		// Log the start
		logger := slog.With(
			slog.String("name", jobName),
			slog.String("jobID", jobID),
		)
		start := time.Now()
		logger.InfoContext(ctx, "Starting job")

		// Run the job
		err = job(ctx)
		d := time.Since(start)
		if err != nil {
			logger.ErrorContext(ctx, "Job failed", slog.Any("error", err), slog.Duration("duration", d))
			return err
		}

		logger.InfoContext(ctx, "Job run successfully", slog.Duration("duration", d))
		return nil
	}
}

func jobWithBackOff(job jobFn, bo backoff.BackOff) jobFn {
	return func(ctx context.Context) error {
		jobName, _ := (ctx.Value(jobNameKey{})).(string)
		jobID, _ := (ctx.Value(jobIDKey{})).(string)

		_, err := backoff.Retry(
			ctx,
			func() (struct{}, error) {
				return struct{}{}, job(ctx)
			},
			backoff.WithBackOff(bo),
			backoff.WithNotify(func(err error, d time.Duration) {
				slog.WarnContext(ctx, "Job failed, retrying",
					slog.String("name", jobName),
					slog.String("jobID", jobID),
					slog.Any("error", err),
					slog.Duration("retryIn", d),
				)
			}),
		)
		return err
	}
}

func jobDefWithJitter(interval time.Duration) gocron.JobDefinition {
	const jitter = 5 * time.Minute

	return gocron.DurationRandomJob(interval-jitter, interval+jitter)
}