summercms/modules/conga/worker.go

package conga

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"runtime/debug"
	"strings"
	"sync/atomic"
	"time"

	"git.golem15.com/golem15/summercms/modules/backpack"
	"git.golem15.com/golem15/summercms/modules/lagoon"
	"git.golem15.com/golem15/summercms/modules/pact"
	"git.golem15.com/golem15/summercms/modules/party"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/riverqueue/river"
	"github.com/riverqueue/river/riverdriver/riverdatabasesql"
	"github.com/riverqueue/river/rivertype"
)

// WorkerOptions selects what a worker runs.
type WorkerOptions struct {
	// Queues limits the worker to these queues; nil means every known queue.
	Queues []string

	// pollInterval overrides River's FetchPollInterval (tests only).
	pollInterval time.Duration
	// pollOnly builds a driver without the LISTEN pool (tests only).
	pollOnly bool
	// retryPolicy overrides River's retry schedule (tests only).
	retryPolicy river.ClientRetryPolicy
}

// Worker is a running River work client. Stop it with Stop.
type Worker struct {
	m        *Manager
	client   *river.Client[*sql.Tx]
	listener *pgxpool.Pool
	queues   []string
	counters *workerCounters
	// concurrency is each queue's MaxWorkers, for the startup summary.
	concurrency map[string]int
}

// WorkerStats counts the job attempts a worker ran since it started.
type WorkerStats struct {
	// Done is the number of attempts that returned no error.
	Done int64
	// Failed is the number of attempts that returned an error, other than
	// those cancelled by Stop (HandedBack).
	Failed int64
	// InFlight is the number of attempts running now.
	InFlight int64
	// HandedBack is the number of attempts cancelled by Stop after its
	// deadline (River's StopAndCancel). River records the cancellation as
	// an errored attempt, so the job is retried later, or fails for good
	// when it was its last attempt.
	HandedBack int64
}

// workerCounters are the live counters behind WorkerStats.
type workerCounters struct {
	done, failed, inFlight, handedBack atomic.Int64
	// stopping is set by Stop before running jobs are cancelled.
	stopping atomic.Bool
}

// Stats returns the worker's attempt counters. A nil worker has none.
func (w *Worker) Stats() WorkerStats {
	if w == nil || w.counters == nil {
		return WorkerStats{}
	}
	c := w.counters
	return WorkerStats{
		Done:       c.done.Load(),
		Failed:     c.failed.Load(),
		InFlight:   c.inFlight.Load(),
		HandedBack: c.handedBack.Load(),
	}
}

// Queues returns the sorted queue names the worker runs.
func (w *Worker) Queues() []string {
	if w == nil {
		return nil
	}
	return append([]string(nil), w.queues...)
}

