// Package jobs runs background work on River, a Postgres-backed job queue. package jobs import ( "context" "errors" "fmt" "log" "strconv" "time" "xprem/internal/database" "xprem/internal/database/postgres" "xprem/internal/repository" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" "github.com/riverqueue/river" "github.com/riverqueue/river/riverdriver/riverpgxv5" "github.com/riverqueue/river/rivermigrate" "github.com/riverqueue/river/rivertype" ) var ErrAlreadyRunning = errors.New("a job of this kind is already running for this scope") // QueueBSDiff runs the bundle patch jobs. Each one holds two bundles and the // patch in memory, so the queue stays narrow. const QueueBSDiff = "bsdiff" // QueueSourcemapIndex runs the source map index jobs. Each one decodes a // whole map in memory, so one at a time. const QueueSourcemapIndex = "sourcemap-index" type Client struct { pool *pgxpool.Pool workers *river.Workers periodic []*river.PeriodicJob riverClient *river.Client[pgx.Tx] } func NewClient(engine *database.Engine) (*Client, error) { pool, ok := engine.DB.(*pgxpool.Pool) if !ok { return nil, fmt.Errorf("the jobs client needs the pgx pool, got %T", engine.DB) } return &Client{pool: pool, workers: river.NewWorkers()}, nil } // Workers is the registry to add workers to, before Start. func (c *Client) Workers() *river.Workers { return c.workers } // AddPeriodic schedules a job, before Start. River inserts it from one // replica only, the elected leader. func (c *Client) AddPeriodic(job *river.PeriodicJob) { c.periodic = append(c.periodic, job) } func (c *Client) Start(ctx context.Context) error { driver := riverpgxv5.New(c.pool) migrator, err := rivermigrate.New(driver, nil) if err != nil { return fmt.Errorf("failed to prepare the river migrator: %w", err) } release, err := postgres.AcquireAdvisoryLock(ctx, c.pool.Config().ConnConfig, postgres.RiverMigrationLockID, "river migration") if err != nil { return err } _, err = migrator.Migrate(ctx, rivermigrate.DirectionUp, nil) release() if err != nil { return fmt.Errorf("failed to migrate the river schema: %w", err) } riverClient, err := river.NewClient(driver, &river.Config{ Queues: map[string]river.QueueConfig{ river.QueueDefault: {MaxWorkers: 10}, QueueBSDiff: {MaxWorkers: 2}, QueueSourcemapIndex: {MaxWorkers: 1}, }, Workers: c.workers, PeriodicJobs: c.periodic, }) if err != nil { return fmt.Errorf("failed to build the river client: %w", err) } if err := riverClient.Start(ctx); err != nil { return fmt.Errorf("failed to start the river client: %w", err) } c.riverClient = riverClient log.Println("🛠️ [JOBS] River job client started") return nil } func (c *Client) Stop() { if c == nil || c.riverClient == nil { return } ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() if err := c.riverClient.Stop(ctx); err != nil { log.Printf("[jobs] river client did not stop cleanly: %v", err) } } // Enqueue inserts args as a job; ErrAlreadyRunning when a unique constraint // skipped the insert. func (c *Client) Enqueue(ctx context.Context, args river.JobArgs) (string, error) { if c == nil || c.riverClient == nil { return "", repository.ErrNotSupportedInStatelessMode } inserted, err := c.riverClient.Insert(ctx, args, nil) if err != nil { return "", fmt.Errorf("failed to enqueue the job: %w", err) } if inserted.UniqueSkippedAsDuplicate { return "", ErrAlreadyRunning } return strconv.FormatInt(inserted.Job.ID, 10), nil } func (c *Client) Get(ctx context.Context, jobId string) (*rivertype.JobRow, error) { if c == nil || c.riverClient == nil { return nil, nil } id, err := strconv.ParseInt(jobId, 10, 64) if err != nil { return nil, nil } job, err := c.riverClient.JobGet(ctx, id) if errors.Is(err, rivertype.ErrNotFound) { return nil, nil } if err != nil { return nil, fmt.Errorf("failed to read the job: %w", err) } return job, nil } // LatestByArg returns the newest job of the kind whose args carry the value // at the key, or nil. func (c *Client) LatestByArg(ctx context.Context, kind string, argKey string, argValue string) (*rivertype.JobRow, error) { if c == nil || c.riverClient == nil { return nil, nil } params := river.NewJobListParams(). Kinds(kind). Where("args->>@key = @value", river.NamedArgs{"key": argKey, "value": argValue}). OrderBy(river.JobListOrderByID, river.SortOrderDesc). First(1) listed, err := c.riverClient.JobList(ctx, params) if err != nil { return nil, fmt.Errorf("failed to list the scope's jobs: %w", err) } if len(listed.Jobs) == 0 { return nil, nil } return listed.Jobs[0], nil } // Cancel stops a job: immediately when it still waits, at its next context // check when it runs. False means the job is unknown. func (c *Client) Cancel(ctx context.Context, jobId string) (bool, error) { if c == nil || c.riverClient == nil { return false, nil } id, err := strconv.ParseInt(jobId, 10, 64) if err != nil { return false, nil } _, err = c.riverClient.JobCancel(ctx, id) if errors.Is(err, rivertype.ErrNotFound) { return false, nil } if err != nil { return false, fmt.Errorf("failed to cancel the job: %w", err) } return true, nil }