// StartWorker registers every pact.HasJobs job of plugins and starts one
// River client on the shared *sql.DB. Its driver is
// riverdatabasesql.NewWithPgxListener: every query runs on the shared pool
// and only Postgres LISTEN runs on a dedicated pgx pool (MaxConns 1) opened
// from database.dsn, so jobs committed by any process are picked up without
// waiting for the poll interval. The worker keeps running after ctx is
// cancelled; stop it with Worker.Stop.
//
// Every worker also carries the River periodic jobs of plugins'
// pact.HasSchedule entries (only the elected leader enqueues them) and runs
// the built-in scheduled command job on QueueScheduled when that queue is
// selected. An invalid schedule entry or app.timezone fails the start.
func StartWorker(ctx context.Context, app *backpack.App, plugins []party.Plugin, o WorkerOptions) (*Worker, error) {
	m, err := From(app)
	if err != nil {
		return nil, err
	}
	for _, p := range plugins {
		hj, ok := p.(pact.HasJobs)
		if !ok {
			continue
		}
		for _, j := range hj.Jobs() {
			if err := m.Register(j); err != nil {
				return nil, fmt.Errorf("conga: plugin %s: %w", p.ID(), err)
			}
		}
	}
	sqlDB, ok := app.Lookup[*sql.DB]()
	if !ok || sqlDB == nil {
		return nil, ErrNoDatabase
	}
	s := settingsFromApp(app)
	log := riverLogger(app)
	periodic, table, err := periodicJobs(app, plugins)
	if err != nil {
		return nil, err
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	if m.worker != nil {
		return nil, fmt.Errorf("conga: a worker is already running")
	}
	if err := m.registerLocked(m.scheduledJobLocked()); err != nil {
		return nil, fmt.Errorf("conga: register scheduled command job: %w", err)
	}
	known := knownQueues(s, m.jobs)
	if _, configured := s.queues[QueueScheduled]; !configured {
		known[QueueScheduled] = scheduledMaxWorkers
	}
	queues, err := selectQueues(known, o.Queues)
	if err != nil {
		return nil, err
	}
	workers := river.NewWorkers()
	for _, j := range m.jobs {
		if err := j.register(workers, m); err != nil {
			return nil, fmt.Errorf("conga: register %s: %w", j.kind(), err)
		}
	}
	if len(m.jobs) == 0 {
		// River refuses to start a client without workers; an app with no
		// jobs yet still gets a worker that starts and idles.
		if err := river.AddWorkerSafely[idleArgs](workers, idleWorker{}); err != nil {
			return nil, fmt.Errorf("conga: register idle worker: %w", err)
		}
	}
	cfg := baseConfig(s, log)
	cfg.Workers = workers
	// Every worker carries the periodic jobs whatever its queue filter:
	// only the elected leader enqueues them.
	cfg.PeriodicJobs = periodic
	cfg.Queues = map[string]river.QueueConfig{}
	for name, n := range queues {
		cfg.Queues[name] = river.QueueConfig{MaxWorkers: n}
	}
	if o.pollInterval > 0 {
		cfg.FetchPollInterval = o.pollInterval
	}
	if o.retryPolicy != nil {
		cfg.RetryPolicy = o.retryPolicy
	}

	var listener *pgxpool.Pool
	var driver *riverdatabasesql.Driver
	if o.pollOnly {
		driver = riverdatabasesql.New(sqlDB)
	} else {
		listener, err = listenerPool(ctx, app)
		if err != nil {
			return nil, err
		}
		driver = riverdatabasesql.NewWithPgxListener(sqlDB, listener)
	}
	client, err := river.NewClient(driver, cfg)
	if err != nil {
		closePool(listener)
		return nil, fmt.Errorf("conga: river client: %w", err)
	}
	if err := client.Start(context.WithoutCancel(ctx)); err != nil {
		closePool(listener)
		return nil, fmt.Errorf("conga: start worker: %w", err)
	}
	counters := &workerCounters{}
	m.worker = client
	m.schedule = table
	m.counters.Store(counters)
	return &Worker{m: m, client: client, listener: listener, queues: sortedQueueNames(queues), counters: counters, concurrency: queues}, nil
}

// StartServeWorker starts the in-process worker of the serve command on
// every known queue. It returns nil, nil when queue.work_in_serve is false,
// for deployments that run a separate queue:work process.
func StartServeWorker(ctx context.Context, app *backpack.App, plugins []party.Plugin) (*Worker, error) {
	if !settingsFromApp(app).workInServe {
		return nil, nil
	}
	return StartWorker(ctx, app, plugins, WorkerOptions{})
}

// Stop stops fetching jobs and waits for running jobs to finish. When ctx
// ends first, running jobs are cancelled. The listener pool is closed.
func (w *Worker) Stop(ctx context.Context) error {
	if w == nil || w.client == nil {
		return nil
	}
	err := w.client.Stop(ctx)
	if err != nil && ctx.Err() != nil {
		if w.counters != nil {
			w.counters.stopping.Store(true)
		}
		hardCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
		err = w.client.StopAndCancel(hardCtx)
		cancel()
	}
	closePool(w.listener)
	w.m.mu.Lock()
	if w.m.worker == w.client {
		w.m.worker = nil
	}
	w.m.mu.Unlock()
	if err != nil {
		return fmt.Errorf("conga: stop worker: %w", err)
	}
	return nil
}

func listenerPool(ctx context.Context, app *backpack.App) (*pgxpool.Pool, error) {
	dsn := ""
	if app != nil {
		dsn = lagoon.DSN(app.Config)
	}
	if dsn == "" {
		return nil, fmt.Errorf("conga: database.dsn is empty; the worker's LISTEN pool needs it")
	}
	cfg, err := pgxpool.ParseConfig(dsn)
	if err != nil {
		return nil, fmt.Errorf("conga: listener pool config: %w", err)
	}
	cfg.MaxConns = 1
	cfg.MinConns = 0
	pool, err := pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		return nil, fmt.Errorf("conga: listener pool: %w", err)
	}
	return pool, nil
}

func closePool(p *pgxpool.Pool) {
	if p != nil {
		p.Close()
	}
}

func selectQueues(known map[string]int, want []string) (map[string]int, error) {
	if len(want) == 0 {
		return known, nil
	}
	out := map[string]int{}
	for _, name := range want {
		name = strings.TrimSpace(name)
		n, ok := known[name]
		if !ok {
			names := sortedQueueNames(known)
			return nil, fmt.Errorf("%w %q (known queues: %s)", ErrUnknownQueue, name, strings.Join(names, ", "))
		}
		out[name] = n
	}
	return out, nil
}

// runAttempt runs one River attempt of a job, counts it in the running
// worker's WorkerStats and applies the summer_jobs rules: the row id goes into ctx; a panic becomes an error; an error on a
// STOPPED or canceled row cancels the River job; an error on the final
// attempt sets StatusError with the error text under metadata "error". An
// error on an earlier attempt leaves the row as it is so River can retry.
func (m *Manager) runAttempt(ctx context.Context, row *rivertype.JobRow, fn func(context.Context) error) error {
	id, hasRow := summerJobID(row)
	if hasRow {
		ctx = withJobID(ctx, id)
	}
	c := m.counters.Load()
	if c != nil {
		c.inFlight.Add(1)
	}
	err := m.call(ctx, row, fn)
	if c != nil {
		c.inFlight.Add(-1)
		switch {
		case err == nil:
			c.done.Add(1)
		case c.stopping.Load() && ctx.Err() != nil:
			c.handedBack.Add(1)
		default:
			c.failed.Add(1)
		}
	}
	if err == nil || !hasRow {
		return err
	}
	dbCtx := context.WithoutCancel(ctx)
	if rec, gerr := m.Get(dbCtx, id); gerr == nil && (rec.Status == StatusStopped || rec.IsCanceled) {
		return river.JobCancel(err)
	}
	if row.Attempt >= row.MaxAttempts {
		if ferr := m.failWithError(dbCtx, id, err); ferr != nil {
			loggerFromApp(m.app).Error("conga: record job failure", "job_id", id, "kind", row.Kind, "error", ferr)
		}
	}
	return err
}

func (m *Manager) call(ctx context.Context, row *rivertype.JobRow, fn func(context.Context) error) (err error) {
	defer func() {
		if r := recover(); r != nil {
			loggerFromApp(m.app).Error("conga: job panicked", "kind", row.Kind, "river_job_id", row.ID, "panic", fmt.Sprint(r), "stack", string(debug.Stack()))
			err = fmt.Errorf("conga: job %s panicked: %v", row.Kind, r)
		}
	}()
	return fn(ctx)
}

func summerJobID(row *rivertype.JobRow) (uint, bool) {
	if row == nil || len(row.Metadata) == 0 {
		return 0, false
	}
	var meta struct {
		SummerJobID uint `json:"summer_job_id"`
	}
	if err := json.Unmarshal(row.Metadata, &meta); err != nil || meta.SummerJobID == 0 {
		return 0, false
	}
	return meta.SummerJobID, true
}

// idleArgs is the kind of the placeholder worker registered when an app has
// no jobs. Nothing inserts it.
type idleArgs struct{}

func (idleArgs) Kind() string { return "summercms_conga_idle" }

type idleWorker struct {
	river.WorkerDefaults[idleArgs]
}

func (idleWorker) Work(context.Context, *river.Job[idleArgs]) error { return nil }